def _fill_trainval_infos(nusc, train_scenes, val_scenes, test=False, max_sweeps=10): train_nusc_infos = [] val_nusc_infos = [] from pyquaternion import Quaternion for sample in prog_bar(nusc.sample): lidar_token = sample["data"]["LIDAR_TOP"] cam_front_token = sample["data"]["CAM_FRONT"] sd_rec = nusc.get('sample_data', sample['data']["LIDAR_TOP"]) cs_record = nusc.get('calibrated_sensor', sd_rec['calibrated_sensor_token']) pose_record = nusc.get('ego_pose', sd_rec['ego_pose_token']) lidar_path, boxes, _ = nusc.get_sample_data(lidar_token) cam_path, _, cam_intrinsic = nusc.get_sample_data(cam_front_token) assert Path(lidar_path).exists(), ( "you must download all trainval data, key-frame only dataset performs far worse than sweeps." ) info = { "lidar_path": lidar_path, "cam_front_path": cam_path, "token": sample["token"], "sweeps": [], "lidar2ego_translation": cs_record['translation'], "lidar2ego_rotation": cs_record['rotation'], "ego2global_translation": pose_record['translation'], "ego2global_rotation": pose_record['rotation'], "timestamp": sample["timestamp"], } l2e_r = info["lidar2ego_rotation"] l2e_t = info["lidar2ego_translation"] e2g_r = info["ego2global_rotation"] e2g_t = info["ego2global_translation"] l2e_r_mat = Quaternion(l2e_r).rotation_matrix e2g_r_mat = Quaternion(e2g_r).rotation_matrix sd_rec = nusc.get('sample_data', sample['data']["LIDAR_TOP"]) sweeps = [] while len(sweeps) < max_sweeps: if not sd_rec['prev'] == "": sd_rec = nusc.get('sample_data', sd_rec['prev']) cs_record = nusc.get('calibrated_sensor', sd_rec['calibrated_sensor_token']) pose_record = nusc.get('ego_pose', sd_rec['ego_pose_token']) lidar_path = nusc.get_sample_data_path(sd_rec['token']) sweep = { "lidar_path": lidar_path, "sample_data_token": sd_rec['token'], "lidar2ego_translation": cs_record['translation'], "lidar2ego_rotation": cs_record['rotation'], "ego2global_translation": pose_record['translation'], "ego2global_rotation": pose_record['rotation'], "timestamp": sd_rec["timestamp"] } l2e_r_s = sweep["lidar2ego_rotation"] l2e_t_s = sweep["lidar2ego_translation"] e2g_r_s = sweep["ego2global_rotation"] e2g_t_s = sweep["ego2global_translation"] # sweep->ego->global->ego'->lidar l2e_r_s_mat = Quaternion(l2e_r_s).rotation_matrix e2g_r_s_mat = Quaternion(e2g_r_s).rotation_matrix R = (l2e_r_s_mat.T @ e2g_r_s_mat.T) @ ( np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T) T = (l2e_t_s @ e2g_r_s_mat.T + e2g_t_s) @ ( np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T) T -= e2g_t @ (np.linalg.inv(e2g_r_mat).T @ np.linalg.inv( l2e_r_mat).T) + l2e_t @ np.linalg.inv(l2e_r_mat).T sweep["sweep2lidar_rotation"] = R.T # points @ R.T + T sweep["sweep2lidar_translation"] = T sweeps.append(sweep) else: break info["sweeps"] = sweeps if not test: annotations = [ nusc.get('sample_annotation', token) for token in sample['anns'] ] locs = np.array([b.center for b in boxes]).reshape(-1, 3) dims = np.array([b.wlh for b in boxes]).reshape(-1, 3) rots = np.array([b.orientation.yaw_pitch_roll[0] for b in boxes]).reshape(-1, 1) velocity = np.array( [nusc.box_velocity(token)[:2] for token in sample['anns']]) # convert velo from global to lidar for i in range(len(boxes)): velo = np.array([*velocity[i], 0.0]) velo = velo @ np.linalg.inv(e2g_r_mat).T @ np.linalg.inv( l2e_r_mat).T velocity[i] = velo[:2] names = [b.name for b in boxes] for i in range(len(names)): if names[i] in NuScenesDataset.NameMapping: names[i] = NuScenesDataset.NameMapping[names[i]] names = np.array(names) # we need to convert rot to SECOND format. # change the rot format will break all checkpoint, so... gt_boxes = np.concatenate([locs, dims, -rots - np.pi / 2], axis=1) assert len(gt_boxes) == len( annotations), f"{len(gt_boxes)}, {len(annotations)}" info["gt_boxes"] = gt_boxes info["gt_names"] = names info["gt_velocity"] = velocity.reshape(-1, 2) info["num_lidar_pts"] = np.array( [a["num_lidar_pts"] for a in annotations]) info["num_radar_pts"] = np.array( [a["num_radar_pts"] for a in annotations]) if sample["scene_token"] in train_scenes: train_nusc_infos.append(info) else: val_nusc_infos.append(info) return train_nusc_infos, val_nusc_infos
def _fill_trainval_infos(lyft_ds, train_scenes, val_scenes, do_test=False, max_sweeps=10): train_nusc_infos = [] val_nusc_infos = [] for sample in prog_bar(lyft_ds.sample): lidar_token = sample['data']['LIDAR_TOP'] cam_front_token = sample['data']['CAM_FRONT'] sd_rec = lyft_ds.get('sample_data', lidar_token) cs_record = lyft_ds.get('calibrated_sensor', sd_rec['calibrated_sensor_token']) pose_record = lyft_ds.get('ego_pose', sd_rec['ego_pose_token']) lidar_path, boxes, _ = lyft_ds.get_sample_data(lidar_token) cam_path, _, cam_intrinsic = lyft_ds.get_sample_data(cam_front_token) assert Path(lidar_path).exists(), ( 'you must download all trainval data, key-frame only dataset performs far worse than sweeps.' ) info = { 'lidar_path': lidar_path, 'cam_front_path': cam_path, 'token': sample['token'], 'sweeps': [], 'lidar2ego_translation': cs_record['translation'], 'lidar2ego_rotation': cs_record['rotation'], 'ego2global_translation': pose_record['translation'], 'ego2global_rotation': pose_record['rotation'], 'timestamp': sample['timestamp'], } l2e_r = info['lidar2ego_rotation'] l2e_t = info['lidar2ego_translation'] e2g_r = info['ego2global_rotation'] e2g_t = info['ego2global_translation'] l2e_r_mat = Quaternion(l2e_r).rotation_matrix e2g_r_mat = Quaternion(e2g_r).rotation_matrix sweeps = [] while len(sweeps) < max_sweeps: if not sd_rec['prev'] == '': sd_rec = lyft_ds.get('sample_data', sd_rec['prev']) cs_record = lyft_ds.get('calibrated_sensor', sd_rec['calibrated_sensor_token']) pose_record = lyft_ds.get('ego_pose', sd_rec['ego_pose_token']) lidar_path = lyft_ds.get_sample_data_path(sd_rec['token']) sweep = { 'lidar_path': lidar_path, 'sample_data_token': sd_rec['token'], 'lidar2ego_translation': cs_record['translation'], 'lidar2ego_rotation': cs_record['rotation'], 'ego2global_translation': pose_record['translation'], 'ego2global_rotation': pose_record['rotation'], 'timestamp': sd_rec['timestamp'] } l2e_r_s = sweep['lidar2ego_rotation'] l2e_t_s = sweep['lidar2ego_translation'] e2g_r_s = sweep['ego2global_rotation'] e2g_t_s = sweep['ego2global_translation'] # sweep->ego->global->ego'->lidar l2e_r_s_mat = Quaternion(l2e_r_s).rotation_matrix e2g_r_s_mat = Quaternion(e2g_r_s).rotation_matrix R = (l2e_r_s_mat.T @ e2g_r_s_mat.T) @ ( np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T) T = (l2e_t_s @ e2g_r_s_mat.T + e2g_t_s) @ ( np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T) T -= e2g_t @ (np.linalg.inv(e2g_r_mat).T @ np.linalg.inv( l2e_r_mat).T) + l2e_t @ np.linalg.inv(l2e_r_mat).T sweep['sweep2lidar_rotation'] = R.T # points @ R.T + T sweep['sweep2lidar_translation'] = T sweeps.append(sweep) else: break info['sweeps'] = sweeps if not do_test: annotations = [ lyft_ds.get('sample_annotation', token) for token in sample['anns'] ] locs = np.array([b.center for b in boxes]).reshape(-1, 3) dims = np.array([b.wlh for b in boxes]).reshape(-1, 3) rots = np.array([b.orientation.yaw_pitch_roll[0] for b in boxes]).reshape(-1, 1) velocity = np.array( [lyft_ds.box_velocity(token)[:2] for token in sample['anns']]) # convert velo from global to lidar for i in range(len(boxes)): velo = np.array([*velocity[i], 0.0]) velo = velo @ np.linalg.inv(e2g_r_mat).T @ np.linalg.inv( l2e_r_mat).T velocity[i] = velo[:2] names = np.array([b.name for b in boxes]) # we need to convert rot to SECOND format. # change the rot format will break all checkpoint, so... gt_boxes = np.concatenate([locs, dims, -rots - np.pi / 2], axis=1) assert len(gt_boxes) == len( annotations), f'{len(gt_boxes)}, {len(annotations)}' info['gt_boxes'] = gt_boxes info['gt_names'] = names info['gt_velocity'] = velocity.reshape(-1, 2) info['num_lidar_pts'] = np.array( [a['num_lidar_pts'] for a in annotations]) info['num_radar_pts'] = np.array( [a['num_radar_pts'] for a in annotations]) if sample['scene_token'] in train_scenes: train_nusc_infos.append(info) else: val_nusc_infos.append(info) return train_nusc_infos, val_nusc_infos
def create_groundtruth_database(dataset_class_name, data_path, info_path=None, used_classes=None, database_save_path=None, db_info_save_path=None, relative_path=True, add_rgb=False, lidar_only=False, bev_only=False, coors_range=None): dataset = get_dataset_class(dataset_class_name)( info_path=info_path, root_path=data_path, ) root_path = Path(data_path) if database_save_path is None: database_save_path = root_path / 'gt_database' else: database_save_path = 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) all_db_infos = {} group_counter = 0 for j in prog_bar(list(range(len(dataset)))): image_idx = j sensor_data = dataset.get_sensor_data(j) if "image_idx" in sensor_data["metadata"]: image_idx = sensor_data["metadata"]["image_idx"] points = sensor_data["lidar"]["points"] 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 = database_save_path / filename gt_points = points[point_indices[:, i]] gt_points[:, :3] -= gt_boxes[i, :3] with open(filepath, 'w') as f: gt_points.tofile(f) if (used_classes is None) or 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": 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] 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, info_path=None, used_classes=None, database_save_path=None, db_info_save_path=None, relative_path=True, lidar_only=True, 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"]) #print(points) annos = info["annos"] names = annos["name"] bboxes = annos["bbox"] difficulty = annos["difficulty"] #print(difficulty) gt_idxes = annos["index"] num_obj = np.sum(annos["index"] >= 0) rbbox_cam = kitti.anno_to_rbboxes(annos)[:num_obj] #print(rbbox_cam) rbbox_lidar = box_np_ops.box_camera_to_lidar(rbbox_cam, rect, Trv2c) #print(rbbox_lidar) 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] #print(rbbox_lidar) 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) #print(group_ids) point_indices = box_np_ops.points_in_rbbox(points, rbbox_lidar) #print(point_indices) #print(num_obj) 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]] #print(gt_points) gt_points[:, :3] -= rbbox_lidar[i, :3] #print(gt_points) 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 evaluation_nusc(self, detections, output_dir, clean_after=True): version = self.version eval_set_map = { "v1.0-mini": "mini_train", "v1.0-trainval": "val", } nusc_annos = {} mapped_class_names = self._class_names token2info = {} for info in self._nusc_infos: token2info[info["token"]] = info for det in prog_bar(detections): annos = [] boxes = _second_det_to_nusc_box(det) for i, box in enumerate(boxes): name = mapped_class_names[box.label] velocity = box.velocity[:2].tolist() if len(token2info[det["metadata"]["token"]]["sweeps"]) == 0: velocity = (np.nan, np.nan) box.velocity = np.array([*velocity, 0.0]) boxes = _lidar_nusc_box_to_global( token2info[det["metadata"]["token"]], boxes, mapped_class_names, "cvpr_2019") for i, box in enumerate(boxes): name = mapped_class_names[box.label] velocity = box.velocity[:2].tolist() nusc_anno = { "sample_token": det["metadata"]["token"], "translation": box.center.tolist(), "size": box.wlh.tolist(), "rotation": box.orientation.elements.tolist(), "velocity": velocity, "detection_name": name, "detection_score": box.score, "attribute_name": NuScenesDataset.DefaultAttribute[name], } annos.append(nusc_anno) nusc_annos[det["metadata"]["token"]] = annos nusc_submissions = { "meta": { "use_camera": False, "use_lidar": True, "use_radar": False, "use_map": False, "use_external": False, }, "results": nusc_annos, } res_path = Path(output_dir) / "results_nusc.json" with open(res_path, "w") as f: json.dump(nusc_submissions, f) print(f"evaluation_nusc(): dumping results to {res_path}") gt_annos = self.ground_truth_annotations if gt_annos is None: return None eval_main_file = Path(__file__).resolve().parent / "nusc_eval.py" # why add \"{}\"? to support path with spaces. cmd = f"python {str(eval_main_file)} --root_path=\"{str(self._root_path)}\"" cmd += f" --version={self.version} --eval_version={self.eval_version}" cmd += f" --res_path=\"{str(res_path)}\" --eval_set={eval_set_map[self.version]}" cmd += f" --output_dir=\"{output_dir}\"" # use subprocess can release all nusc memory after evaluation subprocess.check_output(cmd, shell=True) with open(Path(output_dir) / "metrics_summary.json", "r") as f: metrics = json.load(f) detail = {} if clean_after: # only delete results_nusc.json when needed res_path.unlink() # delete results_nusc.json since it's very large result = f"Nusc {version} Evaluation\n" for name in mapped_class_names: detail[name] = {} for k, v in metrics["label_aps"][name].items(): detail[name][f"dist@{k}"] = v tp_errs = [] tp_names = [] for k, v in metrics["label_tp_errors"][name].items(): detail[name][k] = v tp_errs.append(f"{v:.4f}") tp_names.append(k) threshs = ', '.join(list(metrics["label_aps"][name].keys())) scores = list(metrics["label_aps"][name].values()) scores = ', '.join([f"{s * 100:.2f}" for s in scores]) result += f"{name} Nusc dist AP@{threshs} and TP errors\n" result += scores result += "\n" result += ', '.join(tp_names) + ": " + ', '.join(tp_errs) result += "\n" return { "results": { "nusc": result }, "detail": { "nusc": detail }, }
def _fill_trainval_infos(nusc, train_scenes, val_scenes, test=False, max_sweeps=9): train_nusc_infos = [] val_nusc_infos = [] from pyquaternion import Quaternion for sample in prog_bar(nusc.sample): lidar_token = sample["data"]["LIDAR_TOP"] cam_front_token = sample["data"]["CAM_FRONT"] sd_rec = nusc.get('sample_data', sample['data']["LIDAR_TOP"]) cs_record = nusc.get('calibrated_sensor', sd_rec['calibrated_sensor_token']) pose_record = nusc.get('ego_pose', sd_rec['ego_pose_token']) lidar_path, boxes, _ = nusc.get_sample_data(lidar_token) cam_path, _, cam_intrinsic = nusc.get_sample_data(cam_front_token) # if Path(lidar_path).exists(): info = { "lidar_path": lidar_path, "cam_front_path": cam_path, "token": sample["token"], "sweeps": [], "lidar2ego_translation": cs_record['translation'], "lidar2ego_rotation": cs_record['rotation'], "ego2global_translation": pose_record['translation'], "ego2global_rotation": pose_record['rotation'], "timestamp": sample["timestamp"], } l2e_r = info["lidar2ego_rotation"] l2e_t = info["lidar2ego_translation"] e2g_r = info["ego2global_rotation"] e2g_t = info["ego2global_translation"] l2e_r_mat = Quaternion(l2e_r).rotation_matrix e2g_r_mat = Quaternion(e2g_r).rotation_matrix sd_rec = nusc.get('sample_data', sample['data']["LIDAR_TOP"]) sweeps = [] while len(sweeps) < max_sweeps: if not sd_rec['prev'] == "": sd_rec = nusc.get('sample_data', sd_rec['prev']) cs_record = nusc.get('calibrated_sensor', sd_rec['calibrated_sensor_token']) pose_record = nusc.get('ego_pose', sd_rec['ego_pose_token']) lidar_path = nusc.get_sample_data_path(sd_rec['token']) sweep = { "lidar_path": lidar_path, "sample_data_token": sd_rec['token'], "lidar2ego_translation": cs_record['translation'], "lidar2ego_rotation": cs_record['rotation'], "ego2global_translation": pose_record['translation'], "ego2global_rotation": pose_record['rotation'], "timestamp": sd_rec["timestamp"] } l2e_r_s = sweep["lidar2ego_rotation"] l2e_t_s = sweep["lidar2ego_translation"] e2g_r_s = sweep["ego2global_rotation"] e2g_t_s = sweep["ego2global_translation"] # sweep->ego->global->ego'->lidar l2e_r_s_mat = Quaternion(l2e_r_s).rotation_matrix e2g_r_s_mat = Quaternion(e2g_r_s).rotation_matrix R = (l2e_r_s_mat.T @ e2g_r_s_mat.T) @ ( np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T) T = (l2e_t_s @ e2g_r_s_mat.T + e2g_t_s) @ ( np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(l2e_r_mat).T) T -= e2g_t @ (np.linalg.inv(e2g_r_mat).T @ np.linalg.inv( l2e_r_mat).T) + l2e_t @ np.linalg.inv(l2e_r_mat).T sweep["sweep2lidar_rotation"] = R sweep["sweep2lidar_translation"] = T sweeps.append(sweep) else: break info["sweeps"] = sweeps if not test: locs = np.array([b.center for b in boxes]).reshape(-1, 3) dims = np.array([b.wlh for b in boxes]).reshape(-1, 3) rots = np.array([b.orientation.yaw_pitch_roll[0] for b in boxes]).reshape(-1, 1) names = np.array([b.name for b in boxes]) gt_boxes = np.concatenate([locs, dims, -rots - np.pi / 2], axis=1) info["gt_boxes"] = gt_boxes info["gt_names"] = names if sample["scene_token"] in train_scenes: train_nusc_infos.append(info) else: val_nusc_infos.append(info) return train_nusc_infos, val_nusc_infos
def _create_reduced_point_cloud(data_path, info_path, save_path=None, color=True, back=False): bgr_path = None with open(info_path, 'rb') as f: kitti_infos = pickle.load(f) for info in prog_bar(kitti_infos): pc_info = info["point_cloud"] image_info = info["image"] image_path = image_info['image_path'] calib = info["calib"] v_path = pc_info['velodyne_path'] v_path = Path(data_path) / v_path points_v = np.fromfile(str(v_path), dtype=np.float32, count=-1).reshape([-1, 4]) rect = calib['R0_rect'] P2 = calib['P2'] Trv2c = calib['Tr_velo_to_cam'] # first remove z < 0 points # keep = points_v[:, -1] > 0 # points_v = points_v[keep] # then remove outside. if back: points_v[:, 0] = -points_v[:, 0] points_v = box_np_ops.remove_outside_points(points_v, rect, Trv2c, P2, image_info["image_shape"]) points_xyz = lidar_to_camera(points_v, rect, Trv2c) points_xy = project_to_image(points_xyz, P2) image_path = v_path.parent.parent.parent / image_path BGR_image = cv2.imread(str(image_path)) points_bgr = np.zeros((len(points_v), 3)) image_shape = image_info["image_shape"] for i in range(len(points_bgr)): new_y = int(points_xy[i][0]) if new_y >= image_shape[1]: new_y = image_shape[1] - 1 new_x = int(points_xy[i][1]) if new_x >= image_shape[0]: new_x = image_shape[0] - 1 points_bgr[i] = BGR_image[new_x][new_y] / 255 if bgr_path is None: save_bgr_filename = v_path.parent.parent / (v_path.parent.stem + "_bgr") / v_path.name # save_filename = str(v_path) + '_reduced' if back: save_filename += "_bgr_back" else: save_bgr_filename = str(Path(bgr_path) / v_path.name) if back: save_bgr_filename += "_bgr_back" with open(save_bgr_filename, 'w') as f: points_bgr.tofile(f) if save_path is None: save_filename = v_path.parent.parent / (v_path.parent.stem + "_bgr") / v_path.name # save_filename = str(v_path) + '_reduced' if back: save_filename += "_back" else: save_filename = str(Path(save_path) / v_path.name) if back: save_filename += "_back" with open(save_filename, 'w') as f: points_v.tofile(f)
def create_my_groundtruth_database(dataset_class_name, data_path, info_path=None, used_classes=None, database_save_path=None, db_info_save_path=None, relative_path=True): dataset = get_dataset_class(dataset_class_name)( info_path=info_path, root_path=data_path, ) root_path = Path(data_path) if database_save_path is None: database_save_path = root_path / 'gt_database' else: database_save_path = Path(database_save_path) if db_info_save_path is None: db_info_save_path = root_path / "my_dbinfos_train.pkl" database_save_path.mkdir(parents=True, exist_ok=True) all_db_infos = {} group_counter = 0 for j in prog_bar(list(range(len(dataset)))): image_idx = j sensor_data = dataset.get_sensor_data(j) if "image_idx" in sensor_data["metadata"]: image_idx = sensor_data["metadata"]["image_idx"] points = sensor_data["lidar"]["points"] annos = sensor_data["lidar"]["annotations"] gt_boxes = annos["boxes"] names = annos["names"] group_dict = {} group_ids = np.arange(gt_boxes.shape[0], dtype=np.int64) difficulty = np.zeros(gt_boxes.shape[0], dtype=np.int32) num_obj = gt_boxes.shape[0] point_indices = None if num_obj > 0: point_indices = box_np_ops.points_in_rbbox(points, gt_boxes) for i in range(num_obj): filename = "{}_{}_{}.bin".format(image_idx, names[i], i) filepath = database_save_path / filename gt_points = points[point_indices[:, i]] if gt_points.shape[0] >= 5: gt_points[:, :3] -= gt_boxes[i, :3] with open(filepath, 'w') as f: gt_points.tofile(f) if (used_classes is None) or 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": i, "box3d_lidar": gt_boxes[i], "num_points_in_gt": gt_points.shape[0], "difficulty": difficulty[i], } 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] if names[i] in all_db_infos: all_db_infos[names[i]].append(db_info) else: all_db_infos[names[i]] = [db_info] for k, v in all_db_infos.items(): print("load {} {} database infos".format(len(v), k)) with open(db_info_save_path, 'wb') as f: pickle.dump(all_db_infos, f)
def _fill_trainval_infos(lyftdata, use_flat_vehicle_coordinates, use_second_format_direction, calc_num_points, is_test=False): train_infos, val_infos = [], [] train_scene_tokens = [ scene['token'] for scene in lyftdata.scene if scene['name'] in splits.train ] for sample in prog_bar(sorted(lyftdata.sample, key=lambda s: s['timestamp'])): lidar_token = sample['data']['LIDAR_TOP'] cam_front_token = sample['data']['CAM_FRONT'] sd_record = lyftdata.get('sample_data', lidar_token) cs_record = lyftdata.get('calibrated_sensor', sd_record['calibrated_sensor_token']) pose_record = lyftdata.get('ego_pose', sd_record['ego_pose_token']) lidar_path, boxes_lidar, _ = lyftdata.get_sample_data(lidar_token) cam_path, _, cam_intrinsic = lyftdata.get_sample_data(cam_front_token) info = { 'lidar_path': lidar_path, 'cam_front_path': cam_path, 'token': sample['token'], 'lidar2ego_translation': cs_record['translation'], 'lidar2ego_rotation': cs_record['rotation'], 'ego2global_translation': pose_record['translation'], 'ego2global_rotation': pose_record['rotation'], 'timestamp': sample['timestamp'], } if not is_test: locs = np.array([b.center for b in boxes_lidar]).reshape(-1, 3) dims = np.array([b.wlh for b in boxes_lidar]).reshape(-1, 3) rots = np.array([b.orientation.yaw_pitch_roll[0] for b in boxes_lidar]).reshape(-1, 1) gt_boxes = np.concatenate([locs, dims, rots], axis=1) gt_names = np.array([b.name for b in boxes_lidar]) if calc_num_points: try: pointcloud = LidarPointCloud.from_file(lidar_path) except Exception as e: print('ERROR:', e, lidar_path) continue indices = box_np_ops.points_in_rbbox(pointcloud.points.T[:, :3], gt_boxes) num_points_in_gt = indices.sum(0) info['num_lidar_pts'] = num_points_in_gt.astype(np.int32) if use_flat_vehicle_coordinates: _, boxes_flat_vehicle, _ = lyftdata.get_sample_data( lidar_token, use_flat_vehicle_coordinates=True ) locs = np.array([b.center for b in boxes_flat_vehicle]).reshape(-1, 3) dims = np.array([b.wlh for b in boxes_flat_vehicle]).reshape(-1, 3) rots = np.array([b.orientation.yaw_pitch_roll[0] for b in boxes_flat_vehicle]).reshape(-1, 1) gt_boxes = np.concatenate([locs, dims, rots], axis=1) gt_names = np.array([b.name for b in boxes_flat_vehicle]) if use_second_format_direction: gt_boxes[:, 6] = np.apply_along_axis( lambda r: -ds_utils.wrap_to_pi(r + np.pi / 2), axis=1, arr=gt_boxes[:, 6:]) annotations = [ lyftdata.get('sample_annotation', token) for token in sample['anns'] ] assert len(gt_boxes) == len(annotations), f"{len(gt_boxes)} != {len(annotations)}" info['gt_boxes'] = gt_boxes info['gt_names'] = gt_names if sample['scene_token'] in train_scene_tokens: train_infos.append(info) else: val_infos.append(info) return train_infos, val_infos
def create_groundtruth_database_with_sweep_info(dataset_class_name, data_path, info_path=None, used_classes=None, database_save_path=None, db_info_save_path=None, relative_path=True, add_rgb=False, lidar_only=False, bev_only=False, coors_range=None): dataset = get_dataset_class(dataset_class_name)( info_path=info_path, root_path=data_path, ) root_path = Path(data_path) if database_save_path is None: database_save_path = root_path / 'gt_database_with_sweep_info' else: database_save_path = Path(database_save_path) if db_info_save_path is None: db_info_save_path = root_path / "dbinfos_train_with_sweep_info.pkl" database_save_path.mkdir(parents=True, exist_ok=True) all_db_infos = {} group_counter = 0 for j in prog_bar(list(range(len(dataset)))): image_idx = j sensor_data = dataset.get_sensor_data(j) if "image_idx" in sensor_data["metadata"]: image_idx = sensor_data["metadata"]["image_idx"] assert("sweep_points_list" in sensor_data["lidar"]) sweep_points_list = sensor_data["lidar"]["sweep_points_list"] assert("sweep_annotations" in sensor_data["lidar"]) sweep_gt_boxes_list = sensor_data["lidar"]["sweep_annotations"]["boxes_list"] sweep_gt_tokens_list = sensor_data["lidar"]["sweep_annotations"]["tokens_list"] sweep_gt_names_list = sensor_data["lidar"]["sweep_annotations"]["names_list"] # we focus on the objects in the first frame # and find the bounding box index of the same object in every frame points = sensor_data["lidar"]["points"] annos = sensor_data["lidar"]["annotations"] names = annos["names"] gt_boxes = annos["boxes"] attrs = annos["attrs"] tokens = sweep_gt_tokens_list[0] # sanity check with redundancy assert(len(sweep_gt_boxes_list) == len(sweep_gt_tokens_list) == len(sweep_gt_names_list)) assert(len(gt_boxes) == len(attrs) == len(tokens)) assert(np.all(names == sweep_gt_names_list[0])) assert(np.all(gt_boxes == sweep_gt_boxes_list[0])) # on every frame, we mask points inside each bounding box # but we are not looking at every bounding box sweep_point_indices_list = [] for sweep_points, sweep_gt_boxes in zip(sweep_points_list, sweep_gt_boxes_list): sweep_point_indices_list.append(box_np_ops.points_in_rbbox(sweep_points, sweep_gt_boxes)) # crop point cloud based on boxes in the current frame point_indices = box_np_ops.points_in_rbbox(points, gt_boxes) # 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] for i in range(num_obj): filename = f"{image_idx}_{names[i]}_{i}.bin" filepath = database_save_path / filename # only use non-key frame boxes when the object is moving if attrs[i] in ['vehicle.moving', 'cycle.with_rider', 'pedestrian.moving']: gt_points_list = [] for t in range(len(sweep_points_list)): # fast pass for most frames if i < len(sweep_gt_tokens_list[t]) and sweep_gt_tokens_list[t][i] == tokens[i]: box_idx = i else: I = np.flatnonzero(tokens[i] == sweep_gt_tokens_list[t]) if len(I) == 0: continue elif len(I) == 1: box_idx = I[0] else: raise ValueError('Identical object tokens') gt_points_list.append(sweep_points_list[t][sweep_point_indices_list[t][:, box_idx]]) gt_points = np.concatenate(gt_points_list, axis=0)[:, [0, 1, 2, 4]] else: gt_points = points[point_indices[:, i]] # offset points based on the bounding box in the current frame gt_points[:, :3] -= gt_boxes[i, :3] with open(filepath, 'w') as f: gt_points.tofile(f) if (used_classes is None) or 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": i, "box3d_lidar": gt_boxes[i], "num_points_in_gt": gt_points.shape[0], "difficulty": difficulty[i] } 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] 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] 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 _fill_train_infos(root_path): train_udi_infos = [] lidar_root_path = root_path+ "/lidar" label_root_path = root_path + "/label" img_root_path = root_path + "/image" filenames = os.listdir(lidar_root_path) for filename in prog_bar(filenames): index = filename.split(".")[0] lidar_path = lidar_root_path + "/" + index + ".bin" cam_path = img_root_path + "/" + index + ".jpg" label_path = label_root_path + "/" + index + "_bin.json" assert Path(lidar_path).exists() assert Path(cam_path).exists() assert Path(label_path).exists() with open(label_path, encoding='utf-8') as f: res = f.read() result = json.loads(res) boxes = result["elem"] info = { "lidar_path": lidar_path, "cam_front_path": cam_path, "filename": filename, "token": int(index), } gt_locs_list = [] gt_dims_list = [] print("label file path:", label_path) for box in boxes: box_loc = box["position"] box_size = box["size"] box_loc_ = np.array([box_loc["x"],box_loc["y"], box_loc["z"]], dtype=np.float) box_size_ = np.array([box_size["width"],box_size["depth"],box_size["height"]], dtype=np.float) box_loc_ = box_loc_.reshape(-1, 3) box_size_ = box_size_.reshape(-1, 3) gt_locs_list.append(box_loc_) gt_dims_list.append(box_size_) locs = np.concatenate(gt_locs_list, axis=0) dims = np.concatenate(gt_dims_list, axis=0) rots = np.array([b["yaw"] for b in boxes]).reshape(-1, 1) names = [b["class"] for b in boxes] for i in range(len(names)): if names[i] in UDIDataset.NameMapping: names[i] = UDIDataset.NameMapping[names[i]] names = np.array(names) # we need to convert rot to SECOND format. # change the rot format will break all checkpoint. gt_boxes = np.concatenate([locs, dims, -rots - np.pi / 2], axis=1) info["gt_boxes"] = gt_boxes info["gt_names"] = names train_udi_infos.append(info) return train_udi_infos