# OpenPCDet PyTorch Dataloader and Evaluation Tools for Waymo Open Dataset # Reference https://github.com/open-mmlab/OpenPCDet # Written by Shaoshuai Shi, Chaoxu Guo # All Rights Reserved. import os import pickle import copy import numpy as np import torch import multiprocessing import SharedArray import torch.distributed as dist from tqdm import tqdm from pathlib import Path from functools import partial from ...ops.roiaware_pool3d import roiaware_pool3d_utils from ...utils import box_utils, common_utils from ..dataset import DatasetTemplate class WaymoDataset(DatasetTemplate): def __init__(self, dataset_cfg, class_names, training=True, root_path=None, logger=None): super().__init__( dataset_cfg=dataset_cfg, class_names=class_names, training=training, root_path=root_path, logger=logger ) self.data_path = self.root_path / self.dataset_cfg.PROCESSED_DATA_TAG self.split = self.dataset_cfg.DATA_SPLIT[self.mode] split_dir = self.root_path / 'ImageSets' / (self.split + '.txt') self.sample_sequence_list = [x.strip() for x in open(split_dir).readlines()] self.infos = [] self.seq_name_to_infos = self.include_waymo_data(self.mode) self.use_shared_memory = self.dataset_cfg.get('USE_SHARED_MEMORY', False) and self.training if self.use_shared_memory: self.shared_memory_file_limit = self.dataset_cfg.get('SHARED_MEMORY_FILE_LIMIT', 0x7FFFFFFF) self.load_data_to_shared_memory() if self.dataset_cfg.get('USE_PREDBOX', False): self.pred_boxes_dict = self.load_pred_boxes_to_dict( pred_boxes_path=self.dataset_cfg.ROI_BOXES_PATH[self.mode] ) else: self.pred_boxes_dict = {} def set_split(self, split): super().__init__( dataset_cfg=self.dataset_cfg, class_names=self.class_names, training=self.training, root_path=self.root_path, logger=self.logger ) self.split = split split_dir = self.root_path / 'ImageSets' / (self.split + '.txt') self.sample_sequence_list = [x.strip() for x in open(split_dir).readlines()] self.infos = [] self.seq_name_to_infos = self.include_waymo_data(self.mode) def include_waymo_data(self, mode): self.logger.info('Loading Waymo dataset') waymo_infos = [] seq_name_to_infos = {} num_skipped_infos = 0 for k in range(len(self.sample_sequence_list)): sequence_name = os.path.splitext(self.sample_sequence_list[k])[0] info_path = self.data_path / sequence_name / ('%s.pkl' % sequence_name) info_path = self.check_sequence_name_with_all_version(info_path) if not info_path.exists(): num_skipped_infos += 1 continue with open(info_path, 'rb') as f: infos = pickle.load(f) waymo_infos.extend(infos) seq_name_to_infos[infos[0]['point_cloud']['lidar_sequence']] = infos self.infos.extend(waymo_infos[:]) self.logger.info('Total skipped info %s' % num_skipped_infos) self.logger.info('Total samples for Waymo dataset: %d' % (len(waymo_infos))) if self.dataset_cfg.SAMPLED_INTERVAL[mode] > 1: sampled_waymo_infos = [] for k in range(0, len(self.infos), self.dataset_cfg.SAMPLED_INTERVAL[mode]): sampled_waymo_infos.append(self.infos[k]) self.infos = sampled_waymo_infos self.logger.info('Total sampled samples for Waymo dataset: %d' % len(self.infos)) use_sequence_data = self.dataset_cfg.get('SEQUENCE_CONFIG', None) is not None and self.dataset_cfg.SEQUENCE_CONFIG.ENABLED if not use_sequence_data: seq_name_to_infos = None return seq_name_to_infos def load_pred_boxes_to_dict(self, pred_boxes_path): self.logger.info(f'Loading and reorganizing pred_boxes to dict from path: {pred_boxes_path}') with open(pred_boxes_path, 'rb') as f: pred_dicts = pickle.load(f) pred_boxes_dict = {} for index, box_dict in enumerate(pred_dicts): seq_name = box_dict['frame_id'][:-4].replace('training_', '').replace('validation_', '') sample_idx = int(box_dict['frame_id'][-3:]) if seq_name not in pred_boxes_dict: pred_boxes_dict[seq_name] = {} pred_labels = np.array([self.class_names.index(box_dict['name'][k]) + 1 for k in range(box_dict['name'].shape[0])]) pred_boxes = np.concatenate((box_dict['boxes_lidar'], box_dict['score'][:, np.newaxis], pred_labels[:, np.newaxis]), axis=-1) pred_boxes_dict[seq_name][sample_idx] = pred_boxes self.logger.info(f'Predicted boxes has been loaded, total sequences: {len(pred_boxes_dict)}') return pred_boxes_dict def load_data_to_shared_memory(self): self.logger.info(f'Loading training data to shared memory (file limit={self.shared_memory_file_limit})') cur_rank, num_gpus = common_utils.get_dist_info() all_infos = self.infos[:self.shared_memory_file_limit] \ if self.shared_memory_file_limit < len(self.infos) else self.infos cur_infos = all_infos[cur_rank::num_gpus] for info in cur_infos: pc_info = info['point_cloud'] sequence_name = pc_info['lidar_sequence'] sample_idx = pc_info['sample_idx'] sa_key = f'{sequence_name}___{sample_idx}' if os.path.exists(f"/dev/shm/{sa_key}"): continue points = self.get_lidar(sequence_name, sample_idx) common_utils.sa_create(f"shm://{sa_key}", points) dist.barrier() self.logger.info('Training data has been saved to shared memory') def clean_shared_memory(self): self.logger.info(f'Clean training data from shared memory (file limit={self.shared_memory_file_limit})') cur_rank, num_gpus = common_utils.get_dist_info() all_infos = self.infos[:self.shared_memory_file_limit] \ if self.shared_memory_file_limit < len(self.infos) else self.infos cur_infos = all_infos[cur_rank::num_gpus] for info in cur_infos: pc_info = info['point_cloud'] sequence_name = pc_info['lidar_sequence'] sample_idx = pc_info['sample_idx'] sa_key = f'{sequence_name}___{sample_idx}' if not os.path.exists(f"/dev/shm/{sa_key}"): continue SharedArray.delete(f"shm://{sa_key}") if num_gpus > 1: dist.barrier() self.logger.info('Training data has been deleted from shared memory') @staticmethod def check_sequence_name_with_all_version(sequence_file): if not sequence_file.exists(): found_sequence_file = sequence_file for pre_text in ['training', 'validation', 'testing']: if not sequence_file.exists(): temp_sequence_file = Path(str(sequence_file).replace('segment', pre_text + '_segment')) if temp_sequence_file.exists(): found_sequence_file = temp_sequence_file break if not found_sequence_file.exists(): found_sequence_file = Path(str(sequence_file).replace('_with_camera_labels', '')) if found_sequence_file.exists(): sequence_file = found_sequence_file return sequence_file def get_infos(self, raw_data_path, save_path, num_workers=multiprocessing.cpu_count(), has_label=True, sampled_interval=1, update_info_only=False): from . import waymo_utils print('---------------The waymo sample interval is %d, total sequecnes is %d-----------------' % (sampled_interval, len(self.sample_sequence_list))) process_single_sequence = partial( waymo_utils.process_single_sequence, save_path=save_path, sampled_interval=sampled_interval, has_label=has_label, update_info_only=update_info_only ) sample_sequence_file_list = [ self.check_sequence_name_with_all_version(raw_data_path / sequence_file) for sequence_file in self.sample_sequence_list ] # process_single_sequence(sample_sequence_file_list[0]) with multiprocessing.Pool(num_workers) as p: sequence_infos = list(tqdm(p.imap(process_single_sequence, sample_sequence_file_list), total=len(sample_sequence_file_list))) all_sequences_infos = [item for infos in sequence_infos for item in infos] return all_sequences_infos def get_lidar(self, sequence_name, sample_idx): lidar_file = self.data_path / sequence_name / ('%04d.npy' % sample_idx) point_features = np.load(lidar_file) # (N, 7): [x, y, z, intensity, elongation, NLZ_flag] points_all, NLZ_flag = point_features[:, 0:5], point_features[:, 5] if not self.dataset_cfg.get('DISABLE_NLZ_FLAG_ON_POINTS', False): points_all = points_all[NLZ_flag == -1] if self.dataset_cfg.get('POINTS_TANH_DIM', None) is None: points_all[:, 3] = np.tanh(points_all[:, 3]) else: for dim_idx in self.dataset_cfg.POINTS_TANH_DIM: points_all[:, dim_idx] = np.tanh(points_all[:, dim_idx]) return points_all @staticmethod def transform_prebox_to_current(pred_boxes3d, pose_pre, pose_cur): """ Args: pred_boxes3d (N, 9 or 11): [x, y, z, dx, dy, dz, raw, score, label] pose_pre (4, 4): pose_cur (4, 4): Returns: """ assert pred_boxes3d.shape[-1] in [9, 11] pred_boxes3d = pred_boxes3d.copy() expand_bboxes = np.concatenate([pred_boxes3d[:, :3], np.ones((pred_boxes3d.shape[0], 1))], axis=-1) bboxes_global = np.dot(expand_bboxes, pose_pre.T)[:, :3] expand_bboxes_global = np.concatenate([bboxes_global[:, :3],np.ones((bboxes_global.shape[0], 1))], axis=-1) bboxes_pre2cur = np.dot(expand_bboxes_global, np.linalg.inv(pose_cur.T))[:, :3] pred_boxes3d[:, 0:3] = bboxes_pre2cur if pred_boxes3d.shape[-1] == 11: expand_vels = np.concatenate([pred_boxes3d[:, 7:9], np.zeros((pred_boxes3d.shape[0], 1))], axis=-1) vels_global = np.dot(expand_vels, pose_pre[:3, :3].T) vels_pre2cur = np.dot(vels_global, np.linalg.inv(pose_cur[:3, :3].T))[:,:2] pred_boxes3d[:, 7:9] = vels_pre2cur pred_boxes3d[:, 6] = pred_boxes3d[..., 6] + np.arctan2(pose_pre[..., 1, 0], pose_pre[..., 0, 0]) pred_boxes3d[:, 6] = pred_boxes3d[..., 6] - np.arctan2(pose_cur[..., 1, 0], pose_cur[..., 0, 0]) return pred_boxes3d @staticmethod def reorder_rois_for_refining(pred_bboxes): num_max_rois = max([len(bbox) for bbox in pred_bboxes]) num_max_rois = max(1, num_max_rois) # at least one faked rois to avoid error ordered_bboxes = np.zeros([len(pred_bboxes), num_max_rois, pred_bboxes[0].shape[-1]], dtype=np.float32) for bs_idx in range(ordered_bboxes.shape[0]): ordered_bboxes[bs_idx, :len(pred_bboxes[bs_idx])] = pred_bboxes[bs_idx] return ordered_bboxes def get_sequence_data(self, info, points, sequence_name, sample_idx, sequence_cfg, load_pred_boxes=False): """ Args: info: points: sequence_name: sample_idx: sequence_cfg: Returns: """ def remove_ego_points(points, center_radius=1.0): mask = ~((np.abs(points[:, 0]) < center_radius) & (np.abs(points[:, 1]) < center_radius)) return points[mask] def load_pred_boxes_from_dict(sequence_name, sample_idx): """ boxes: (N, 11) [x, y, z, dx, dy, dn, raw, vx, vy, score, label] """ sequence_name = sequence_name.replace('training_', '').replace('validation_', '') load_boxes = self.pred_boxes_dict[sequence_name][sample_idx] assert load_boxes.shape[-1] == 11 load_boxes[:, 7:9] = -0.1 * load_boxes[:, 7:9] # transfer speed to negtive motion from t to t-1 return load_boxes pose_cur = info['pose'].reshape((4, 4)) num_pts_cur = points.shape[0] sample_idx_pre_list = np.clip(sample_idx + np.arange(sequence_cfg.SAMPLE_OFFSET[0], sequence_cfg.SAMPLE_OFFSET[1]), 0, 0x7FFFFFFF) sample_idx_pre_list = sample_idx_pre_list[::-1] if sequence_cfg.get('ONEHOT_TIMESTAMP', False): onehot_cur = np.zeros((points.shape[0], len(sample_idx_pre_list) + 1)).astype(points.dtype) onehot_cur[:, 0] = 1 points = np.hstack([points, onehot_cur]) else: points = np.hstack([points, np.zeros((points.shape[0], 1)).astype(points.dtype)]) points_pre_all = [] num_points_pre = [] pose_all = [pose_cur] pred_boxes_all = [] if load_pred_boxes: pred_boxes = load_pred_boxes_from_dict(sequence_name, sample_idx) pred_boxes_all.append(pred_boxes) sequence_info = self.seq_name_to_infos[sequence_name] for idx, sample_idx_pre in enumerate(sample_idx_pre_list): points_pre = self.get_lidar(sequence_name, sample_idx_pre) pose_pre = sequence_info[sample_idx_pre]['pose'].reshape((4, 4)) expand_points_pre = np.concatenate([points_pre[:, :3], np.ones((points_pre.shape[0], 1))], axis=-1) points_pre_global = np.dot(expand_points_pre, pose_pre.T)[:, :3] expand_points_pre_global = np.concatenate([points_pre_global, np.ones((points_pre_global.shape[0], 1))], axis=-1) points_pre2cur = np.dot(expand_points_pre_global, np.linalg.inv(pose_cur.T))[:, :3] points_pre = np.concatenate([points_pre2cur, points_pre[:, 3:]], axis=-1) if sequence_cfg.get('ONEHOT_TIMESTAMP', False): onehot_vector = np.zeros((points_pre.shape[0], len(sample_idx_pre_list) + 1)) onehot_vector[:, idx + 1] = 1 points_pre = np.hstack([points_pre, onehot_vector]) else: # add timestamp points_pre = np.hstack([points_pre, 0.1 * (sample_idx - sample_idx_pre) * np.ones((points_pre.shape[0], 1)).astype(points_pre.dtype)]) # one frame 0.1s points_pre = remove_ego_points(points_pre, 1.0) points_pre_all.append(points_pre) num_points_pre.append(points_pre.shape[0]) pose_all.append(pose_pre) if load_pred_boxes: pose_pre = sequence_info[sample_idx_pre]['pose'].reshape((4, 4)) pred_boxes = load_pred_boxes_from_dict(sequence_name, sample_idx_pre) pred_boxes = self.transform_prebox_to_current(pred_boxes, pose_pre, pose_cur) pred_boxes_all.append(pred_boxes) points = np.concatenate([points] + points_pre_all, axis=0).astype(np.float32) num_points_all = np.array([num_pts_cur] + num_points_pre).astype(np.int32) poses = np.concatenate(pose_all, axis=0).astype(np.float32) if load_pred_boxes: temp_pred_boxes = self.reorder_rois_for_refining(pred_boxes_all) pred_boxes = temp_pred_boxes[:, :, 0:9] pred_scores = temp_pred_boxes[:, :, 9] pred_labels = temp_pred_boxes[:, :, 10] else: pred_boxes = pred_scores = pred_labels = None return points, num_points_all, sample_idx_pre_list, poses, pred_boxes, pred_scores, pred_labels def __len__(self): if self._merge_all_iters_to_one_epoch: return len(self.infos) * self.total_epochs return len(self.infos) def __getitem__(self, index): if self._merge_all_iters_to_one_epoch: index = index % len(self.infos) info = copy.deepcopy(self.infos[index]) pc_info = info['point_cloud'] sequence_name = pc_info['lidar_sequence'] sample_idx = pc_info['sample_idx'] input_dict = { 'sample_idx': sample_idx } if self.use_shared_memory and index < self.shared_memory_file_limit: sa_key = f'{sequence_name}___{sample_idx}' points = SharedArray.attach(f"shm://{sa_key}").copy() else: points = self.get_lidar(sequence_name, sample_idx) if self.dataset_cfg.get('SEQUENCE_CONFIG', None) is not None and self.dataset_cfg.SEQUENCE_CONFIG.ENABLED: points, num_points_all, sample_idx_pre_list, poses, pred_boxes, pred_scores, pred_labels = self.get_sequence_data( info, points, sequence_name, sample_idx, self.dataset_cfg.SEQUENCE_CONFIG, load_pred_boxes=self.dataset_cfg.get('USE_PREDBOX', False) ) input_dict['poses'] = poses if self.dataset_cfg.get('USE_PREDBOX', False): input_dict.update({ 'roi_boxes': pred_boxes, 'roi_scores': pred_scores, 'roi_labels': pred_labels, }) input_dict.update({ 'points': points, 'frame_id': info['frame_id'], }) if 'annos' in info: annos = info['annos'] annos = common_utils.drop_info_with_name(annos, name='unknown') if self.dataset_cfg.get('INFO_WITH_FAKELIDAR', False): gt_boxes_lidar = box_utils.boxes3d_kitti_fakelidar_to_lidar(annos['gt_boxes_lidar']) else: gt_boxes_lidar = annos['gt_boxes_lidar'] if self.dataset_cfg.get('TRAIN_WITH_SPEED', False): assert gt_boxes_lidar.shape[-1] == 9 else: gt_boxes_lidar = gt_boxes_lidar[:, 0:7] if self.training and self.dataset_cfg.get('FILTER_EMPTY_BOXES_FOR_TRAIN', False): mask = (annos['num_points_in_gt'] > 0) # filter empty boxes annos['name'] = annos['name'][mask] gt_boxes_lidar = gt_boxes_lidar[mask] annos['num_points_in_gt'] = annos['num_points_in_gt'][mask] input_dict.update({ 'gt_names': annos['name'], 'gt_boxes': gt_boxes_lidar, 'num_points_in_gt': annos.get('num_points_in_gt', None) }) data_dict = self.prepare_data(data_dict=input_dict) data_dict['metadata'] = info.get('metadata', info['frame_id']) data_dict.pop('num_points_in_gt', None) return data_dict def evaluation(self, det_annos, class_names, **kwargs): if 'annos' not in self.infos[0].keys(): return 'No ground-truth boxes for evaluation', {} def kitti_eval(eval_det_annos, eval_gt_annos): from ..kitti.kitti_object_eval_python import eval as kitti_eval from ..kitti import kitti_utils map_name_to_kitti = { 'Vehicle': 'Car', 'Pedestrian': 'Pedestrian', 'Cyclist': 'Cyclist', 'Sign': 'Sign', 'Car': 'Car' } kitti_utils.transform_annotations_to_kitti_format(eval_det_annos, map_name_to_kitti=map_name_to_kitti) kitti_utils.transform_annotations_to_kitti_format( eval_gt_annos, map_name_to_kitti=map_name_to_kitti, info_with_fakelidar=self.dataset_cfg.get('INFO_WITH_FAKELIDAR', False) ) kitti_class_names = [map_name_to_kitti[x] for x in class_names] ap_result_str, ap_dict = kitti_eval.get_official_eval_result( gt_annos=eval_gt_annos, dt_annos=eval_det_annos, current_classes=kitti_class_names ) return ap_result_str, ap_dict def waymo_eval(eval_det_annos, eval_gt_annos): from .waymo_eval import OpenPCDetWaymoDetectionMetricsEstimator eval = OpenPCDetWaymoDetectionMetricsEstimator() ap_dict = eval.waymo_evaluation( eval_det_annos, eval_gt_annos, class_name=class_names, distance_thresh=1000, fake_gt_infos=self.dataset_cfg.get('INFO_WITH_FAKELIDAR', False) ) ap_result_str = '\n' for key in ap_dict: ap_dict[key] = ap_dict[key][0] ap_result_str += '%s: %.4f \n' % (key, ap_dict[key]) return ap_result_str, ap_dict eval_det_annos = copy.deepcopy(det_annos) eval_gt_annos = [copy.deepcopy(info['annos']) for info in self.infos] if kwargs['eval_metric'] == 'kitti': ap_result_str, ap_dict = kitti_eval(eval_det_annos, eval_gt_annos) elif kwargs['eval_metric'] == 'waymo': ap_result_str, ap_dict = waymo_eval(eval_det_annos, eval_gt_annos) else: raise NotImplementedError return ap_result_str, ap_dict def create_groundtruth_database(self, info_path, save_path, used_classes=None, split='train', sampled_interval=10, processed_data_tag=None): use_sequence_data = self.dataset_cfg.get('SEQUENCE_CONFIG', None) is not None and self.dataset_cfg.SEQUENCE_CONFIG.ENABLED if use_sequence_data: st_frame, ed_frame = self.dataset_cfg.SEQUENCE_CONFIG.SAMPLE_OFFSET[0], self.dataset_cfg.SEQUENCE_CONFIG.SAMPLE_OFFSET[1] self.dataset_cfg.SEQUENCE_CONFIG.SAMPLE_OFFSET[0] = min(-4, st_frame) # at least we use 5 frames for generating gt database to support various sequence configs (<= 5 frames) st_frame = self.dataset_cfg.SEQUENCE_CONFIG.SAMPLE_OFFSET[0] database_save_path = save_path / ('%s_gt_database_%s_sampled_%d_multiframe_%s_to_%s' % (processed_data_tag, split, sampled_interval, st_frame, ed_frame)) db_info_save_path = save_path / ('%s_waymo_dbinfos_%s_sampled_%d_multiframe_%s_to_%s.pkl' % (processed_data_tag, split, sampled_interval, st_frame, ed_frame)) db_data_save_path = save_path / ('%s_gt_database_%s_sampled_%d_multiframe_%s_to_%s_global.npy' % (processed_data_tag, split, sampled_interval, st_frame, ed_frame)) else: database_save_path = save_path / ('%s_gt_database_%s_sampled_%d' % (processed_data_tag, split, sampled_interval)) db_info_save_path = save_path / ('%s_waymo_dbinfos_%s_sampled_%d.pkl' % (processed_data_tag, split, sampled_interval)) db_data_save_path = save_path / ('%s_gt_database_%s_sampled_%d_global.npy' % (processed_data_tag, split, sampled_interval)) database_save_path.mkdir(parents=True, exist_ok=True) all_db_infos = {} with open(info_path, 'rb') as f: infos = pickle.load(f) point_offset_cnt = 0 stacked_gt_points = [] for k in tqdm(range(0, len(infos), sampled_interval)): # print('gt_database sample: %d/%d' % (k + 1, len(infos))) info = infos[k] pc_info = info['point_cloud'] sequence_name = pc_info['lidar_sequence'] sample_idx = pc_info['sample_idx'] points = self.get_lidar(sequence_name, sample_idx) if use_sequence_data: points, num_points_all, sample_idx_pre_list, _, _, _, _ = self.get_sequence_data( info, points, sequence_name, sample_idx, self.dataset_cfg.SEQUENCE_CONFIG ) annos = info['annos'] names = annos['name'] difficulty = annos['difficulty'] gt_boxes = annos['gt_boxes_lidar'] if k % 4 != 0 and len(names) > 0: mask = (names == 'Vehicle') names = names[~mask] difficulty = difficulty[~mask] gt_boxes = gt_boxes[~mask] if k % 2 != 0 and len(names) > 0: mask = (names == 'Pedestrian') names = names[~mask] difficulty = difficulty[~mask] gt_boxes = gt_boxes[~mask] num_obj = gt_boxes.shape[0] if num_obj == 0: continue box_idxs_of_pts = roiaware_pool3d_utils.points_in_boxes_gpu( torch.from_numpy(points[:, 0:3]).unsqueeze(dim=0).float().cuda(), torch.from_numpy(gt_boxes[:, 0:7]).unsqueeze(dim=0).float().cuda() ).long().squeeze(dim=0).cpu().numpy() for i in range(num_obj): filename = '%s_%04d_%s_%d.bin' % (sequence_name, sample_idx, names[i], i) filepath = database_save_path / filename gt_points = points[box_idxs_of_pts == i] gt_points[:, :3] -= gt_boxes[i, :3] if (used_classes is None) or names[i] in used_classes: gt_points = gt_points.astype(np.float32) assert gt_points.dtype == np.float32 with open(filepath, 'w') as f: gt_points.tofile(f) db_path = str(filepath.relative_to(self.root_path)) # gt_database/xxxxx.bin db_info = {'name': names[i], 'path': db_path, 'sequence_name': sequence_name, 'sample_idx': sample_idx, 'gt_idx': i, 'box3d_lidar': gt_boxes[i], 'num_points_in_gt': gt_points.shape[0], 'difficulty': difficulty[i]} # it will be used if you choose to use shared memory for gt sampling stacked_gt_points.append(gt_points) db_info['global_data_offset'] = [point_offset_cnt, point_offset_cnt + gt_points.shape[0]] point_offset_cnt += gt_points.shape[0] 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('Database %s: %d' % (k, len(v))) with open(db_info_save_path, 'wb') as f: pickle.dump(all_db_infos, f) # it will be used if you choose to use shared memory for gt sampling stacked_gt_points = np.concatenate(stacked_gt_points, axis=0) np.save(db_data_save_path, stacked_gt_points) def create_gt_database_of_single_scene(self, info_with_idx, database_save_path=None, use_sequence_data=False, used_classes=None, total_samples=0, use_cuda=False, crop_gt_with_tail=False): info, info_idx = info_with_idx print('gt_database sample: %d/%d' % (info_idx, total_samples)) all_db_infos = {} pc_info = info['point_cloud'] sequence_name = pc_info['lidar_sequence'] sample_idx = pc_info['sample_idx'] points = self.get_lidar(sequence_name, sample_idx) if use_sequence_data: points, num_points_all, sample_idx_pre_list, _, _, _, _ = self.get_sequence_data( info, points, sequence_name, sample_idx, self.dataset_cfg.SEQUENCE_CONFIG ) annos = info['annos'] names = annos['name'] difficulty = annos['difficulty'] gt_boxes = annos['gt_boxes_lidar'] if info_idx % 4 != 0 and len(names) > 0: mask = (names == 'Vehicle') names = names[~mask] difficulty = difficulty[~mask] gt_boxes = gt_boxes[~mask] if info_idx % 2 != 0 and len(names) > 0: mask = (names == 'Pedestrian') names = names[~mask] difficulty = difficulty[~mask] gt_boxes = gt_boxes[~mask] num_obj = gt_boxes.shape[0] if num_obj == 0: return {} if use_sequence_data and crop_gt_with_tail: assert gt_boxes.shape[1] == 9 speed = gt_boxes[:, 7:9] sequence_cfg = self.dataset_cfg.SEQUENCE_CONFIG assert sequence_cfg.SAMPLE_OFFSET[1] == 0 assert sequence_cfg.SAMPLE_OFFSET[0] < 0 num_frames = sequence_cfg.SAMPLE_OFFSET[1] - sequence_cfg.SAMPLE_OFFSET[0] + 1 assert num_frames > 1 latest_center = gt_boxes[:, 0:2] oldest_center = latest_center - speed * (num_frames - 1) * 0.1 new_center = (latest_center + oldest_center) * 0.5 new_length = gt_boxes[:, 3] + np.linalg.norm(latest_center - oldest_center, axis=-1) gt_boxes_crop = gt_boxes.copy() gt_boxes_crop[:, 0:2] = new_center gt_boxes_crop[:, 3] = new_length else: gt_boxes_crop = gt_boxes if use_cuda: box_idxs_of_pts = roiaware_pool3d_utils.points_in_boxes_gpu( torch.from_numpy(points[:, 0:3]).unsqueeze(dim=0).float().cuda(), torch.from_numpy(gt_boxes_crop[:, 0:7]).unsqueeze(dim=0).float().cuda() ).long().squeeze(dim=0).cpu().numpy() else: box_point_mask = roiaware_pool3d_utils.points_in_boxes_cpu( torch.from_numpy(points[:, 0:3]).float(), torch.from_numpy(gt_boxes_crop[:, 0:7]).float() ).long().numpy() # (num_boxes, num_points) for i in range(num_obj): filename = '%s_%04d_%s_%d.bin' % (sequence_name, sample_idx, names[i], i) filepath = database_save_path / filename if use_cuda: gt_points = points[box_idxs_of_pts == i] else: gt_points = points[box_point_mask[i] > 0] gt_points[:, :3] -= gt_boxes[i, :3] if (used_classes is None) or names[i] in used_classes: gt_points = gt_points.astype(np.float32) assert gt_points.dtype == np.float32 with open(filepath, 'w') as f: gt_points.tofile(f) db_path = str(filepath.relative_to(self.root_path)) # gt_database/xxxxx.bin db_info = {'name': names[i], 'path': db_path, 'sequence_name': sequence_name, 'sample_idx': sample_idx, 'gt_idx': i, 'box3d_lidar': gt_boxes[i], 'num_points_in_gt': gt_points.shape[0], 'difficulty': difficulty[i], 'box3d_crop': gt_boxes_crop[i]} if names[i] in all_db_infos: all_db_infos[names[i]].append(db_info) else: all_db_infos[names[i]] = [db_info] return all_db_infos def create_groundtruth_database_parallel(self, info_path, save_path, used_classes=None, split='train', sampled_interval=10, processed_data_tag=None, num_workers=16, crop_gt_with_tail=False): use_sequence_data = self.dataset_cfg.get('SEQUENCE_CONFIG', None) is not None and self.dataset_cfg.SEQUENCE_CONFIG.ENABLED if use_sequence_data: st_frame, ed_frame = self.dataset_cfg.SEQUENCE_CONFIG.SAMPLE_OFFSET[0], self.dataset_cfg.SEQUENCE_CONFIG.SAMPLE_OFFSET[1] self.dataset_cfg.SEQUENCE_CONFIG.SAMPLE_OFFSET[0] = min(-4, st_frame) # at least we use 5 frames for generating gt database to support various sequence configs (<= 5 frames) st_frame = self.dataset_cfg.SEQUENCE_CONFIG.SAMPLE_OFFSET[0] database_save_path = save_path / ('%s_gt_database_%s_sampled_%d_multiframe_%s_to_%s_%sparallel' % (processed_data_tag, split, sampled_interval, st_frame, ed_frame, 'tail_' if crop_gt_with_tail else '')) db_info_save_path = save_path / ('%s_waymo_dbinfos_%s_sampled_%d_multiframe_%s_to_%s_%sparallel.pkl' % (processed_data_tag, split, sampled_interval, st_frame, ed_frame, 'tail_' if crop_gt_with_tail else '')) else: database_save_path = save_path / ('%s_gt_database_%s_sampled_%d_parallel' % (processed_data_tag, split, sampled_interval)) db_info_save_path = save_path / ('%s_waymo_dbinfos_%s_sampled_%d_parallel.pkl' % (processed_data_tag, split, sampled_interval)) database_save_path.mkdir(parents=True, exist_ok=True) with open(info_path, 'rb') as f: infos = pickle.load(f) print(f'Number workers: {num_workers}') create_gt_database_of_single_scene = partial( self.create_gt_database_of_single_scene, use_sequence_data=use_sequence_data, database_save_path=database_save_path, used_classes=used_classes, total_samples=len(infos), use_cuda=False, crop_gt_with_tail=crop_gt_with_tail ) # create_gt_database_of_single_scene((infos[300], 0)) with multiprocessing.Pool(num_workers) as p: all_db_infos_list = list(p.map(create_gt_database_of_single_scene, zip(infos, np.arange(len(infos))))) all_db_infos = {} for cur_db_infos in all_db_infos_list: for key, val in cur_db_infos.items(): if key not in all_db_infos: all_db_infos[key] = val else: all_db_infos[key].extend(val) for k, v in all_db_infos.items(): print('Database %s: %d' % (k, len(v))) with open(db_info_save_path, 'wb') as f: pickle.dump(all_db_infos, f) def create_waymo_infos(dataset_cfg, class_names, data_path, save_path, raw_data_tag='raw_data', processed_data_tag='waymo_processed_data', workers=min(16, multiprocessing.cpu_count()), update_info_only=False): dataset = WaymoDataset( dataset_cfg=dataset_cfg, class_names=class_names, root_path=data_path, training=False, logger=common_utils.create_logger() ) train_split, val_split = 'train', 'val' train_filename = save_path / ('%s_infos_%s.pkl' % (processed_data_tag, train_split)) val_filename = save_path / ('%s_infos_%s.pkl' % (processed_data_tag, val_split)) os.environ["CUDA_VISIBLE_DEVICES"] = "-1" print('---------------Start to generate data infos---------------') dataset.set_split(train_split) waymo_infos_train = dataset.get_infos( raw_data_path=data_path / raw_data_tag, save_path=save_path / processed_data_tag, num_workers=workers, has_label=True, sampled_interval=1, update_info_only=update_info_only ) with open(train_filename, 'wb') as f: pickle.dump(waymo_infos_train, f) print('----------------Waymo info train file is saved to %s----------------' % train_filename) dataset.set_split(val_split) waymo_infos_val = dataset.get_infos( raw_data_path=data_path / raw_data_tag, save_path=save_path / processed_data_tag, num_workers=workers, has_label=True, sampled_interval=1, update_info_only=update_info_only ) with open(val_filename, 'wb') as f: pickle.dump(waymo_infos_val, f) print('----------------Waymo info val file is saved to %s----------------' % val_filename) if update_info_only: return print('---------------Start create groundtruth database for data augmentation---------------') os.environ["CUDA_VISIBLE_DEVICES"] = "0" dataset.set_split(train_split) dataset.create_groundtruth_database( info_path=train_filename, save_path=save_path, split='train', sampled_interval=1, used_classes=['Vehicle', 'Pedestrian', 'Cyclist'], processed_data_tag=processed_data_tag ) print('---------------Data preparation Done---------------') def create_waymo_gt_database( dataset_cfg, class_names, data_path, save_path, processed_data_tag='waymo_processed_data', workers=min(16, multiprocessing.cpu_count()), use_parallel=False, crop_gt_with_tail=False): dataset = WaymoDataset( dataset_cfg=dataset_cfg, class_names=class_names, root_path=data_path, training=False, logger=common_utils.create_logger() ) train_split = 'train' train_filename = save_path / ('%s_infos_%s.pkl' % (processed_data_tag, train_split)) print('---------------Start create groundtruth database for data augmentation---------------') dataset.set_split(train_split) if use_parallel: dataset.create_groundtruth_database_parallel( info_path=train_filename, save_path=save_path, split='train', sampled_interval=1, used_classes=['Vehicle', 'Pedestrian', 'Cyclist'], processed_data_tag=processed_data_tag, num_workers=workers, crop_gt_with_tail=crop_gt_with_tail ) else: dataset.create_groundtruth_database( info_path=train_filename, save_path=save_path, split='train', sampled_interval=1, used_classes=['Vehicle', 'Pedestrian', 'Cyclist'], processed_data_tag=processed_data_tag ) print('---------------Data preparation Done---------------') if __name__ == '__main__': import argparse import yaml from easydict import EasyDict parser = argparse.ArgumentParser(description='arg parser') parser.add_argument('--cfg_file', type=str, default=None, help='specify the config of dataset') parser.add_argument('--func', type=str, default='create_waymo_infos', help='') parser.add_argument('--processed_data_tag', type=str, default='waymo_processed_data_v0_5_0', help='') parser.add_argument('--update_info_only', action='store_true', default=False, help='') parser.add_argument('--use_parallel', action='store_true', default=False, help='') parser.add_argument('--wo_crop_gt_with_tail', action='store_true', default=False, help='') args = parser.parse_args() ROOT_DIR = (Path(__file__).resolve().parent / '../../../').resolve() if args.func == 'create_waymo_infos': try: yaml_config = yaml.safe_load(open(args.cfg_file), Loader=yaml.FullLoader) except: yaml_config = yaml.safe_load(open(args.cfg_file)) dataset_cfg = EasyDict(yaml_config) dataset_cfg.PROCESSED_DATA_TAG = args.processed_data_tag create_waymo_infos( dataset_cfg=dataset_cfg, class_names=['Vehicle', 'Pedestrian', 'Cyclist'], data_path=ROOT_DIR / 'data' / 'waymo', save_path=ROOT_DIR / 'data' / 'waymo', raw_data_tag='raw_data', processed_data_tag=args.processed_data_tag, update_info_only=args.update_info_only ) elif args.func == 'create_waymo_gt_database': try: yaml_config = yaml.safe_load(open(args.cfg_file), Loader=yaml.FullLoader) except: yaml_config = yaml.safe_load(open(args.cfg_file)) dataset_cfg = EasyDict(yaml_config) dataset_cfg.PROCESSED_DATA_TAG = args.processed_data_tag create_waymo_gt_database( dataset_cfg=dataset_cfg, class_names=['Vehicle', 'Pedestrian', 'Cyclist'], data_path=ROOT_DIR / 'data' / 'waymo', save_path=ROOT_DIR / 'data' / 'waymo', processed_data_tag=args.processed_data_tag, use_parallel=args.use_parallel, crop_gt_with_tail=not args.wo_crop_gt_with_tail ) else: raise NotImplementedError