def data_augmentation(self, aug_pts_rect, aug_gt_boxes3d, mustaug=False, stage=1):
"""
:param aug_pts_rect: (N, 3)
:param aug_gt_boxes3d: (N, 7)
:param gt_alpha: (N)
:return:
"""
aug_list = cfg.AUG_METHOD_LIST
aug_enable = 1 - np.random.rand(3)
if mustaug is True:
aug_enable[0] = -1
aug_enable[1] = -1
aug_method = []
if 'rotation' in aug_list and aug_enable[0] < cfg.AUG_METHOD_PROB[0]:
angle = np.random.uniform(-np.pi / cfg.AUG_ROT_RANGE, np.pi / cfg.AUG_ROT_RANGE)
aug_pts_rect = kitti_utils.rotate_pc_along_y(aug_pts_rect, rot_angle=angle)
aug_gt_boxes3d = kitti_utils.rotate_pc_along_y(aug_gt_boxes3d, rot_angle=angle)
aug_method.append(['rotation', angle])
if 'scaling' in aug_list and aug_enable[1] < cfg.AUG_METHOD_PROB[1]:
scale = np.random.uniform(0.95, 1.05)
aug_pts_rect = aug_pts_rect * scale
aug_gt_boxes3d[:, 0:6] = aug_gt_boxes3d[:, 0:6] * scale
aug_method.append(['scaling', scale])
if 'flip' in aug_list and aug_enable[2] < cfg.AUG_METHOD_PROB[2]:
# flip horizontal
aug_pts_rect[:, 0] = -aug_pts_rect[:, 0]
aug_gt_boxes3d[:, 0] = -aug_gt_boxes3d[:, 0]
# flip orientation: ry > 0: pi - ry, ry < 0: -pi - ry
aug_method.append('flip')
return aug_pts_rect, aug_gt_boxes3d, aug_method
def data_augmentation(self, aug_pts_rect, aug_gt_boxes3d, gt_alpha, sample_id=None, mustaug=False, stage=1):
"""
:param aug_pts_rect: (N, 3)
:param aug_gt_boxes3d: (N, 7)
:param gt_alpha: (N)
:return:
"""
aug_list = cfg.AUG_METHOD_LIST
aug_enable = 1 - np.random.rand(3)
if mustaug is True:
aug_enable[0] = -1
aug_enable[1] = -1
aug_method = []
if 'rotation' in aug_list and aug_enable[0] < cfg.AUG_METHOD_PROB[0]:
angle = np.random.uniform(-np.pi / cfg.AUG_ROT_RANGE, np.pi / cfg.AUG_ROT_RANGE)
aug_pts_rect = kitti_utils.rotate_pc_along_y(aug_pts_rect, rot_angle=angle)
if stage == 1:
# xyz change, hwl unchange
aug_gt_boxes3d = kitti_utils.rotate_pc_along_y(aug_gt_boxes3d, rot_angle=angle)
# calculate the ry after rotation
x, z = aug_gt_boxes3d[:, 0], aug_gt_boxes3d[:, 2]
beta = np.arctan2(z, x)
new_ry = np.sign(beta) * np.pi / 2 + gt_alpha - beta
aug_gt_boxes3d[:, 6] = new_ry # TODO: not in [-np.pi / 2, np.pi / 2]
elif stage == 2:
# for debug stage-2, this implementation has little float precision difference with the above one
assert aug_gt_boxes3d.shape[0] == 2
aug_gt_boxes3d[0] = self.rotate_box3d_along_y(aug_gt_boxes3d[0], angle)
aug_gt_boxes3d[1] = self.rotate_box3d_along_y(aug_gt_boxes3d[1], angle)
else:
raise NotImplementedError
aug_method.append(['rotation', angle])
if 'scaling' in aug_list and aug_enable[1] < cfg.AUG_METHOD_PROB[1]:
scale = np.random.uniform(0.95, 1.05)
aug_pts_rect = aug_pts_rect * scale
aug_gt_boxes3d[:, 0:6] = aug_gt_boxes3d[:, 0:6] * scale
aug_method.append(['scaling', scale])
if 'flip' in aug_list and aug_enable[2] < cfg.AUG_METHOD_PROB[2]:
# flip horizontal
aug_pts_rect[:, 0] = -aug_pts_rect[:, 0]
aug_gt_boxes3d[:, 0] = -aug_gt_boxes3d[:, 0]
# flip orientation: ry > 0: pi - ry, ry < 0: -pi - ry
if stage == 1:
aug_gt_boxes3d[:, 6] = np.sign(aug_gt_boxes3d[:, 6]) * np.pi - aug_gt_boxes3d[:, 6]
elif stage == 2:
assert aug_gt_boxes3d.shape[0] == 2
aug_gt_boxes3d[0, 6] = np.sign(aug_gt_boxes3d[0, 6]) * np.pi - aug_gt_boxes3d[0, 6]
aug_gt_boxes3d[1, 6] = np.sign(aug_gt_boxes3d[1, 6]) * np.pi - aug_gt_boxes3d[1, 6]
else:
raise NotImplementedError
aug_method.append('flip')
return aug_pts_rect, aug_gt_boxes3d, aug_method
def canonical_transform(pts_input, roi_box3d, gt_box3d):
roi_ry = roi_box3d[6] % (2 * np.pi) # 0 ~ 2pi
roi_center = roi_box3d[0:3]
# shift to center
pts_input[:, [0, 1, 2]] = pts_input[:, [0, 1, 2]] - roi_center
gt_box3d_ct = np.copy(gt_box3d)
gt_box3d_ct[0:3] = gt_box3d_ct[0:3] - roi_center
# rotate to the direction of head
gt_box3d_ct = kitti_utils.rotate_pc_along_y(gt_box3d_ct.reshape(1, 7), roi_ry).reshape(7)
gt_box3d_ct[6] = gt_box3d_ct[6] - roi_ry
pts_input = kitti_utils.rotate_pc_along_y(pts_input, roi_ry)
return pts_input, gt_box3d_ct
def rotate_box3d_along_y(self, box3d, rot_angle):
old_x, old_z, ry = box3d[0], box3d[2], box3d[6]
old_beta = np.arctan2(old_z, old_x)
alpha = -np.sign(old_beta) * np.pi / 2 + old_beta + ry
box3d = kitti_utils.rotate_pc_along_y(box3d.reshape(1, 7), rot_angle=rot_angle)[0]
new_x, new_z = box3d[0], box3d[2]
new_beta = np.arctan2(new_z, new_x)
box3d[6] = np.sign(new_beta) * np.pi / 2 + alpha - new_beta
return box3d
def roipool3d_cpu(boxes3d,
pts,
pts_feature,
pts_extra_input,
pool_extra_width,
sampled_pt_num=512,
canonical_transform=True):
"""
:param boxes3d: (N, 7)
:param pts: (N, 3)
:param pts_feature: (N, C)
:param pts_extra_input: (N, C2)
:param pool_extra_width: constant
:param sampled_pt_num: constant
:return:
"""
pooled_boxes3d = kitti_utils.enlarge_box3d(boxes3d, pool_extra_width)
pts_feature_all = np.concatenate((pts_extra_input, pts_feature), axis=1)
# Note: if pooled_empty_flag[i] > 0, the pooled_pts[i], pooled_features[i] will be zero
pooled_pts, pooled_features, pooled_empty_flag = \
roipool_pc_cpu(torch.from_numpy(pts), torch.from_numpy(pts_feature_all),
torch.from_numpy(pooled_boxes3d), sampled_pt_num)
extra_input_len = pts_extra_input.shape[1]
sampled_pts_input = torch.cat(
(pooled_pts, pooled_features[:, :, 0:extra_input_len]), dim=2).numpy()
sampled_pts_feature = pooled_features[:, :, extra_input_len:].numpy()
if canonical_transform:
# Translate to the roi coordinates
roi_ry = boxes3d[:, 6] % (2 * np.pi) # 0~2pi
roi_center = boxes3d[:, 0:3]
# shift to center
sampled_pts_input[:, :,
0:3] = sampled_pts_input[:, :, 0:
3] - roi_center[:, np.
newaxis, :]
for k in range(sampled_pts_input.shape[0]):
sampled_pts_input[k] = kitti_utils.rotate_pc_along_y(
sampled_pts_input[k], roi_ry[k])
return sampled_pts_input, sampled_pts_feature
return sampled_pts_input, sampled_pts_feature, pooled_empty_flag.numpy()