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()