def _calculate_num_points_in_gt(data_path, infos, relative_path, remove_outside=True, num_features=4): for info in infos: if relative_path: v_path = str(pathlib.Path(data_path) / info["velodyne_path"]) else: v_path = info["velodyne_path"] points_v = np.fromfile(v_path, dtype=np.float32, count=-1).reshape([-1, num_features]) rect = info['calib/R0_rect'] Trv2c = info['calib/Tr_velo_to_cam'] P2 = info['calib/P2'] if remove_outside: points_v = remove_outside_points(points_v, rect, Trv2c, P2, info["img_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_camera_to_lidar(gt_boxes_camera, rect, Trv2c) indices = 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 prepare_train_img(self, idx): sample_id = self.sample_ids[idx] # load image img = mmcv.imread(osp.join(self.img_prefix, '%06d.png' % sample_id)) img, img_shape, pad_shape, scale_factor = self.img_transform( img, 1, False) objects = read_label( osp.join(self.label_prefix, '%06d.txt' % sample_id)) calib = Calibration(osp.join(self.calib_prefix, '%06d.txt' % sample_id)) gt_bboxes = [ object.box3d for object in objects if object.type not in ["DontCare"] ] gt_bboxes = np.array(gt_bboxes, dtype=np.float32) gt_types = [ object.type for object in objects if object.type not in ["DontCare"] ] # transfer from cam to lidar coordinates if len(gt_bboxes) != 0: gt_bboxes[:, :3] = project_rect_to_velo(gt_bboxes[:, :3], calib) img_meta = dict(img_shape=img_shape, sample_idx=sample_id, calib=calib) data = dict(img=to_tensor(img), img_meta=DC(img_meta, cpu_only=True)) if self.anchors is not None: data['anchors'] = DC(to_tensor(self.anchors.astype(np.float32))) if self.with_mask: NotImplemented if self.with_point: points = read_lidar( osp.join(self.lidar_prefix, '%06d.bin' % sample_id)) if self.augmentor is not None and self.test_mode is False: sampled_gt_boxes, sampled_gt_types, sampled_points = self.augmentor.sample_all( gt_bboxes, gt_types) assert sampled_points.dtype == np.float32 gt_bboxes = np.concatenate([gt_bboxes, sampled_gt_boxes]) gt_types = gt_types + sampled_gt_types assert len(gt_types) == len(gt_bboxes) # to avoid overlapping point (option) masks = points_in_rbbox(points, sampled_gt_boxes) #masks = points_op_cpu.points_in_bbox3d_np(points[:,:3], sampled_gt_boxes) points = points[np.logical_not(masks.any(-1))] # paste sampled points to the scene points = np.concatenate([sampled_points, points], axis=0) # select the interest classes selected = [ i for i in range(len(gt_types)) if gt_types[i] in self.class_names ] gt_bboxes = gt_bboxes[selected, :] gt_types = [ gt_types[i] for i in range(len(gt_types)) if gt_types[i] in self.class_names ] # force van to have same label as car gt_types = ['Car' if n == 'Van' else n for n in gt_types] gt_labels = np.array( [self.class_names.index(n) + 1 for n in gt_types], dtype=np.int64) self.augmentor.noise_per_object_(gt_bboxes, points, num_try=100) gt_bboxes, points = self.augmentor.random_flip(gt_bboxes, points) gt_bboxes, points = self.augmentor.global_rotation( gt_bboxes, points) gt_bboxes, points = self.augmentor.global_scaling( gt_bboxes, points) if isinstance(self.generator, VoxelGenerator): #voxels, coordinates, num_points = self.generator.generate(points) voxel_size = self.generator.voxel_size pc_range = self.generator.point_cloud_range grid_size = self.generator.grid_size keep = points_op_cpu.points_bound_kernel(points, pc_range[:3], pc_range[3:]) voxels = points[keep, :] coordinates = ( (voxels[:, [2, 1, 0]] - np.array(pc_range[[2, 1, 0]], dtype=np.float32)) / np.array(voxel_size[::-1], dtype=np.float32)).astype(np.int32) num_points = np.ones(len(keep)).astype(np.int32) data['voxels'] = DC(to_tensor(voxels.astype(np.float32))) data['coordinates'] = DC(to_tensor(coordinates)) data['num_points'] = DC(to_tensor(num_points)) if self.anchor_area_threshold >= 0 and self.anchors is not None: dense_voxel_map = sparse_sum_for_anchors_mask( coordinates, tuple(grid_size[::-1][1:])) dense_voxel_map = dense_voxel_map.cumsum(0) dense_voxel_map = dense_voxel_map.cumsum(1) anchors_area = fused_get_anchors_area(dense_voxel_map, self.anchors_bv, voxel_size, pc_range, grid_size) anchors_mask = anchors_area > self.anchor_area_threshold data['anchors_mask'] = DC( to_tensor(anchors_mask.astype(np.uint8))) # filter gt_bbox out of range bv_range = self.generator.point_cloud_range[[0, 1, 3, 4]] mask = filter_gt_box_outside_range(gt_bboxes, bv_range) gt_bboxes = gt_bboxes[mask] gt_labels = gt_labels[mask] else: NotImplementedError # skip the image if there is no valid gt bbox if len(gt_bboxes) == 0: return None # limit rad to [-pi, pi] gt_bboxes[:, 6] = limit_period(gt_bboxes[:, 6], offset=0.5, period=2 * np.pi) if self.with_label: data['gt_labels'] = DC(to_tensor(gt_labels)) data['gt_bboxes'] = DC(to_tensor(gt_bboxes)) return data
def create_groundtruth_database(data_path, info_path=None, used_classes=None, database_save_path=None, db_info_save_path=None, relative_path=True, lidar_only=False, bev_only=False, coors_range=None): root_path = pathlib.Path(data_path) if info_path is None: info_path = root_path / 'kitti_infos_train.pkl' if database_save_path is None: database_save_path = root_path / 'gt_database' else: database_save_path = pathlib.Path(database_save_path) if db_info_save_path is None: db_info_save_path = root_path / "kitti_dbinfos_train.pkl" database_save_path.mkdir(parents=True, exist_ok=True) with open(info_path, 'rb') as f: kitti_infos = pickle.load(f) all_db_infos = {} if used_classes is None: used_classes = list(kitti.get_classes()) used_classes.pop(used_classes.index('DontCare')) for name in used_classes: all_db_infos[name] = [] group_counter = 0 for info in prog_bar(kitti_infos): velodyne_path = info['velodyne_path'] if relative_path: # velodyne_path = str(root_path / velodyne_path) + "_reduced" velodyne_path = str(root_path / velodyne_path) num_features = 4 if 'pointcloud_num_features' in info: num_features = info['pointcloud_num_features'] points = np.fromfile(velodyne_path, dtype=np.float32, count=-1).reshape([-1, num_features]) image_idx = info["image_idx"] rect = info['calib/R0_rect'] P2 = info['calib/P2'] Trv2c = info['calib/Tr_velo_to_cam'] if not lidar_only: points = remove_outside_points(points, rect, Trv2c, P2, info["img_shape"]) annos = info["annos"] names = annos["name"] bboxes = annos["bbox"] difficulty = annos["difficulty"] gt_idxes = annos["index"] num_obj = np.sum(annos["index"] >= 0) rbbox_cam = kitti.anno_to_rbboxes(annos)[:num_obj] rbbox_lidar = box_camera_to_lidar(rbbox_cam, rect, Trv2c) if bev_only: # set z and h to limits assert coors_range is not None rbbox_lidar[:, 2] = coors_range[2] rbbox_lidar[:, 5] = coors_range[5] - coors_range[2] group_dict = {} group_ids = np.full([bboxes.shape[0]], -1, dtype=np.int64) if "group_ids" in annos: group_ids = annos["group_ids"] else: group_ids = np.arange(bboxes.shape[0], dtype=np.int64) point_indices = points_in_rbbox(points, rbbox_lidar) for i in range(num_obj): filename = f"{image_idx}_{names[i]}_{gt_idxes[i]}.bin" filepath = database_save_path / filename gt_points = points[point_indices[:, i]] gt_points[:, :3] -= rbbox_lidar[i, :3] with open(filepath, 'w') as f: gt_points.tofile(f) if names[i] in used_classes: if relative_path: db_path = str(database_save_path.stem + "/" + filename) else: db_path = str(filepath) db_info = { "name": names[i], "path": db_path, "image_idx": image_idx, "gt_idx": gt_idxes[i], "box3d_lidar": rbbox_lidar[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] all_db_infos[names[i]].append(db_info) for k, v in all_db_infos.items(): print(f"load {len(v)} {k} database infos") with open(db_info_save_path, 'wb') as f: pickle.dump(all_db_infos, f)