def rotato_patch_to_center(self, patch, angle):
# Use np.copy to avoid corrupting original data
w, h, c = patch.shape
xyz = np.copy(patch).reshape(-1, 3)
xyz_updated = rotate_pc_along_y(xyz, angle)
patch_updated = xyz_updated.reshape(w, h, c)
return patch_updated
def from_prediction_to_label_format(center, angle_class, angle_res,\
size_class, size_res, rot_angle):
''' Convert predicted box parameters to label format. '''
l, w, h = class2size(size_class, size_res)
ry = class2angle(angle_class, angle_res, NUM_HEADING_BIN) + rot_angle
tx, ty, tz = rotate_pc_along_y(np.expand_dims(center, 0),
-rot_angle).squeeze()
ty += h / 2.0
return h, w, l, tx, ty, tz, ry
def get_center_view_point_set(self, index):
''' Frustum rotation of point clouds.
NxC points with first 3 channels as XYZ
z is facing forward, x is left ward, y is downward
'''
# Use np.copy to avoid corrupting original data
point_set = np.copy(self.input_list[index])
return rotate_pc_along_y(point_set,
self.get_center_view_rot_angle(index))
def __getitem__(self, index):
''' Get index-th element from the picked file dataset. '''
# ------------------------------ INPUTS ----------------------------
rot_angle = self.get_center_view_rot_angle(index)
patch = self.patch_xyz_list[index]
if self.add_rgb:
rgb = self.patch_rgb_list[index]
patch = np.concatenate((patch, rgb), axis=-1)
self.rotate_to_center = True
if self.rotate_to_center:
patch[:, :,
0:3] = self.rotato_patch_to_center(patch[:, :, 0:3],
rot_angle)
# transport patch to fixed size
# and change shape to C * H * W from H * W * C
patch = torch.from_numpy(patch) # H * W * C
patch = patch.unsqueeze(0) # 1 * H * W * C
patch = patch.transpose(2, 3).transpose(
1, 2) # 1 * H * W * C -> 1 * H * C * W -> 1 * C * H * W
patch = F.interpolate(patch,
self.patch_size,
mode='bilinear',
align_corners=True).squeeze(0).numpy()
cls_type = self.type_list[index]
assert (cls_type in ['Car', 'Pedestrian', 'Cyclist'])
one_hot_vec = np.zeros((3), dtype=np.float32)
one_hot_vec[self.dataset_helper.type2onehot[cls_type]] = 1
if self.from_rgb_detection:
return patch, rot_angle, self.prob_list[index], self.id_list[index], \
self.type_list[index], self.box2d_list[index], one_hot_vec
# ------------------------------ LABELS ----------------------------
# size labels
size_class, size_residual = self.dataset_helper.size2class(
self.box3d_size_list[index], self.type_list[index])
# center labels
center = self.box3d_center_list[index]
if self.rotate_to_center:
center = rotate_pc_along_y(np.expand_dims(
center, 0), self.get_center_view_rot_angle(index)).squeeze()
# heading labels
heading_angle = self.heading_list[index]
if self.rotate_to_center:
heading_angle = heading_angle - rot_angle
if self.random_flip:
if np.random.random() > 0.5: # 50% chance flipping
patch[0, :, :] *= -1
center[0] *= -1
heading_angle = np.pi - heading_angle
if self.random_shift: # random shift object center
dist = np.sqrt(np.sum(center[0]**2 + center[2]**2))
shift = np.clip(np.random.randn() * dist * 0.2, -dist * 0.2,
dist * 0.2)
patch[2, :, :] += shift
center[2] += shift
angle_class, angle_residual = self.dataset_helper.angle2class(
heading_angle)
return patch, center, angle_class, angle_residual, size_class, size_residual, rot_angle, one_hot_vec
def get_center_view_box3d_center(self, index):
''' Frustum rotation of 3D bounding box center. '''
box3d_center = (self.box3d_list[index][0, :] +
self.box3d_list[index][6, :]) / 2.0
return rotate_pc_along_y(np.expand_dims(
box3d_center, 0), self.get_center_view_rot_angle(index)).squeeze()