import pickle import copy import numpy as np from skimage import io from ...utils import box_utils, common_utils, calibration_kitti, object3d_kitti from ..dataset import DatasetTemplate class KittiDataset(DatasetTemplate): def __init__(self, dataset_cfg, class_names, training=True, root_path=None, logger=None): """ Args: root_path: dataset_cfg: class_names: training: logger: """ super().__init__( dataset_cfg=dataset_cfg, class_names=class_names, training=training, root_path=root_path, logger=logger ) self.split = self.dataset_cfg.DATA_SPLIT[self.mode] self.root_split_path = self.root_path / ('training' if self.split != 'test' else 'testing') split_dir = self.root_path / 'ImageSets' / (self.split + '.txt') self.sample_id_list = [x.strip() for x in open(split_dir).readlines()] if split_dir.exists() else None self.kitti_infos = [] self.include_kitti_data(self.mode) def include_kitti_data(self, mode): if self.logger is not None: self.logger.info('Loading KITTI dataset') kitti_infos = [] for info_path in self.dataset_cfg.INFO_PATH[mode]: info_path = self.root_path / info_path with open(info_path, 'rb') as f: infos = pickle.load(f) kitti_infos.extend(infos) self.kitti_infos.extend(kitti_infos) if self.logger is not None: self.logger.info('Total samples for KITTI dataset: %d' % (len(kitti_infos))) def get_lidar(self, idx): lidar_file = self.root_split_path / 'velodyne' / ('%s.bin' % idx) assert lidar_file.exists() return np.fromfile(str(lidar_file), dtype=np.float32).reshape(-1, 4) def get_image_shape(self, idx): img_file = self.root_split_path / 'image_2' / ('%s.png' % idx) assert img_file.exists() return np.array(io.imread(img_file).shape[:2], dtype=np.int32) def get_label(self, idx): label_file = self.root_split_path / 'label_2' / ('%s.txt' % idx) assert label_file.exists() return object3d_kitti.get_objects_from_label(label_file) def get_calib(self, idx): calib_file = self.root_split_path / 'calib' / ('%s.txt' % idx) assert calib_file.exists() return calibration_kitti.Calibration(calib_file) def get_road_plane(self, idx): plane_file = self.root_split_path / 'planes' / ('%s.txt' % idx) if not plane_file.exists(): return None with open(plane_file, 'r') as f: lines = f.readlines() lines = [float(i) for i in lines[3].split()] plane = np.asarray(lines) # Ensure normal is always facing up, this is in the rectified camera coordinate if plane[1] > 0: plane = -plane norm = np.linalg.norm(plane[0:3]) plane = plane / norm return plane @staticmethod def get_fov_flag(pts_rect, img_shape, calib): """ Args: pts_rect: img_shape: calib: Returns: """ pts_img, pts_rect_depth = calib.rect_to_img(pts_rect) val_flag_1 = np.logical_and(pts_img[:, 0] >= 0, pts_img[:, 0] < img_shape[1]) val_flag_2 = np.logical_and(pts_img[:, 1] >= 0, pts_img[:, 1] < img_shape[0]) val_flag_merge = np.logical_and(val_flag_1, val_flag_2) pts_valid_flag = np.logical_and(val_flag_merge, pts_rect_depth >= 0) return pts_valid_flag @staticmethod def generate_prediction_dicts(batch_dict, pred_dicts, class_names, output_path=None): """ Args: batch_dict: frame_id: pred_dicts: list of pred_dicts pred_boxes: (N, 7), Tensor pred_scores: (N), Tensor pred_labels: (N), Tensor class_names: output_path: Returns: """ def get_template_prediction(num_samples): ret_dict = { 'name': np.zeros(num_samples), 'truncated': np.zeros(num_samples), 'occluded': np.zeros(num_samples), 'alpha': np.zeros(num_samples), 'bbox': np.zeros([num_samples, 4]), 'dimensions': np.zeros([num_samples, 3]), 'location': np.zeros([num_samples, 3]), 'rotation_y': np.zeros(num_samples), 'score': np.zeros(num_samples), 'boxes_lidar': np.zeros([num_samples, 7]) } return ret_dict def generate_single_sample_dict(batch_index, box_dict): pred_scores = box_dict['pred_scores'].cpu().numpy() pred_boxes = box_dict['pred_boxes'].cpu().numpy() pred_labels = box_dict['pred_labels'].cpu().numpy() pred_dict = get_template_prediction(pred_scores.shape[0]) if pred_scores.shape[0] == 0: return pred_dict calib = batch_dict['calib'][batch_index] image_shape = batch_dict['image_shape'][batch_index] pred_boxes_camera = box_utils.boxes3d_lidar_to_kitti_camera(pred_boxes, calib) pred_boxes_img = box_utils.boxes3d_kitti_camera_to_imageboxes( pred_boxes_camera, calib, image_shape=image_shape ) pred_dict['name'] = np.array(class_names)[pred_labels - 1] pred_dict['alpha'] = -np.arctan2(-pred_boxes[:, 1], pred_boxes[:, 0]) + pred_boxes_camera[:, 6] pred_dict['bbox'] = pred_boxes_img pred_dict['dimensions'] = pred_boxes_camera[:, 3:6] pred_dict['location'] = pred_boxes_camera[:, 0:3] pred_dict['rotation_y'] = pred_boxes_camera[:, 6] pred_dict['score'] = pred_scores pred_dict['boxes_lidar'] = pred_boxes return pred_dict annos = [] for index, box_dict in enumerate(pred_dicts): frame_id = batch_dict['frame_id'][index] single_pred_dict = generate_single_sample_dict(index, box_dict) single_pred_dict['frame_id'] = frame_id annos.append(single_pred_dict) if output_path is not None: cur_det_file = output_path / ('%s.txt' % frame_id) with open(cur_det_file, 'w') as f: bbox = single_pred_dict['bbox'] loc = single_pred_dict['location'] dims = single_pred_dict['dimensions'] # lhw -> hwl for idx in range(len(bbox)): print('%s -1 -1 %.4f %.4f %.4f %.4f %.4f %.4f %.4f %.4f %.4f %.4f %.4f %.4f %.4f' % (single_pred_dict['name'][idx], single_pred_dict['alpha'][idx], bbox[idx][0], bbox[idx][1], bbox[idx][2], bbox[idx][3], dims[idx][1], dims[idx][2], dims[idx][0], loc[idx][0], loc[idx][1], loc[idx][2], single_pred_dict['rotation_y'][idx], single_pred_dict['score'][idx]), file=f) return annos def evaluation(self, det_annos, class_names, **kwargs): assert 'annos' in self.kitti_infos[0].keys() from .kitti_object_eval_python import eval as kitti_eval if 'annos' not in self.kitti_infos[0]: return 'None', {} eval_det_annos = copy.deepcopy(det_annos) eval_gt_annos = [copy.deepcopy(info['annos']) for info in self.kitti_infos] ap_result_str, ap_dict = kitti_eval.get_official_eval_result(eval_gt_annos, eval_det_annos, class_names) return ap_result_str, ap_dict def __len__(self): return len(self.kitti_infos) def __getitem__(self, index): # index = 4 info = copy.deepcopy(self.kitti_infos[index]) sample_idx = info['point_cloud']['lidar_idx'] points = self.get_lidar(sample_idx) calib = self.get_calib(sample_idx) img_shape = info['image']['image_shape'] if self.dataset_cfg.FOV_POINTS_ONLY: pts_rect = calib.lidar_to_rect(points[:, 0:3]) fov_flag = self.get_fov_flag(pts_rect, img_shape, calib) points = points[fov_flag] input_dict = { 'points': points, 'frame_id': sample_idx, 'calib': calib, } if 'annos' in info: annos = info['annos'] annos = common_utils.drop_info_with_name(annos, name='DontCare') loc, dims, rots = annos['location'], annos['dimensions'], annos['rotation_y'] gt_names = annos['name'] gt_boxes_camera = np.concatenate([loc, dims, rots[..., np.newaxis]], axis=1).astype(np.float32) gt_boxes_lidar = box_utils.boxes3d_kitti_camera_to_lidar(gt_boxes_camera, calib) input_dict.update({ 'gt_names': gt_names, 'gt_boxes': gt_boxes_lidar }) road_plane = self.get_road_plane(sample_idx) if road_plane is not None: input_dict['road_plane'] = road_plane data_dict = self.prepare_data(data_dict=input_dict) data_dict['image_shape'] = img_shape return data_dict def create_kitti_infos(data_path, save_path, workers=4): pass if __name__ == '__main__': pass