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): RGB_embedding = True 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: if RGB_embedding: db_info_save_path = root_path / "kitti_dbinfos_train_RGB.pkl" else: 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 = box_np_ops.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_np_ops.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] if RGB_embedding: RGB_image = cv2.imread(str(root_path / info['img_path'])) points_camera = box_np_ops.box_lidar_to_camera( points[:, :3], rect, Trv2c) points_to_image_idx = box_np_ops.project_to_image( points_camera, P2) points_to_image_idx = points_to_image_idx.astype(int) mask = box_np_ops.remove_points_outside_image( RGB_image, points_to_image_idx) points = points[mask] points_to_image_idx = points_to_image_idx[mask] BGR = RGB_image[points_to_image_idx[:, 1], points_to_image_idx[:, 0]] points = np.concatenate((points, BGR), axis=1) 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 = box_np_ops.points_in_rbbox(points, rbbox_lidar) for i in range(num_obj): if RGB_embedding: filename = f"{image_idx}_{names[i]}_{gt_idxes[i]}_RGB.bin" else: 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)
def create_groundtruth_database(data_path, train_or_test, 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): # custom dataset parameter custom_dataset = True # if true, it indicates that we are not operating on the kitti data sample_val_dataset_mode = True if train_or_test == "test" else False # if true, we are creating a gt database for the test set (instead of the train set) # ------------------------------------------------------------------------------------------------------ # create path where the gt boxes are stored # ----------------------------------------------------------------------------------------------------- root_path = pathlib.Path(data_path) if info_path is None: if sample_val_dataset_mode: info_path = root_path / 'kitti_infos_val.pkl' else: info_path = root_path / 'kitti_infos_train.pkl' if database_save_path is None: if sample_val_dataset_mode: database_save_path = root_path / 'gt_database_val' else: database_save_path = root_path / 'gt_database' else: database_save_path = pathlib.Path(database_save_path) if db_info_save_path is None: if sample_val_dataset_mode: db_info_save_path = root_path / "kitti_dbinfos_val.pkl" else: db_info_save_path = root_path / "kitti_dbinfos_train.pkl" database_save_path.mkdir(parents=True, exist_ok=True) # ------------------------------------------------------------------------------------------------------ # load kitti infos # ----------------------------------------------------------------------------------------------------- with open(info_path, 'rb') as f: kitti_infos = pickle.load(f) all_db_infos = {} # get the classnames we are intered in 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 # ------------------------------------------------------------------------------------------------------ # iterate over kitti_infos # ----------------------------------------------------------------------------------------------------- for info in prog_bar(kitti_infos): # ------------------------------------------------------------------------------------------------------ # load pc # ----------------------------------------------------------------------------------------------------- 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'] if custom_dataset: with open(str(velodyne_path)[:-3] + "pkl", 'rb') as file: points = pickle.load(file, encoding='latin1') else: points = np.fromfile(str(velodyne_path), dtype=np.float32, count=-1).reshape([-1, 4]) image_idx = info["image_idx"] rect = info['calib/R0_rect'] P2 = info['calib/P2'] Trv2c = info['calib/Tr_velo_to_cam'] # ------------------------------------------------------------------------------------------------------ # remove boxes outside the frustum of the image # ----------------------------------------------------------------------------------------------------- if not lidar_only and not custom_dataset: points = box_np_ops.remove_outside_points(points, rect, Trv2c, P2, info["img_shape"]) # ------------------------------------------------------------------------------------------------------ # get the bboxes and transform (annos not points) # ----------------------------------------------------------------------------------------------------- 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_np_ops.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] # ------------------------------------------------------------------------------------------------------ # other stuff # ----------------------------------------------------------------------------------------------------- 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) # ------------------------------------------------------------------------------------------------------ # test which points are in bboxes # ----------------------------------------------------------------------------------------------------- point_indices = box_np_ops.points_in_rbbox(points, rbbox_lidar) # ------------------------------------------------------------------------------------------------------ # iterate over all objects in the annos # ----------------------------------------------------------------------------------------------------- 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] # ------------------------------------------------------------------------------------------------------ # save points of gt boxes to files # ----------------------------------------------------------------------------------------------------- with open(str(filepath)[:-3] + "pkl", 'wb') as file: pickle.dump(np.array(gt_points), file, 2) # debug # with open("/home/makr/Documents/uni/TU/3.Master/experiments/own/tf_3dRGB_pc/debug_rviz/points/bbox_pixels.pkl", 'wb') as file: # pickle.dump(np.array(gt_points), file, 2) # ------------------------------------------------------------------------------------------------------ # save infos of gt boxes to single file # ----------------------------------------------------------------------------------------------------- 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)
def create_groundtruth_database( data_path='/mrtstorage/datasets/kitti/object_detection', info_path='/home/zhwang/second.pytorch/second/data/sets/kitti_second' '/kitti_infos_train.pkl', used_classes=None, database_save_path='/home/zhwang/second.pytorch/second/data/sets/kitti_second' '/gt_database', db_info_save_path='/home/zhwang/second.pytorch/second/data/sets/kitti_second' '/kitti_dbinfos_train.pkl', 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 = box_np_ops.remove_outside_points(points, rect, Trv2c, P2, info["img_shape"]) """ fusion camera data to lidar """ image_label = _get_semantic_segmentation_result( image_idx, data_dir='/home/zhwang/semantic-segmentation' '/kitti_train_results') points = _add_class_score(image_label, points, rect, Trv2c, P2) 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_np_ops.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 = box_np_ops.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)
def create_inference_dataset(det_anno_path, calib_path, velodyne_path, box3d_expansion, bbox_expansion, output_file): # get annos annos = get_label_annos(det_anno_path) # create examples examples = [] for anno in tqdm(annos): # continue if zero object detected if anno['image_idx'].shape[0] == 0: continue # get scene index scene_idx = str(anno['image_idx'][0]).zfill(6) # get calib calib = get_kitti_calib(os.path.join(calib_path, scene_idx + '.txt'), True) # get box3d_camera box3d_camera = anno_to_rbboxes(anno) box3d_lidar = box_camera_to_lidar(box3d_camera, calib["R0_rect"], calib["Tr_velo_to_cam"]) box3d_lidar[:, 2] += box3d_lidar[:, 5] / 2 # get expanded box3d box3d_lidar_expanded = expand_box3d(box3d_lidar, box3d_expansion) # get bbox bbox = box3d_to_bbox(box3d_lidar_expanded, calib["R0_rect"], calib["Tr_velo_to_cam"], calib['P2']) bbox_expanded = expand_bbox(bbox, bbox_expansion).astype(np.int) bbox_expanded[:, 0] = np.clip(bbox_expanded[:, 0], 0, 1242) bbox_expanded[:, 1] = np.clip(bbox_expanded[:, 1], 0, 375) bbox_expanded[:, 2] = np.clip(bbox_expanded[:, 2], 0, 1242) bbox_expanded[:, 3] = np.clip(bbox_expanded[:, 3], 0, 375) # read scene pts pts = read_bin(os.path.join(velodyne_path, scene_idx + '.bin'))[:, :3] # create example for idx in range(box3d_lidar_expanded.shape[0]): filtered_pts = points_in_rbbox( pts, box3d_lidar_expanded[idx][np.newaxis, ...]) res = { 'rgb': {}, 'point': {}, 'calib': calib, } res['rgb']['bbox'] = bbox_expanded[idx] res['rgb']['img_id'] = anno['image_idx'][0] res['point']['box3d_lidar'] = box3d_lidar_expanded[idx] res['point']['point_dict'] = { 'source': [pts[filtered_pts.reshape(-1)]], 'gt': [], 3: [], 4: [], 5: [], 7: [], 9: [], } examples.append(res) # write to file with open(output_file, 'wb') as f: pickle.dump(examples, f)