def noise_per_object_v4_(gt_boxes, points=None, valid_mask=None, rotation_perturb=np.pi / 4, center_noise_std=1.0,\ global_random_rot_range=np.pi / 4, num_try=5, group_ids=None, data_aug_with_context=-1.0,\ data_aug_random_drop=-1.0): """ perform random rotation and translation on each groundtruth independently. Args: gt_boxes: [N, 7], gt box in lidar, [x, y, z, w, l, h, ry] points: [M, 4], point cloud in lidar, all points in the scene. rotation_perturb: [-0.785, 0.785], center_noise_std: [1.0, 1.0, 0.5], global_random_rot_range: [0, 0] num_try: 100 """ num_boxes = gt_boxes.shape[0] center_noise_std = np.array(center_noise_std, dtype=gt_boxes.dtype) loc_noises = np.random.normal(scale=center_noise_std, size=[num_boxes, num_try, 3]) rot_noises = np.random.uniform(rotation_perturb[0], rotation_perturb[1], size=[num_boxes, num_try]) origin = [0.5, 0.5, 0.5] offset = [0.0, 0.0, 0.0, 0.0, 0.0] if data_aug_with_context > 0: # to enlarge boxes and keep context offset = [0.0, 0.0, data_aug_with_context, data_aug_with_context, 0.0] # gt_boxes: x,y,z(lidar), w, l, h, ry(cam). gt_box_corners = box_np_ops.center_to_corner_box3d(gt_boxes[:, :3], gt_boxes[:, 3:6] + offset[2:5], gt_boxes[:, 6], origin=origin, axis=2) # noise on bev_box (x, y, w, l, ry). # todo: seems without noise in z-axis velo, because only cal bev iou, but z-axis noise added later. selected_noise = noise_per_box(gt_boxes[:, [0, 1, 3, 4, 6]] + offset, valid_mask, loc_noises, rot_noises) # loc_transforms = _select_transform(loc_noises, selected_noise) rot_transforms = _select_transform(rot_noises, selected_noise) surfaces = box_np_ops.corner_to_surfaces_3d_jit(gt_box_corners) if points is not None: point_masks = points_in_convex_polygon_3d_jit(points[:, :3], surfaces) # todo: perform random drop here. if data_aug_random_drop > 0.0: point_masks = point_masks points_transform_( points, gt_boxes[:, :3], point_masks, loc_transforms, rot_transforms, valid_mask, ) box3d_transform_(gt_boxes, loc_transforms, rot_transforms, valid_mask)
def noise_per_object_v2_( gt_boxes, points=None, valid_mask=None, rotation_perturb=np.pi / 4, center_noise_std=1.0, global_random_rot_range=np.pi / 4, num_try=100, ): """random rotate or remove each groundtrutn independently. use kitti viewer to test this function points_transform_ Args: gt_boxes: [N, 7], gt box in lidar.points_transform_ points: [M, 4], point cloud in lidar. """ num_boxes = gt_boxes.shape[0] if not isinstance(rotation_perturb, (list, tuple, np.ndarray)): rotation_perturb = [-rotation_perturb, rotation_perturb] if not isinstance(global_random_rot_range, (list, tuple, np.ndarray)): global_random_rot_range = [ -global_random_rot_range, global_random_rot_range ] if not isinstance(center_noise_std, (list, tuple, np.ndarray)): center_noise_std = [ center_noise_std, center_noise_std, center_noise_std ] if valid_mask is None: valid_mask = np.ones((num_boxes, ), dtype=np.bool_) center_noise_std = np.array(center_noise_std, dtype=gt_boxes.dtype) loc_noises = np.random.normal(scale=center_noise_std, size=[num_boxes, num_try, 3]) # loc_noises = np.random.uniform( # -center_noise_std, center_noise_std, size=[num_boxes, num_try, 3]) rot_noises = np.random.uniform(rotation_perturb[0], rotation_perturb[1], size=[num_boxes, num_try]) gt_grots = np.arctan2(gt_boxes[:, 0], gt_boxes[:, 1]) grot_lowers = global_random_rot_range[0] - gt_grots grot_uppers = global_random_rot_range[1] - gt_grots global_rot_noises = np.random.uniform( grot_lowers[..., np.newaxis], grot_uppers[..., np.newaxis], size=[num_boxes, num_try], ) origin = [0.5, 0.5, 0] gt_box_corners = box_np_ops.center_to_corner_box3d(gt_boxes[:, :3], gt_boxes[:, 3:6], gt_boxes[:, 6], origin=origin, axis=2) if np.abs(global_random_rot_range[0] - global_random_rot_range[1]) < 1e-3: selected_noise = noise_per_box(gt_boxes[:, [0, 1, 3, 4, 6]], valid_mask, loc_noises, rot_noises) else: selected_noise = noise_per_box_v2_( gt_boxes[:, [0, 1, 3, 4, 6]], valid_mask, loc_noises, rot_noises, global_rot_noises, ) loc_transforms = _select_transform(loc_noises, selected_noise) rot_transforms = _select_transform(rot_noises, selected_noise) if points is not None: surfaces = box_np_ops.corner_to_surfaces_3d_jit(gt_box_corners) point_masks = points_in_convex_polygon_3d_jit(points[:, :3], surfaces) points_transform_( points, gt_boxes[:, :3], point_masks, loc_transforms, rot_transforms, valid_mask, ) box3d_transform_(gt_boxes, loc_transforms, rot_transforms, valid_mask)
def __call__(self, res, info): res["mode"] = self.mode if res["type"] in ["KittiDataset", "WaymoDataset", "LyftDataset"]: points = res["lidar"]["points"] elif res["type"] in ["NuScenesDataset"]: points = res["lidar"]["combined"] if self.mode == "train": anno_dict = res["lidar"]["annotations"] gt_dict = { "gt_boxes": anno_dict["boxes"], "gt_names": np.array(anno_dict["names"]).reshape(-1), } if "difficulty" not in anno_dict: difficulty = np.zeros([anno_dict["boxes"].shape[0]], dtype=np.int32) gt_dict["difficulty"] = difficulty else: gt_dict["difficulty"] = anno_dict["difficulty"] if "calib" in res: calib = res["calib"] else: calib = None if self.add_rgb_to_points: assert calib is not None and "image" in res image_path = res["image"]["image_path"] image = (imgio.imread(str(pathlib.Path(root_path) / image_path)).astype(np.float32) / 255) points_rgb = box_np_ops.add_rgb_to_points(points, image, calib["rect"], calib["Trv2c"], calib["P2"]) points = np.concatenate([points, points_rgb], axis=1) num_point_features += 3 if self.reference_detections is not None: assert calib is not None and "image" in res C, R, T = box_np_ops.projection_matrix_to_CRT_kitti(P2) frustums = box_np_ops.get_frustum_v2(reference_detections, C) frustums -= T frustums = np.einsum("ij, akj->aki", np.linalg.inv(R), frustums) frustums = box_np_ops.camera_to_lidar(frustums, rect, Trv2c) surfaces = box_np_ops.corner_to_surfaces_3d_jit(frustums) masks = points_in_convex_polygon_3d_jit(points, surfaces) points = points[masks.any(-1)] if self.remove_outside_points: assert calib is not None image_shape = res["metadata"]["image_shape"] points = box_np_ops.remove_outside_points(points, calib["rect"], calib["Trv2c"], calib["P2"], image_shape) if self.remove_environment is True and self.mode == "train": selected = keep_arrays_by_name(gt_names, target_assigner.classes) _dict_select(gt_dict, selected) masks = box_np_ops.points_in_rbbox(points, gt_dict["gt_boxes"]) points = points[masks.any(-1)] if self.mode == "train": selected = drop_arrays_by_name(gt_dict["gt_names"], ["DontCare", "ignore"]) _dict_select(gt_dict, selected) if self.remove_unknown: remove_mask = gt_dict["difficulty"] == -1 """ gt_boxes_remove = gt_boxes[remove_mask] gt_boxes_remove[:, 3:6] += 0.25 points = prep.remove_points_in_boxes(points, gt_boxes_remove) """ keep_mask = np.logical_not(remove_mask) _dict_select(gt_dict, keep_mask) gt_dict.pop("difficulty") if self.min_points_in_gt > 0: # points_count_rbbox takes 10ms with 10 sweeps nuscenes data point_counts = box_np_ops.points_count_rbbox( points, gt_dict["gt_boxes"]) mask = point_counts >= min_points_in_gt _dict_select(gt_dict, mask) gt_boxes_mask = np.array( [n in self.class_names for n in gt_dict["gt_names"]], dtype=np.bool_) if self.db_sampler: sampled_dict = self.db_sampler.sample_all( res["metadata"]["image_prefix"], gt_dict["gt_boxes"], gt_dict["gt_names"], res["metadata"]["num_point_features"], self.random_crop, gt_group_ids=None, calib=calib, road_planes=None # res["lidar"]["ground_plane"] ) if sampled_dict is not None: sampled_gt_names = sampled_dict["gt_names"] sampled_gt_boxes = sampled_dict["gt_boxes"] sampled_points = sampled_dict["points"] sampled_gt_masks = sampled_dict["gt_masks"] gt_dict["gt_names"] = np.concatenate( [gt_dict["gt_names"], sampled_gt_names], axis=0) gt_dict["gt_boxes"] = np.concatenate( [gt_dict["gt_boxes"], sampled_gt_boxes]) gt_boxes_mask = np.concatenate( [gt_boxes_mask, sampled_gt_masks], axis=0) if self.remove_points_after_sample: masks = box_np_ops.points_in_rbbox( points, sampled_gt_boxes) points = points[np.logical_not(masks.any(-1))] points = np.concatenate([sampled_points, points], axis=0) prep.noise_per_object_v3_( gt_dict["gt_boxes"], points, gt_boxes_mask, rotation_perturb=self.gt_rotation_noise, center_noise_std=self.gt_loc_noise_std, global_random_rot_range=self.global_random_rot_range, group_ids=None, num_try=100, ) _dict_select(gt_dict, gt_boxes_mask) gt_classes = np.array( [self.class_names.index(n) + 1 for n in gt_dict["gt_names"]], dtype=np.int32, ) gt_dict["gt_classes"] = gt_classes iskitti = res["type"] in ["KittiDataset"] if self.kitti_double: assert False, "No more KITTI" gt_dict["gt_boxes"], points = prep.random_flip_both( gt_dict["gt_boxes"], points, flip_coor=70.4 / 2) elif self.flip_single or iskitti: assert False, "nuscenes double flip is better" gt_dict["gt_boxes"], points = prep.random_flip( gt_dict["gt_boxes"], points) else: gt_dict["gt_boxes"], points = prep.random_flip_both( gt_dict["gt_boxes"], points) gt_dict["gt_boxes"], points = prep.global_rotation( gt_dict["gt_boxes"], points, rotation=self.global_rotation_noise) gt_dict["gt_boxes"], points = prep.global_scaling_v2( gt_dict["gt_boxes"], points, *self.global_scaling_noise) if self.shuffle_points: # shuffle is a little slow. np.random.shuffle(points) if self.mode == "train" and self.random_select: if self.npoints < points.shape[0]: pts_depth = points[:, 2] pts_near_flag = pts_depth < 40.0 far_idxs_choice = np.where(pts_near_flag == 0)[0] near_idxs = np.where(pts_near_flag == 1)[0] near_idxs_choice = np.random.choice(near_idxs, self.npoints - len(far_idxs_choice), replace=False) choice = (np.concatenate( (near_idxs_choice, far_idxs_choice), axis=0) if len(far_idxs_choice) > 0 else near_idxs_choice) np.random.shuffle(choice) else: choice = np.arange(0, len(points), dtype=np.int32) if self.npoints > len(points): extra_choice = np.random.choice(choice, self.npoints - len(points), replace=False) choice = np.concatenate((choice, extra_choice), axis=0) np.random.shuffle(choice) points = points[choice] if self.symmetry_intensity: points[:, -1] -= 0.5 # translate intensity to [-0.5, 0.5] # points[:, -1] *= 2 if self.normalize_intensity and res["type"] in ["NuScenesDataset"]: # print(points[:20, 3]) assert 0, "Velocity Accuracy drops 3 percent with normalization.." points[:, 3] /= 255 res["lidar"]["points"] = points if self.mode == "train": res["lidar"]["annotations"] = gt_dict return res, info
def noise_per_object_v3_( gt_boxes, points=None, valid_mask=None, rotation_perturb=np.pi / 4, center_noise_std=1.0,\ global_random_rot_range=np.pi / 4, num_try=5, group_ids=None,): """random rotate or remove each groundtruth independently. use kitti viewer to test this function points_transform_ Args: gt_boxes: [N, 7], gt box in lidar.points_transform_ points: [M, 4], point cloud in lidar. rotation_perturb: [-0.785, 0.785], center_noise_std: [1.0, 1.0, 0.5], global_random_rot_range: [0, 0] num_try: 100 """ num_boxes = gt_boxes.shape[0] if not isinstance(rotation_perturb, (list, tuple, np.ndarray)): # False rotation_perturb = [-rotation_perturb, rotation_perturb] if not isinstance(global_random_rot_range, (list, tuple, np.ndarray)): # False global_random_rot_range = [ -global_random_rot_range, global_random_rot_range ] enable_grot = ( np.abs(global_random_rot_range[0] - global_random_rot_range[1]) >= 1e-3 ) # False if not isinstance(center_noise_std, (list, tuple, np.ndarray)): # False center_noise_std = [ center_noise_std, center_noise_std, center_noise_std ] if valid_mask is None: # False valid_mask = np.ones((num_boxes, ), dtype=np.bool_) center_noise_std = np.array(center_noise_std, dtype=gt_boxes.dtype) loc_noises = np.random.normal(scale=center_noise_std, size=[num_boxes, num_try, 3]) rot_noises = np.random.uniform(rotation_perturb[0], rotation_perturb[1], size=[num_boxes, num_try]) gt_grots = np.arctan2(gt_boxes[:, 0], gt_boxes[:, 1]) grot_lowers = global_random_rot_range[0] - gt_grots grot_uppers = global_random_rot_range[1] - gt_grots global_rot_noises = np.random.uniform( grot_lowers[..., np.newaxis], grot_uppers[..., np.newaxis], size=[num_boxes, num_try], ) if group_ids is not None: # False if enable_grot: set_group_noise_same_v2_(loc_noises, rot_noises, global_rot_noises, group_ids) else: set_group_noise_same_(loc_noises, rot_noises, group_ids) group_centers, group_id_num_dict = get_group_center( gt_boxes[:, :3], group_ids) if enable_grot: group_transform_v2_( loc_noises, rot_noises, gt_boxes[:, :3], gt_boxes[:, 6], group_centers, global_rot_noises, valid_mask, ) else: group_transform_( loc_noises, rot_noises, gt_boxes[:, :3], gt_boxes[:, 6], group_centers, valid_mask, ) group_nums = np.array(list(group_id_num_dict.values()), dtype=np.int64) origin = [0.5, 0.5, 0.5] gt_box_corners = box_np_ops.center_to_corner_box3d(gt_boxes[:, :3], gt_boxes[:, 3:6], gt_boxes[:, 6], origin=origin, axis=2) if group_ids is not None: if not enable_grot: selected_noise = noise_per_box_group( gt_boxes[:, [0, 1, 3, 4, 6]], valid_mask, loc_noises, rot_noises, group_nums, ) else: selected_noise = noise_per_box_group_v2_( gt_boxes[:, [0, 1, 3, 4, 6]], valid_mask, loc_noises, rot_noises, group_nums, global_rot_noises, ) else: if not enable_grot: # True selected_noise = noise_per_box( gt_boxes[:, [0, 1, 3, 4, 6]], valid_mask, loc_noises, rot_noises) # todo: seems without noise in z-axis velo else: selected_noise = noise_per_box_v2_( gt_boxes[:, [0, 1, 3, 4, 6]], valid_mask, loc_noises, rot_noises, global_rot_noises, ) loc_transforms = _select_transform(loc_noises, selected_noise) rot_transforms = _select_transform(rot_noises, selected_noise) surfaces = box_np_ops.corner_to_surfaces_3d_jit(gt_box_corners) if points is not None: point_masks = points_in_convex_polygon_3d_jit(points[:, :3], surfaces) points_transform_( points, gt_boxes[:, :3], point_masks, loc_transforms, rot_transforms, valid_mask, ) box3d_transform_(gt_boxes, loc_transforms, rot_transforms, valid_mask)
def __call__(self, res, info, res_1, info_1, res_2, info_2): res["mode"] = self.mode res_1["mode"] = self.mode res_2["mode"] = self.mode if res["type"] in ["KittiDataset", "LyftDataset", "LvxDataset"]: points = res["lidar"]["points"] points_1 = res_1["lidar"]["points"] points_2 = res_2["lidar"]["points"] elif res["type"] == "NuScenesDataset": points = res["lidar"]["combined"] if self.mode == "train": anno_dict = res["lidar"]["annotations"] gt_dict = { "gt_boxes": anno_dict["boxes"], # 加入之前之后一帧的boxes "gt_boxes_1": anno_dict["boxes_1"], "gt_boxes_2": anno_dict["boxes_2"], "gt_boxes_3": anno_dict["boxes_3"], "gt_names": np.array(anno_dict["names"]).reshape(-1), } if "difficulty" not in anno_dict: difficulty = np.zeros([anno_dict["boxes"].shape[0]], dtype=np.int32) gt_dict["difficulty"] = difficulty else: gt_dict["difficulty"] = anno_dict["difficulty"] if "calib" in res: calib = res["calib"] else: calib = None if self.add_rgb_to_points: assert calib is not None and "image" in res image_path = res["image"]["image_path"] image = (imgio.imread(str(pathlib.Path(root_path) / image_path)).astype(np.float32) / 255) points_rgb = box_np_ops.add_rgb_to_points(points, image, calib["rect"], calib["Trv2c"], calib["P2"]) points = np.concatenate([points, points_rgb], axis=1) num_point_features += 3 if self.reference_detections is not None: assert calib is not None and "image" in res C, R, T = box_np_ops.projection_matrix_to_CRT_kitti(P2) frustums = box_np_ops.get_frustum_v2(reference_detections, C) frustums -= T frustums = np.einsum("ij, akj->aki", np.linalg.inv(R), frustums) frustums = box_np_ops.camera_to_lidar(frustums, rect, Trv2c) surfaces = box_np_ops.corner_to_surfaces_3d_jit(frustums) masks = points_in_convex_polygon_3d_jit(points, surfaces) points = points[masks.any(-1)] if self.remove_outside_points: assert calib is not None image_shape = res["image"]["image_shape"] points = box_np_ops.remove_outside_points(points, calib["rect"], calib["Trv2c"], calib["P2"], image_shape) if self.remove_environment is True and self.mode == "train": selected = kitti.keep_arrays_by_name(gt_names, target_assigner.classes) _dict_select(gt_dict, selected) masks = box_np_ops.points_in_rbbox(points, gt_dict["gt_boxes"]) points = points[masks.any(-1)] if self.mode == "train": selected = kitti.drop_arrays_by_name(gt_dict["gt_names"], ["DontCare", "ignore"]) _dict_select(gt_dict, selected) if self.remove_unknown: remove_mask = gt_dict["difficulty"] == -1 """ gt_boxes_remove = gt_boxes[remove_mask] gt_boxes_remove[:, 3:6] += 0.25 points = prep.remove_points_in_boxes(points, gt_boxes_remove) """ keep_mask = np.logical_not(remove_mask) _dict_select(gt_dict, keep_mask) gt_dict.pop("difficulty") # 保证每个box内点的数量 if self.min_points_in_gt > 0: # points_count_rbbox takes 10ms with 10 sweeps nuscenes data point_counts = box_np_ops.points_count_rbbox( points, gt_dict["gt_boxes"]) mask = point_counts >= min_points_in_gt _dict_select(gt_dict, mask) gt_boxes_mask = np.array( [n in self.class_names for n in gt_dict["gt_names"]], dtype=np.bool_) # 数据增广暂时去掉,有点改不过来 if self.db_sampler: pass assert NotImplementedError sampled_dict = self.db_sampler.sample_all( res["metadata"]["image_prefix"], gt_dict["gt_boxes"], gt_dict["gt_names"], res["metadata"]["num_point_features"], self.random_crop, gt_group_ids=None, calib=calib, ) if sampled_dict is not None: sampled_gt_names = sampled_dict["gt_names"] sampled_gt_boxes = sampled_dict["gt_boxes"] sampled_points = sampled_dict["points"] sampled_gt_masks = sampled_dict["gt_masks"] gt_dict["gt_names"] = np.concatenate( [gt_dict["gt_names"], sampled_gt_names], axis=0) gt_dict["gt_boxes"] = np.concatenate( [gt_dict["gt_boxes"], sampled_gt_boxes]) gt_boxes_mask = np.concatenate( [gt_boxes_mask, sampled_gt_masks], axis=0) if self.remove_points_after_sample: masks = box_np_ops.points_in_rbbox( points, sampled_gt_boxes) points = points[np.logical_not(masks.any(-1))] points = np.concatenate([sampled_points, points], axis=0) # 暂时不加扰动了 # TODO:需要改4帧,t-2,t-1,t的点云,以及t+1的box prep.noise_per_object_v3_( gt_dict["gt_boxes"], gt_dict["gt_boxes_1"], gt_dict["gt_boxes_2"], gt_dict["gt_boxes_3"], points, points_1, points_2, gt_boxes_mask, rotation_perturb=self.gt_rotation_noise, center_noise_std=self.gt_loc_noise_std, global_random_rot_range=self.global_random_rot_range, group_ids=None, num_try=100, ) _dict_select(gt_dict, gt_boxes_mask) # 将类别转换为index gt_classes = np.array( [self.class_names.index(n) + 1 for n in gt_dict["gt_names"]], dtype=np.int32, ) gt_dict["gt_classes"] = gt_classes # 添加了三帧的全局变换 gt_dict["gt_boxes"], gt_dict["gt_boxes_1"], gt_dict[ "gt_boxes_2"], points, points_1, points_2 = prep.random_flip( gt_dict["gt_boxes"], gt_dict["gt_boxes_1"], gt_dict["gt_boxes_2"], points, points_1, points_2) gt_dict["gt_boxes"], gt_dict["gt_boxes_1"], gt_dict[ "gt_boxes_2"], points, points_1, points_2 = prep.global_rotation( gt_dict["gt_boxes"], gt_dict["gt_boxes_1"], gt_dict["gt_boxes_2"], points, points_1, points_2, rotation=self.global_rotation_noise) gt_dict["gt_boxes"], gt_dict["gt_boxes_1"], gt_dict[ "gt_boxes_2"], points, points_1, points_2 = prep.global_scaling_v2( gt_dict["gt_boxes"], gt_dict["gt_boxes_1"], gt_dict["gt_boxes_2"], points, points_1, points_2, *self.global_scaling_noise) if self.shuffle_points: # shuffle is a little slow. np.random.shuffle(points) np.random.shuffle(points_1) np.random.shuffle(points_2) if self.mode == "train" and self.random_select: # 进行点的随机采样,暂时不需要 if self.npoints < points.shape[0]: pts_depth = points[:, 2] pts_near_flag = pts_depth < 40.0 far_idxs_choice = np.where(pts_near_flag == 0)[0] near_idxs = np.where(pts_near_flag == 1)[0] near_idxs_choice = np.random.choice(near_idxs, self.npoints - len(far_idxs_choice), replace=False) choice = (np.concatenate( (near_idxs_choice, far_idxs_choice), axis=0) if len(far_idxs_choice) > 0 else near_idxs_choice) np.random.shuffle(choice) else: choice = np.arange(0, len(points), dtype=np.int32) if self.npoints > len(points): extra_choice = np.random.choice(choice, self.npoints - len(points), replace=False) choice = np.concatenate((choice, extra_choice), axis=0) np.random.shuffle(choice) points = points[choice] if self.symmetry_intensity: points[:, -1] -= 0.5 # translate intensity to [-0.5, 0.5] # points[:, -1] *= 2 res["lidar"]["points"] = points res_1["lidar"]["points"] = points_1 res_2["lidar"]["points"] = points_2 if self.mode == "train": import copy res["lidar"]["annotations"] = gt_dict res_1["lidar"]["annotations"] = copy.deepcopy(gt_dict) res_2["lidar"]["annotations"] = copy.deepcopy(gt_dict) return res, info, res_1,info_1,res_2,info_2
def noise_per_object_v3_( gt_boxes, gt_boxes_1, gt_boxes_2, gt_boxes_3, points=None, points_1=None, points_2=None, valid_mask=None, rotation_perturb=np.pi / 4, center_noise_std=1.0, global_random_rot_range=np.pi / 4, num_try=5, group_ids=None, ): """random rotate or remove each groundtruth independently. use kitti viewer to test this function points_transform_ Args: gt_boxes: [N, 7], gt box in lidar.points_transform_ points: [M, 4], point cloud in lidar. """ num_boxes = gt_boxes.shape[0] if not isinstance(rotation_perturb, (list, tuple, np.ndarray)): rotation_perturb = [-rotation_perturb, rotation_perturb] if not isinstance(global_random_rot_range, (list, tuple, np.ndarray)): global_random_rot_range = [ -global_random_rot_range, global_random_rot_range ] enable_grot = (np.abs(global_random_rot_range[0] - global_random_rot_range[1]) >= 1e-3) if not isinstance(center_noise_std, (list, tuple, np.ndarray)): center_noise_std = [ center_noise_std, center_noise_std, center_noise_std ] if valid_mask is None: valid_mask = np.ones((num_boxes, ), dtype=np.bool_) center_noise_std = np.array(center_noise_std, dtype=gt_boxes.dtype) loc_noises = np.random.normal(scale=center_noise_std, size=[num_boxes, num_try, 3]) # loc_noises = np.random.uniform( # -center_noise_std, center_noise_std, size=[num_boxes, num_try, 3]) rot_noises = np.random.uniform(rotation_perturb[0], rotation_perturb[1], size=[num_boxes, num_try]) gt_grots = np.arctan2(gt_boxes[:, 0], gt_boxes[:, 1]) grot_lowers = global_random_rot_range[0] - gt_grots grot_uppers = global_random_rot_range[1] - gt_grots global_rot_noises = np.random.uniform( grot_lowers[..., np.newaxis], grot_uppers[..., np.newaxis], size=[num_boxes, num_try], ) if group_ids is not None: if enable_grot: set_group_noise_same_v2_(loc_noises, rot_noises, global_rot_noises, group_ids) else: set_group_noise_same_(loc_noises, rot_noises, group_ids) group_centers, group_id_num_dict = get_group_center( gt_boxes[:, :3], group_ids) if enable_grot: group_transform_v2_( loc_noises, rot_noises, gt_boxes[:, :3], gt_boxes[:, 6], group_centers, global_rot_noises, valid_mask, ) else: group_transform_( loc_noises, rot_noises, gt_boxes[:, :3], gt_boxes[:, 6], group_centers, valid_mask, ) group_nums = np.array(list(group_id_num_dict.values()), dtype=np.int64) origin = [0.5, 0.5, 0.5] # TODO: 3个gt_box gt_box_corners = box_np_ops.center_to_corner_box3d(gt_boxes[:, :3], gt_boxes[:, 3:6], gt_boxes[:, 6], origin=origin, axis=2) # T-1 gt_box_corners_1 = box_np_ops.center_to_corner_box3d(gt_boxes_1[:, :3], gt_boxes_1[:, 3:6], gt_boxes_1[:, 6], origin=origin, axis=2) # T+1 gt_box_corners_2 = box_np_ops.center_to_corner_box3d(gt_boxes_2[:, :3], gt_boxes_2[:, 3:6], gt_boxes_2[:, 6], origin=origin, axis=2) # T-2 gt_box_corners_3 = box_np_ops.center_to_corner_box3d(gt_boxes_3[:, :3], gt_boxes_3[:, 3:6], gt_boxes_3[:, 6], origin=origin, axis=2) if group_ids is not None: if not enable_grot: selected_noise = noise_per_box_group( gt_boxes[:, [0, 1, 3, 4, 6]], valid_mask, loc_noises, rot_noises, group_nums, ) else: selected_noise = noise_per_box_group_v2_( gt_boxes[:, [0, 1, 3, 4, 6]], valid_mask, loc_noises, rot_noises, group_nums, global_rot_noises, ) else: if not enable_grot: selected_noise = noise_per_box(gt_boxes[:, [0, 1, 3, 4, 6]], gt_boxes_1[:, [0, 1, 3, 4, 6]], gt_boxes_2[:, [0, 1, 3, 4, 6]], gt_boxes_3[:, [0, 1, 3, 4, 6]], valid_mask, loc_noises, rot_noises) else: selected_noise = noise_per_box_v2_( gt_boxes[:, [0, 1, 3, 4, 6]], valid_mask, loc_noises, rot_noises, global_rot_noises, ) loc_transforms = _select_transform(loc_noises, selected_noise) rot_transforms = _select_transform(rot_noises, selected_noise) surfaces = box_np_ops.corner_to_surfaces_3d_jit(gt_box_corners) if points is not None: point_masks = points_in_convex_polygon_3d_jit(points[:, :3], surfaces) points_transform_( points, gt_boxes[:, :3], point_masks, loc_transforms, rot_transforms, valid_mask, ) # T-1 帧的增广 loc_transforms_1 = get_loc_transform(loc_transforms, rot_transforms, gt_boxes, gt_boxes_1) surfaces = box_np_ops.corner_to_surfaces_3d_jit(gt_box_corners_1) if points_1 is not None: point_masks = points_in_convex_polygon_3d_jit(points_1[:, :3], surfaces) points_transform_( points_1, gt_boxes_1[:, :3], point_masks, loc_transforms_1, rot_transforms, valid_mask, ) box3d_transform_(gt_boxes_1, loc_transforms_1, rot_transforms, valid_mask) # T+1帧的增广 loc_transforms_2 = get_loc_transform(loc_transforms, rot_transforms, gt_boxes, gt_boxes_2) box3d_transform_(gt_boxes_2, loc_transforms_2, rot_transforms, valid_mask) # T-2帧的增广 loc_transforms_3 = get_loc_transform(loc_transforms, rot_transforms, gt_boxes, gt_boxes_3) if points_2 is not None: point_masks = points_in_convex_polygon_3d_jit(points_2[:, :3], surfaces) points_transform_( points_2, gt_boxes_3[:, :3], point_masks, loc_transforms_3, rot_transforms, valid_mask, ) # T帧的box需要做为基准,所以最后变换 box3d_transform_(gt_boxes, loc_transforms, rot_transforms, valid_mask)