def _calculate_num_points_in_gt(data_path, infos, relative_path, remove_outside=False, num_features=3): ''' 计算bounding box内点的数量 ''' for info in infos: pc_info = info["point_cloud"] if relative_path: v_path = str(Path(data_path) / pc_info["velodyne_path"]) else: v_path = pc_info["velodyne_path"] pcd = o3d.io.read_point_cloud(str(v_path)) points_v = np.asarray(pcd.points) if num_features >= 4: normals_v = np.asarray(pcd.normals) points_v = np.concatenate([points_v, normals_v], axis=1)[:, :num_features] annos = info["annos"] dims = annos["dimensions"] loc = annos["location"] rots = annos["rotation_y"] gt_boxes = np.concatenate([loc, dims, rots[..., np.newaxis]], axis=1) indices = box_np_ops.points_in_rbbox(points_v[:, :3], gt_boxes) num_points_in_gt = indices.sum(0) annos["num_points_in_gt"] = num_points_in_gt
def _calculate_num_points_in_gt(data_path, infos, relative_path, remove_outside=True, num_features=4): for info in infos: pc_info = info["point_cloud"] image_info = info["image"] calib = info["calib"] if relative_path: v_path = str(Path(data_path) / pc_info["velodyne_path"]) else: v_path = pc_info["velodyne_path"] points_v = np.fromfile(v_path, dtype=np.float32, count=-1).reshape([-1, num_features]) rect = calib["R0_rect"] Trv2c = calib["Tr_velo_to_cam"] P2 = calib["P2"] if remove_outside: points_v = box_np_ops.remove_outside_points( points_v, rect, Trv2c, P2, image_info["image_shape"]) # points_v = points_v[points_v[:, 0] > 0] annos = info["annos"] num_obj = len([n for n in annos["name"] if n != "DontCare"]) # annos = kitti.filter_kitti_anno(annos, ['DontCare']) dims = annos["dimensions"][:num_obj] loc = annos["location"][:num_obj] rots = annos["rotation_y"][:num_obj] gt_boxes_camera = np.concatenate([loc, dims, rots[..., np.newaxis]], axis=1) gt_boxes_lidar = box_np_ops.box_camera_to_lidar( gt_boxes_camera, rect, Trv2c) indices = box_np_ops.points_in_rbbox(points_v[:, :3], gt_boxes_lidar) num_points_in_gt = indices.sum(0) num_ignored = len(annos["dimensions"]) - num_obj num_points_in_gt = np.concatenate( [num_points_in_gt, -np.ones([num_ignored])]) annos["num_points_in_gt"] = num_points_in_gt.astype(np.int32)
def create_groundtruth_database( dataset_class_name, data_path, info_path=None, used_classes=None, db_path=None, dbinfo_path=None, relative_path=True, add_rgb=False, lidar_only=False, bev_only=False, coors_range=None, **kwargs, ): pipeline = [ { "type": "LoadPointCloudFromFile", "dataset": dataset_name_map[dataset_class_name], }, { "type": "LoadPointCloudAnnotations", "with_bbox": True }, ] if "nsweeps" in kwargs: dataset = get_dataset(dataset_class_name)( info_path=info_path, root_path=data_path, pipeline=pipeline, test_mode=True, nsweeps=kwargs["nsweeps"], ) nsweeps = dataset.nsweeps else: dataset = get_dataset(dataset_class_name)(info_path=info_path, root_path=data_path, test_mode=True, pipeline=pipeline) nsweeps = 1 root_path = Path(data_path) if dataset_class_name == "NUSC": if db_path is None: db_path = root_path / f"gt_database_{nsweeps}sweeps_withvelo" if dbinfo_path is None: dbinfo_path = root_path / f"dbinfos_train_{nsweeps}sweeps_withvelo.pkl" else: if db_path is None: db_path = root_path / "gt_database" if dbinfo_path is None: dbinfo_path = root_path / "dbinfos_train.pkl" if dataset_class_name == "NUSC" or dataset_class_name == "LYFT": point_features = 5 elif dataset_class_name == "KITTI": point_features = 4 db_path.mkdir(parents=True, exist_ok=True) all_db_infos = {} group_counter = 0 # def prepare_single_data(index): for index in tqdm(range(len(dataset))): image_idx = index # modified to nuscenes sensor_data = dataset.get_sensor_data(index) # for nsweep, sensor_data in enumerate(sensor_datas): if "image_idx" in sensor_data["metadata"]: image_idx = sensor_data["metadata"]["image_idx"] if dataset_class_name == "NUSC": points = sensor_data["lidar"]["combined"] elif dataset_class_name == "KITTI": points = sensor_data["lidar"]["points"] elif dataset_class_name == "LYFT": points = sensor_data["lidar"]["points"] print(points.shape) annos = sensor_data["lidar"]["annotations"] gt_boxes = annos["boxes"] names = annos["names"] group_dict = {} group_ids = np.full([gt_boxes.shape[0]], -1, dtype=np.int64) if "group_ids" in annos: group_ids = annos["group_ids"] else: group_ids = np.arange(gt_boxes.shape[0], dtype=np.int64) difficulty = np.zeros(gt_boxes.shape[0], dtype=np.int32) if "difficulty" in annos: difficulty = annos["difficulty"] num_obj = gt_boxes.shape[0] point_indices = box_np_ops.points_in_rbbox(points, gt_boxes) for i in range(num_obj): filename = f"{image_idx}_{names[i]}_{i}.bin" filepath = db_path / filename gt_points = points[point_indices[:, i]] gt_points[:, :3] -= gt_boxes[i, :3] with open(filepath, "w") as f: gt_points[:, :point_features].tofile(f) if (used_classes is None) or names[i] in used_classes: if relative_path: db_dump_path = str(db_path.stem + "/" + filename) else: db_dump_path = str(filepath) db_info = { "name": names[i], "path": db_dump_path, "image_idx": image_idx, "gt_idx": i, "box3d_lidar": gt_boxes[i], "num_points_in_gt": gt_points.shape[0], "difficulty": difficulty[i], # "group_id": -1, # "bbox": bboxes[i], } local_group_id = group_ids[i] # if local_group_id >= 0: if local_group_id not in group_dict: group_dict[local_group_id] = group_counter group_counter += 1 db_info["group_id"] = group_dict[local_group_id] if "score" in annos: db_info["score"] = annos["score"][i] if names[i] in all_db_infos: all_db_infos[names[i]].append(db_info) else: all_db_infos[names[i]] = [db_info] # print(f"Finish {index}th sample") print("dataset length: ", len(dataset)) for k, v in all_db_infos.items(): print(f"load {len(v)} {k} database infos") with open(dbinfo_path, "wb") as f: pickle.dump(all_db_infos, f)
def create_groundtruth_database( dataset_class_name, data_path, info_path=None, used_classes=None, db_path=None, dbinfo_path=None, relative_path=True, add_rgb=False, lidar_only=False, bev_only=False, coors_range=None, gt_aug_with_context=-1.0, **kwargs, ): gt_aug_with_context = gt_aug_with_context pipeline = [ { "type": "LoadPointCloudFromFile", "dataset": dataset_name_map[dataset_class_name], }, { "type": "LoadPointCloudAnnotations", "with_bbox": True, "enable_difficulty_level": True }, ] # get KittiDataset loaded points and annos. dataset = get_dataset(dataset_class_name)(info_path=info_path, root_path=data_path, test_mode=True, pipeline=pipeline) # prepare dbinfo_path and db_path. root_path = Path(data_path) if db_path is None: db_path = root_path / "gt_database" if dbinfo_path is None: dbinfo_path = root_path / "dbinfos_train.pkl" db_path.mkdir(parents=True, exist_ok=True) all_db_infos = {} group_counter = 0 for index in tqdm(range(len(dataset))): image_idx = index sensor_data = dataset.get_sensor_data( index) # see in loading.py after pipelines. if "image_idx" in sensor_data[ "metadata"]: # True, image_idx = file_name (000001) image_idx = sensor_data["metadata"]["image_idx"] points = sensor_data["lidar"]["points"] annos = sensor_data["lidar"]["annotations"] gt_boxes = annos[ "boxes"] # gt_boxes of all classes, dc has been removed in pipeline LoadPointCloudAnnotations. names = annos["names"] # gt_names. group_dict = {} group_ids = np.full([gt_boxes.shape[0]], -1, dtype=np.int64) if "group_ids" in annos: # False group_ids = annos["group_ids"] else: group_ids = np.arange(gt_boxes.shape[0], dtype=np.int64) difficulty = np.zeros(gt_boxes.shape[0], dtype=np.int32) if "difficulty" in annos: # False difficulty = annos["difficulty"] num_obj = gt_boxes.shape[0] # todo: maybe we need add some contexual points here. offset = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] # [x, y, z, w, l, h, ry] if gt_aug_with_context > 0.0: offset = [ 0.0, 0.0, 0.0, gt_aug_with_context, gt_aug_with_context, 0.0, 0.0 ] db_path = root_path / "gt_enlarged_database" dbinfo_path = root_path / "dbinfos_enlarged_train.pkl" db_path.mkdir(parents=True, exist_ok=True) point_indices_for_num = box_np_ops.points_in_rbbox(points, gt_boxes) point_indices = box_np_ops.points_in_rbbox(points, gt_boxes + offset) for i in range(num_obj): # in a single scene. filename = f"{image_idx}_{names[i]}_{i}.bin" filepath = db_path / filename gt_points = points[point_indices[:, i]] num_points_in_gt = point_indices_for_num[:, i].sum() gt_points[:, :3] -= gt_boxes[i, : 3] # only record relative distance with open(filepath, "w") as f: # db: gt points in each gt_box are saved gt_points[:, :4].tofile(f) if (used_classes is None) or names[i] in used_classes: if relative_path: db_dump_path = str(db_path.stem + "/" + filename) else: db_dump_path = str(filepath) db_info = { "name": names[i], "path": db_dump_path, "image_idx": image_idx, "gt_idx": i, "box3d_lidar": gt_boxes[i], "num_points_in_gt": num_points_in_gt, "difficulty": difficulty[i] # todo: not accurate, all are set as 0. } local_group_id = group_ids[i] if local_group_id not in group_dict: group_dict[local_group_id] = group_counter group_counter += 1 db_info["group_id"] = group_dict[ local_group_id] # count from 0 to total_num_of_specific_class[like 13442 for car] if "score" in annos: # False db_info["score"] = annos["score"][i] if names[ i] in all_db_infos: # all_db_infos are grouped by class_names (like car, pedestrian, cyclist, cycle) all_db_infos[names[i]].append( db_info) # all db infos include info of all db else: all_db_infos[names[i]] = [db_info] print("dataset length: ", len(dataset)) for k, v in all_db_infos.items(): print(f"load {len(v)} {k} database infos") with open(dbinfo_path, "wb") as f: pickle.dump(all_db_infos, f)
def create_groundtruth_database( dataset_class_name, data_path, info_path=None, used_classes=None, db_path=None, dbinfo_path=None, relative_path=True, **kwargs, ): pipeline = [ { "type": "LoadPointCloudFromFile", "dataset": dataset_name_map[dataset_class_name], }, { "type": "LoadPointCloudAnnotations", "with_bbox": True }, ] if "nsweeps" in kwargs: dataset = get_dataset(dataset_class_name)( info_path=info_path, root_path=data_path, pipeline=pipeline, test_mode=True, nsweeps=kwargs["nsweeps"], ) nsweeps = dataset.nsweeps else: dataset = get_dataset(dataset_class_name)(info_path=info_path, root_path=data_path, test_mode=True, pipeline=pipeline) nsweeps = 1 root_path = Path(data_path) if dataset_class_name in ["WAYMO", "NUSC"]: if db_path is None: db_path = root_path / f"gt_database_{nsweeps}sweeps_withvelo" if dbinfo_path is None: dbinfo_path = root_path / f"dbinfos_train_{nsweeps}sweeps_withvelo.pkl" else: raise NotImplementedError() if dataset_class_name == "NUSC": point_features = 5 elif dataset_class_name == "WAYMO": point_features = 5 if nsweeps == 1 else 6 else: raise NotImplementedError() db_path.mkdir(parents=True, exist_ok=True) all_db_infos = {} group_counter = 0 for index in tqdm(range(len(dataset))): image_idx = index # modified to nuscenes sensor_data = dataset.get_sensor_data(index) if "image_idx" in sensor_data["metadata"]: image_idx = sensor_data["metadata"]["image_idx"] if nsweeps > 1: points = sensor_data["lidar"]["combined"] else: points = sensor_data["lidar"]["points"] annos = sensor_data["lidar"]["annotations"] gt_boxes = annos["boxes"] names = annos["names"] if dataset_class_name == 'WAYMO': # waymo dataset contains millions of objects and it is not possible to store # all of them into a single folder # we randomly sample a few objects for gt augmentation # We keep all cyclist as they are rare if index % 4 != 0: mask = (names == 'VEHICLE') mask = np.logical_not(mask) names = names[mask] gt_boxes = gt_boxes[mask] if index % 2 != 0: mask = (names == 'PEDESTRIAN') mask = np.logical_not(mask) names = names[mask] gt_boxes = gt_boxes[mask] group_dict = {} group_ids = np.full([gt_boxes.shape[0]], -1, dtype=np.int64) if "group_ids" in annos: group_ids = annos["group_ids"] else: group_ids = np.arange(gt_boxes.shape[0], dtype=np.int64) difficulty = np.zeros(gt_boxes.shape[0], dtype=np.int32) if "difficulty" in annos: difficulty = annos["difficulty"] num_obj = gt_boxes.shape[0] if num_obj == 0: continue point_indices = box_np_ops.points_in_rbbox(points, gt_boxes) for i in range(num_obj): if (used_classes is None) or names[i] in used_classes: filename = f"{image_idx}_{names[i]}_{i}.bin" dirpath = os.path.join(str(db_path), names[i]) os.makedirs(dirpath, exist_ok=True) filepath = os.path.join(str(db_path), names[i], filename) gt_points = points[point_indices[:, i]] gt_points[:, :3] -= gt_boxes[i, :3] with open(filepath, "w") as f: try: gt_points[:, :point_features].tofile(f) except: print("process {} files".format(index)) break if (used_classes is None) or names[i] in used_classes: if relative_path: db_dump_path = os.path.join(db_path.stem, names[i], filename) else: db_dump_path = str(filepath) db_info = { "name": names[i], "path": db_dump_path, "image_idx": image_idx, "gt_idx": i, "box3d_lidar": gt_boxes[i], "num_points_in_gt": gt_points.shape[0], "difficulty": difficulty[i], # "group_id": -1, # "bbox": bboxes[i], } local_group_id = group_ids[i] # if local_group_id >= 0: if local_group_id not in group_dict: group_dict[local_group_id] = group_counter group_counter += 1 db_info["group_id"] = group_dict[local_group_id] if "score" in annos: db_info["score"] = annos["score"][i] if names[i] in all_db_infos: all_db_infos[names[i]].append(db_info) else: all_db_infos[names[i]] = [db_info] print("dataset length: ", len(dataset)) for k, v in all_db_infos.items(): print(f"load {len(v)} {k} database infos") with open(dbinfo_path, "wb") as f: pickle.dump(all_db_infos, f)
def prep_pointcloud_rpn( input_dict, root_path, task_class_names=[], prep_cfg=None, db_sampler=None, remove_outside_points=False, training=True, num_point_features=4, random_crop=False, reference_detections=None, out_dtype=np.float32, min_points_in_gt=-1, logger=None, ): """ convert point cloud to voxels, create targets if ground truths exists. input_dict format: dataset.get_sensor_data format """ assert prep_cfg is not None remove_environment = prep_cfg.REMOVE_UNKOWN_EXAMPLES if training: remove_unknown = prep_cfg.REMOVE_UNKOWN_EXAMPLES gt_rotation_noise = prep_cfg.GT_ROT_NOISE gt_loc_noise_std = prep_cfg.GT_LOC_NOISE global_rotation_noise = prep_cfg.GLOBAL_ROT_NOISE global_scaling_noise = prep_cfg.GLOBAL_SCALE_NOISE global_random_rot_range = prep_cfg.GLOBAL_ROT_PER_OBJ_RANGE global_translate_noise_std = prep_cfg.GLOBAL_TRANS_NOISE gt_points_drop = prep_cfg.GT_DROP_PERCENTAGE gt_drop_max_keep = prep_cfg.GT_DROP_MAX_KEEP_POINTS remove_points_after_sample = prep_cfg.REMOVE_POINTS_AFTER_SAMPLE class_names = list(itertools.chain(*task_class_names)) # points_only = input_dict["lidar"]["points"] # times = input_dict["lidar"]["times"] # points = np.hstack([points_only, times]) points = input_dict["lidar"]["points"] if training: anno_dict = input_dict["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 use_group_id and "group_ids" in anno_dict: # group_ids = anno_dict["group_ids"] # gt_dict["group_ids"] = group_ids calib = None if "calib" in input_dict: calib = input_dict["calib"] if reference_detections is not None: assert calib is not None and "image" in input_dict 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 remove_outside_points: assert calib is not None image_shape = input_dict["image"]["image_shape"] points = box_np_ops.remove_outside_points(points, calib["rect"], calib["Trv2c"], calib["P2"], image_shape) if remove_environment is True and training: 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 training: selected = kitti.drop_arrays_by_name(gt_dict["gt_names"], ["DontCare", "ignore"]) _dict_select(gt_dict, selected) if 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") gt_boxes_mask = np.array( [n in class_names for n in gt_dict["gt_names"]], dtype=np.bool_) # db_sampler = None if db_sampler is not None: group_ids = None # if "group_ids" in gt_dict: # group_ids = gt_dict["group_ids"] sampled_dict = db_sampler.sample_all( root_path, gt_dict["gt_boxes"], gt_dict["gt_names"], num_point_features, random_crop, gt_group_ids=group_ids, 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 group_ids is not None: # sampled_group_ids = sampled_dict["group_ids"] # gt_dict["group_ids"] = np.concatenate( # [gt_dict["group_ids"], sampled_group_ids]) if 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) # group_ids = None # if "group_ids" in gt_dict: # group_ids = gt_dict["group_ids"] prep.noise_per_object_v3_( gt_dict["gt_boxes"], points, gt_boxes_mask, rotation_perturb=gt_rotation_noise, center_noise_std=gt_loc_noise_std, global_random_rot_range=global_random_rot_range, group_ids=None, num_try=100, ) # should remove unrelated objects after noise per object # for k, v in gt_dict.items(): # print(k, v.shape) _dict_select(gt_dict, gt_boxes_mask) gt_classes = np.array( [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"], points = prep.random_flip(gt_dict["gt_boxes"], points) gt_dict["gt_boxes"], points = prep.global_rotation( gt_dict["gt_boxes"], points, rotation=global_rotation_noise) gt_dict["gt_boxes"], points = prep.global_scaling_v2( gt_dict["gt_boxes"], points, *global_scaling_noise) prep.global_translate_(gt_dict["gt_boxes"], points, global_translate_noise_std) task_masks = [] flag = 0 for class_name in task_class_names: task_masks.append([ np.where(gt_dict["gt_classes"] == class_name.index(i) + 1 + flag) for i in class_name ]) flag += len(class_name) task_boxes = [] task_classes = [] task_names = [] flag2 = 0 for idx, mask in enumerate(task_masks): task_box = [] task_class = [] task_name = [] for m in mask: task_box.append(gt_dict["gt_boxes"][m]) task_class.append(gt_dict["gt_classes"][m] - flag2) task_name.append(gt_dict["gt_names"][m]) task_boxes.append(np.concatenate(task_box, axis=0)) task_classes.append(np.concatenate(task_class)) task_names.append(np.concatenate(task_name)) flag2 += len(mask) for task_box in task_boxes: # limit rad to [-pi, pi] task_box[:, -1] = box_np_ops.limit_period(task_box[:, -1], offset=0.5, period=2 * np.pi) # print(gt_dict.keys()) gt_dict["gt_classes"] = task_classes gt_dict["gt_names"] = task_names gt_dict["gt_boxes"] = task_boxes example = { "pts_input": points, "pts_rect": None, "pts_features": None, "gt_boxes3d": gt_dict["gt_boxes"], "rpn_cls_label": [], "rpn_reg_label": [], } if calib is not None: example["calib"] = calib return example
def prep_sequence_pointcloud( input_dict, root_path, voxel_generator, target_assigners, prep_cfg=None, db_sampler=None, remove_outside_points=False, training=True, create_targets=True, num_point_features=4, anchor_cache=None, random_crop=False, reference_detections=None, out_size_factor=2, out_dtype=np.float32, min_points_in_gt=-1, logger=None, ): """ convert point cloud to voxels, create targets if ground truths exists. input_dict format: dataset.get_sensor_data format """ assert prep_cfg is not None remove_environment = prep_cfg.REMOVE_ENVIRONMENT max_voxels = prep_cfg.MAX_VOXELS_NUM shuffle_points = prep_cfg.SHUFFLE anchor_area_threshold = prep_cfg.ANCHOR_AREA_THRES if training: remove_unknown = prep_cfg.REMOVE_UNKOWN_EXAMPLES gt_rotation_noise = prep_cfg.GT_ROT_NOISE gt_loc_noise_std = prep_cfg.GT_LOC_NOISE global_rotation_noise = prep_cfg.GLOBAL_ROT_NOISE global_scaling_noise = prep_cfg.GLOBAL_SCALE_NOISE global_random_rot_range = prep_cfg.GLOBAL_ROT_PER_OBJ_RANGE global_translate_noise_std = prep_cfg.GLOBAL_TRANS_NOISE gt_points_drop = prep_cfg.GT_DROP_PERCENTAGE gt_drop_max_keep = prep_cfg.GT_DROP_MAX_KEEP_POINTS remove_points_after_sample = prep_cfg.REMOVE_POINTS_AFTER_SAMPLE min_points_in_gt = prep_cfg.get("MIN_POINTS_IN_GT", -1) task_class_names = [ target_assigner.classes for target_assigner in target_assigners ] class_names = list(itertools.chain(*task_class_names)) # points_only = input_dict["lidar"]["points"] # times = input_dict["lidar"]["times"] # points = np.hstack([points_only, times]) try: points = input_dict["current_frame"]["lidar"]["combined"] except Exception: points = input_dict["current_frame"]["lidar"]["points"] keyframe_points = input_dict["keyframe"]["lidar"]["combined"] if training: anno_dict = input_dict["current_frame"]["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 use_group_id and "group_ids" in anno_dict: # group_ids = anno_dict["group_ids"] # gt_dict["group_ids"] = group_ids calib = None if "calib" in input_dict: calib = input_dict["current_frame"]["calib"] if reference_detections is not None: assert calib is not None and "image" in input_dict["current_frame"] 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 remove_outside_points: assert calib is not None image_shape = input_dict["current_frame"]["image"]["image_shape"] points = box_np_ops.remove_outside_points(points, calib["rect"], calib["Trv2c"], calib["P2"], image_shape) if remove_environment is True and training: 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 training: # boxes_lidar = gt_dict["gt_boxes"] # cv2.imshow('pre-noise', bev_map) selected = kitti.drop_arrays_by_name(gt_dict["gt_names"], ["DontCare", "ignore"]) _dict_select(gt_dict, selected) if 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 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 class_names for n in gt_dict["gt_names"]], dtype=np.bool_) # db_sampler = None if db_sampler is not None: group_ids = None # if "group_ids" in gt_dict: # group_ids = gt_dict["group_ids"] sampled_dict = db_sampler.sample_all( root_path, gt_dict["gt_boxes"], gt_dict["gt_names"], num_point_features, random_crop, gt_group_ids=group_ids, 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 group_ids is not None: # sampled_group_ids = sampled_dict["group_ids"] # gt_dict["group_ids"] = np.concatenate( # [gt_dict["group_ids"], sampled_group_ids]) if 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) pc_range = voxel_generator.point_cloud_range # group_ids = None # if "group_ids" in gt_dict: # group_ids = gt_dict["group_ids"] # prep.noise_per_object_v3_( # gt_dict["gt_boxes"], # points, # gt_boxes_mask, # rotation_perturb=gt_rotation_noise, # center_noise_std=gt_loc_noise_std, # global_random_rot_range=global_random_rot_range, # group_ids=group_ids, # num_try=100) # should remove unrelated objects after noise per object # for k, v in gt_dict.items(): # print(k, v.shape) _dict_select(gt_dict, gt_boxes_mask) gt_classes = np.array( [class_names.index(n) + 1 for n in gt_dict["gt_names"]], dtype=np.int32) gt_dict["gt_classes"] = gt_classes # concatenate points_current = points.shape[0] points_keyframe = keyframe_points.shape[0] points = np.concatenate((points, keyframe_points), axis=0) # data aug gt_dict["gt_boxes"], points = prep.random_flip(gt_dict["gt_boxes"], points) gt_dict["gt_boxes"], points = prep.global_rotation( gt_dict["gt_boxes"], points, rotation=global_rotation_noise) gt_dict["gt_boxes"], points = prep.global_scaling_v2( gt_dict["gt_boxes"], points, *global_scaling_noise) prep.global_translate_(gt_dict["gt_boxes"], points, global_translate_noise_std) # slice points_keyframe = points[points_current:, :] points = points[:points_current, :] bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]] mask = prep.filter_gt_box_outside_range(gt_dict["gt_boxes"], bv_range) _dict_select(gt_dict, mask) task_masks = [] flag = 0 for class_name in task_class_names: task_masks.append([ np.where(gt_dict["gt_classes"] == class_name.index(i) + 1 + flag) for i in class_name ]) flag += len(class_name) task_boxes = [] task_classes = [] task_names = [] flag2 = 0 for idx, mask in enumerate(task_masks): task_box = [] task_class = [] task_name = [] for m in mask: task_box.append(gt_dict["gt_boxes"][m]) task_class.append(gt_dict["gt_classes"][m] - flag2) task_name.append(gt_dict["gt_names"][m]) task_boxes.append(np.concatenate(task_box, axis=0)) task_classes.append(np.concatenate(task_class)) task_names.append(np.concatenate(task_name)) flag2 += len(mask) for task_box in task_boxes: # limit rad to [-pi, pi] task_box[:, -1] = box_np_ops.limit_period(task_box[:, -1], offset=0.5, period=2 * np.pi) # print(gt_dict.keys()) gt_dict["gt_classes"] = task_classes gt_dict["gt_names"] = task_names gt_dict["gt_boxes"] = task_boxes # if shuffle_points: # # shuffle is a little slow. # np.random.shuffle(points) # [0, -40, -3, 70.4, 40, 1] voxel_size = voxel_generator.voxel_size pc_range = voxel_generator.point_cloud_range grid_size = voxel_generator.grid_size # [352, 400] # points = points[:int(points.shape[0] * 0.1), :] voxels, coordinates, num_points = voxel_generator.generate( points, max_voxels) # res = voxel_generator.generate( # points, max_voxels) # voxels = res["voxels"] # coordinates = res["coordinates"] # num_points = res["num_points_per_voxel"] num_voxels = np.array([voxels.shape[0]], dtype=np.int64) # key frame voxel keyframe_info = voxel_generator.generate(keyframe_points, max_voxels) keyframe_info = keyframe_voxels, keyframe_coordinates, keyframe_num_points keyframe_num_voxels = np.array([keyframe_voxels.shape[0]], dtype=np.int64) example = { "voxels": voxels, "num_points": num_points, "points": points, "coordinates": coordinates, "num_voxels": num_voxels, } example_keyframe = { "voxels": keyframe_voxels, "num_points": keyframe_num_points, "points": keyframe_points, "coordinates": keyframe_coordinates, "num_voxels": keyframe_num_voxels, } if training: example["gt_boxes"] = gt_dict["gt_boxes"] if calib is not None: example["calib"] = calib feature_map_size = grid_size[:2] // out_size_factor feature_map_size = [*feature_map_size, 1][::-1] if anchor_cache is not None: anchorss = anchor_cache["anchors"] anchors_bvs = anchor_cache["anchors_bv"] anchors_dicts = anchor_cache["anchors_dict"] else: rets = [ target_assigner.generate_anchors(feature_map_size) for target_assigner in target_assigners ] anchorss = [ret["anchors"].reshape([-1, 7]) for ret in rets] anchors_dicts = [ target_assigner.generate_anchors_dict(feature_map_size) for target_assigner in target_assigners ] anchors_bvs = [ box_np_ops.rbbox2d_to_near_bbox(anchors[:, [0, 1, 3, 4, 6]]) for anchors in anchorss ] example["anchors"] = anchorss if anchor_area_threshold >= 0: example["anchors_mask"] = [] for idx, anchors_bv in enumerate(anchors_bvs): anchors_mask = None # slow with high resolution. recommend disable this forever. coors = coordinates dense_voxel_map = box_np_ops.sparse_sum_for_anchors_mask( coors, tuple(grid_size[::-1][1:])) dense_voxel_map = dense_voxel_map.cumsum(0) dense_voxel_map = dense_voxel_map.cumsum(1) anchors_area = box_np_ops.fused_get_anchors_area( dense_voxel_map, anchors_bv, voxel_size, pc_range, grid_size) anchors_mask = anchors_area > anchor_area_threshold # example['anchors_mask'] = anchors_mask.astype(np.uint8) example["anchors_mask"].append(anchors_mask) example_sequences = {} example_sequences["current_frame"] = example example_sequences["keyframe"] = example_keyframe if not training: return example_sequences # voxel_labels = box_np_ops.assign_label_to_voxel(gt_boxes, coordinates, # voxel_size, coors_range) """ example.update({ 'gt_boxes': gt_boxes.astype(out_dtype), 'num_gt': np.array([gt_boxes.shape[0]]), # 'voxel_labels': voxel_labels, }) """ if create_targets: targets_dicts = [] for idx, target_assigner in enumerate(target_assigners): if "anchors_mask" in example: anchors_mask = example["anchors_mask"][idx] else: anchors_mask = None targets_dict = target_assigner.assign_v2( anchors_dicts[idx], gt_dict["gt_boxes"][idx], anchors_mask, gt_classes=gt_dict["gt_classes"][idx], gt_names=gt_dict["gt_names"][idx], ) targets_dicts.append(targets_dict) example_sequences["current_frame"].update({ "labels": [targets_dict["labels"] for targets_dict in targets_dicts], "reg_targets": [targets_dict["bbox_targets"] for targets_dict in targets_dicts], "reg_weights": [ targets_dict["bbox_outside_weights"] for targets_dict in targets_dicts ], }) return example_sequences