import pickle import os import copy import numpy as np from skimage import io import torch import SharedArray import torch.distributed as dist from ...ops.iou3d_nms import iou3d_nms_utils from ...utils import box_utils, common_utils class DataBaseSampler(object): def __init__(self, root_path, sampler_cfg, class_names, logger=None): self.root_path = root_path self.class_names = class_names self.sampler_cfg = sampler_cfg self.img_aug_type = sampler_cfg.get('IMG_AUG_TYPE', None) self.img_aug_iou_thresh = sampler_cfg.get('IMG_AUG_IOU_THRESH', 0.5) self.logger = logger self.db_infos = {} for class_name in class_names: self.db_infos[class_name] = [] self.use_shared_memory = sampler_cfg.get('USE_SHARED_MEMORY', False) for db_info_path in sampler_cfg.DB_INFO_PATH: db_info_path = self.root_path.resolve() / db_info_path if not db_info_path.exists(): assert len(sampler_cfg.DB_INFO_PATH) == 1 sampler_cfg.DB_INFO_PATH[0] = sampler_cfg.BACKUP_DB_INFO['DB_INFO_PATH'] sampler_cfg.DB_DATA_PATH[0] = sampler_cfg.BACKUP_DB_INFO['DB_DATA_PATH'] db_info_path = self.root_path.resolve() / sampler_cfg.DB_INFO_PATH[0] sampler_cfg.NUM_POINT_FEATURES = sampler_cfg.BACKUP_DB_INFO['NUM_POINT_FEATURES'] with open(str(db_info_path), 'rb') as f: infos = pickle.load(f) [self.db_infos[cur_class].extend(infos[cur_class]) for cur_class in class_names] for func_name, val in sampler_cfg.PREPARE.items(): self.db_infos = getattr(self, func_name)(self.db_infos, val) self.gt_database_data_key = self.load_db_to_shared_memory() if self.use_shared_memory else None self.sample_groups = {} self.sample_class_num = {} self.limit_whole_scene = sampler_cfg.get('LIMIT_WHOLE_SCENE', False) for x in sampler_cfg.SAMPLE_GROUPS: class_name, sample_num = x.split(':') if class_name not in class_names: continue self.sample_class_num[class_name] = sample_num self.sample_groups[class_name] = { 'sample_num': sample_num, 'pointer': len(self.db_infos[class_name]), 'indices': np.arange(len(self.db_infos[class_name])) } def __getstate__(self): d = dict(self.__dict__) del d['logger'] return d def __setstate__(self, d): self.__dict__.update(d) def __del__(self): if self.use_shared_memory: self.logger.info('Deleting GT database from shared memory') cur_rank, num_gpus = common_utils.get_dist_info() sa_key = self.sampler_cfg.DB_DATA_PATH[0] if cur_rank % num_gpus == 0 and os.path.exists(f"/dev/shm/{sa_key}"): SharedArray.delete(f"shm://{sa_key}") if num_gpus > 1: dist.barrier() self.logger.info('GT database has been removed from shared memory') def load_db_to_shared_memory(self): self.logger.info('Loading GT database to shared memory') cur_rank, world_size, num_gpus = common_utils.get_dist_info(return_gpu_per_machine=True) assert self.sampler_cfg.DB_DATA_PATH.__len__() == 1, 'Current only support single DB_DATA' db_data_path = self.root_path.resolve() / self.sampler_cfg.DB_DATA_PATH[0] sa_key = self.sampler_cfg.DB_DATA_PATH[0] if cur_rank % num_gpus == 0 and not os.path.exists(f"/dev/shm/{sa_key}"): gt_database_data = np.load(db_data_path) common_utils.sa_create(f"shm://{sa_key}", gt_database_data) if num_gpus > 1: dist.barrier() self.logger.info('GT database has been saved to shared memory') return sa_key def filter_by_difficulty(self, db_infos, removed_difficulty): new_db_infos = {} for key, dinfos in db_infos.items(): pre_len = len(dinfos) new_db_infos[key] = [ info for info in dinfos if info['difficulty'] not in removed_difficulty ] if self.logger is not None: self.logger.info('Database filter by difficulty %s: %d => %d' % (key, pre_len, len(new_db_infos[key]))) return new_db_infos def filter_by_min_points(self, db_infos, min_gt_points_list): for name_num in min_gt_points_list: name, min_num = name_num.split(':') min_num = int(min_num) if min_num > 0 and name in db_infos.keys(): filtered_infos = [] for info in db_infos[name]: if info['num_points_in_gt'] >= min_num: filtered_infos.append(info) if self.logger is not None: self.logger.info('Database filter by min points %s: %d => %d' % (name, len(db_infos[name]), len(filtered_infos))) db_infos[name] = filtered_infos return db_infos def sample_with_fixed_number(self, class_name, sample_group): """ Args: class_name: sample_group: Returns: """ sample_num, pointer, indices = int(sample_group['sample_num']), sample_group['pointer'], sample_group['indices'] if pointer >= len(self.db_infos[class_name]): indices = np.random.permutation(len(self.db_infos[class_name])) pointer = 0 sampled_dict = [self.db_infos[class_name][idx] for idx in indices[pointer: pointer + sample_num]] pointer += sample_num sample_group['pointer'] = pointer sample_group['indices'] = indices return sampled_dict @staticmethod def put_boxes_on_road_planes(gt_boxes, road_planes, calib): """ Only validate in KITTIDataset Args: gt_boxes: (N, 7 + C) [x, y, z, dx, dy, dz, heading, ...] road_planes: [a, b, c, d] calib: Returns: """ a, b, c, d = road_planes center_cam = calib.lidar_to_rect(gt_boxes[:, 0:3]) cur_height_cam = (-d - a * center_cam[:, 0] - c * center_cam[:, 2]) / b center_cam[:, 1] = cur_height_cam cur_lidar_height = calib.rect_to_lidar(center_cam)[:, 2] mv_height = gt_boxes[:, 2] - gt_boxes[:, 5] / 2 - cur_lidar_height gt_boxes[:, 2] -= mv_height # lidar view return gt_boxes, mv_height def copy_paste_to_image_kitti(self, data_dict, crop_feat, gt_number, point_idxes=None): kitti_img_aug_type = 'by_depth' kitti_img_aug_use_type = 'annotation' image = data_dict['images'] boxes3d = data_dict['gt_boxes'] boxes2d = data_dict['gt_boxes2d'] corners_lidar = box_utils.boxes_to_corners_3d(boxes3d) if 'depth' in kitti_img_aug_type: paste_order = boxes3d[:,0].argsort() paste_order = paste_order[::-1] else: paste_order = np.arange(len(boxes3d),dtype=np.int) if 'reverse' in kitti_img_aug_type: paste_order = paste_order[::-1] paste_mask = -255 * np.ones(image.shape[:2], dtype=np.int) fg_mask = np.zeros(image.shape[:2], dtype=np.int) overlap_mask = np.zeros(image.shape[:2], dtype=np.int) depth_mask = np.zeros((*image.shape[:2], 2), dtype=np.float) points_2d, depth_2d = data_dict['calib'].lidar_to_img(data_dict['points'][:,:3]) points_2d[:,0] = np.clip(points_2d[:,0], a_min=0, a_max=image.shape[1]-1) points_2d[:,1] = np.clip(points_2d[:,1], a_min=0, a_max=image.shape[0]-1) points_2d = points_2d.astype(np.int) for _order in paste_order: _box2d = boxes2d[_order] image[_box2d[1]:_box2d[3],_box2d[0]:_box2d[2]] = crop_feat[_order] overlap_mask[_box2d[1]:_box2d[3],_box2d[0]:_box2d[2]] += \ (paste_mask[_box2d[1]:_box2d[3],_box2d[0]:_box2d[2]] > 0).astype(np.int) paste_mask[_box2d[1]:_box2d[3],_box2d[0]:_box2d[2]] = _order if 'cover' in kitti_img_aug_use_type: # HxWx2 for min and max depth of each box region depth_mask[_box2d[1]:_box2d[3],_box2d[0]:_box2d[2],0] = corners_lidar[_order,:,0].min() depth_mask[_box2d[1]:_box2d[3],_box2d[0]:_box2d[2],1] = corners_lidar[_order,:,0].max() # foreground area of original point cloud in image plane if _order < gt_number: fg_mask[_box2d[1]:_box2d[3],_box2d[0]:_box2d[2]] = 1 data_dict['images'] = image # if not self.joint_sample: # return data_dict new_mask = paste_mask[points_2d[:,1], points_2d[:,0]]==(point_idxes+gt_number) if False: # self.keep_raw: raw_mask = (point_idxes == -1) else: raw_fg = (fg_mask == 1) & (paste_mask >= 0) & (paste_mask < gt_number) raw_bg = (fg_mask == 0) & (paste_mask < 0) raw_mask = raw_fg[points_2d[:,1], points_2d[:,0]] | raw_bg[points_2d[:,1], points_2d[:,0]] keep_mask = new_mask | raw_mask data_dict['points_2d'] = points_2d if 'annotation' in kitti_img_aug_use_type: data_dict['points'] = data_dict['points'][keep_mask] data_dict['points_2d'] = data_dict['points_2d'][keep_mask] elif 'projection' in kitti_img_aug_use_type: overlap_mask[overlap_mask>=1] = 1 data_dict['overlap_mask'] = overlap_mask if 'cover' in kitti_img_aug_use_type: data_dict['depth_mask'] = depth_mask return data_dict def sample_gt_boxes_2d(self, data_dict, sampled_boxes, valid_mask): mv_height = None if self.img_aug_type == 'kitti': sampled_boxes2d, mv_height, ret_valid_mask = self.sample_gt_boxes_2d_kitti(data_dict, sampled_boxes, valid_mask) else: raise NotImplementedError return sampled_boxes2d, mv_height, ret_valid_mask def initilize_image_aug_dict(self, data_dict, gt_boxes_mask): img_aug_gt_dict = None if self.img_aug_type is None: pass elif self.img_aug_type == 'kitti': obj_index_list, crop_boxes2d = [], [] gt_number = gt_boxes_mask.sum().astype(np.int) gt_boxes2d = data_dict['gt_boxes2d'][gt_boxes_mask].astype(np.int) gt_crops2d = [data_dict['images'][_x[1]:_x[3],_x[0]:_x[2]] for _x in gt_boxes2d] img_aug_gt_dict = { 'obj_index_list': obj_index_list, 'gt_crops2d': gt_crops2d, 'gt_boxes2d': gt_boxes2d, 'gt_number': gt_number, 'crop_boxes2d': crop_boxes2d } else: raise NotImplementedError return img_aug_gt_dict def collect_image_crops(self, img_aug_gt_dict, info, data_dict, obj_points, sampled_gt_boxes, sampled_gt_boxes2d, idx): if self.img_aug_type == 'kitti': new_box, img_crop2d, obj_points, obj_idx = self.collect_image_crops_kitti(info, data_dict, obj_points, sampled_gt_boxes, sampled_gt_boxes2d, idx) img_aug_gt_dict['crop_boxes2d'].append(new_box) img_aug_gt_dict['gt_crops2d'].append(img_crop2d) img_aug_gt_dict['obj_index_list'].append(obj_idx) else: raise NotImplementedError return img_aug_gt_dict, obj_points def copy_paste_to_image(self, img_aug_gt_dict, data_dict, points): if self.img_aug_type == 'kitti': obj_points_idx = np.concatenate(img_aug_gt_dict['obj_index_list'], axis=0) point_idxes = -1 * np.ones(len(points), dtype=np.int) point_idxes[:obj_points_idx.shape[0]] = obj_points_idx data_dict['gt_boxes2d'] = np.concatenate([img_aug_gt_dict['gt_boxes2d'], np.array(img_aug_gt_dict['crop_boxes2d'])], axis=0) data_dict = self.copy_paste_to_image_kitti(data_dict, img_aug_gt_dict['gt_crops2d'], img_aug_gt_dict['gt_number'], point_idxes) if 'road_plane' in data_dict: data_dict.pop('road_plane') else: raise NotImplementedError return data_dict def add_sampled_boxes_to_scene(self, data_dict, sampled_gt_boxes, total_valid_sampled_dict, mv_height=None, sampled_gt_boxes2d=None): gt_boxes_mask = data_dict['gt_boxes_mask'] gt_boxes = data_dict['gt_boxes'][gt_boxes_mask] gt_names = data_dict['gt_names'][gt_boxes_mask] points = data_dict['points'] if self.sampler_cfg.get('USE_ROAD_PLANE', False) and mv_height is None: sampled_gt_boxes, mv_height = self.put_boxes_on_road_planes( sampled_gt_boxes, data_dict['road_plane'], data_dict['calib'] ) data_dict.pop('calib') data_dict.pop('road_plane') obj_points_list = [] # convert sampled 3D boxes to image plane img_aug_gt_dict = self.initilize_image_aug_dict(data_dict, gt_boxes_mask) if self.use_shared_memory: gt_database_data = SharedArray.attach(f"shm://{self.gt_database_data_key}") gt_database_data.setflags(write=0) else: gt_database_data = None for idx, info in enumerate(total_valid_sampled_dict): if self.use_shared_memory: start_offset, end_offset = info['global_data_offset'] obj_points = copy.deepcopy(gt_database_data[start_offset:end_offset]) else: file_path = self.root_path / info['path'] obj_points = np.fromfile(str(file_path), dtype=np.float32).reshape( [-1, self.sampler_cfg.NUM_POINT_FEATURES]) if obj_points.shape[0] != info['num_points_in_gt']: obj_points = np.fromfile(str(file_path), dtype=np.float64).reshape(-1, self.sampler_cfg.NUM_POINT_FEATURES) assert obj_points.shape[0] == info['num_points_in_gt'] obj_points[:, :3] += info['box3d_lidar'][:3].astype(np.float32) if self.sampler_cfg.get('USE_ROAD_PLANE', False): # mv height obj_points[:, 2] -= mv_height[idx] if self.img_aug_type is not None: img_aug_gt_dict, obj_points = self.collect_image_crops( img_aug_gt_dict, info, data_dict, obj_points, sampled_gt_boxes, sampled_gt_boxes2d, idx ) obj_points_list.append(obj_points) obj_points = np.concatenate(obj_points_list, axis=0) sampled_gt_names = np.array([x['name'] for x in total_valid_sampled_dict]) if self.sampler_cfg.get('FILTER_OBJ_POINTS_BY_TIMESTAMP', False) or obj_points.shape[-1] != points.shape[-1]: if self.sampler_cfg.get('FILTER_OBJ_POINTS_BY_TIMESTAMP', False): min_time = min(self.sampler_cfg.TIME_RANGE[0], self.sampler_cfg.TIME_RANGE[1]) max_time = max(self.sampler_cfg.TIME_RANGE[0], self.sampler_cfg.TIME_RANGE[1]) else: assert obj_points.shape[-1] == points.shape[-1] + 1 # transform multi-frame GT points to single-frame GT points min_time = max_time = 0.0 time_mask = np.logical_and(obj_points[:, -1] < max_time + 1e-6, obj_points[:, -1] > min_time - 1e-6) obj_points = obj_points[time_mask] large_sampled_gt_boxes = box_utils.enlarge_box3d( sampled_gt_boxes[:, 0:7], extra_width=self.sampler_cfg.REMOVE_EXTRA_WIDTH ) points = box_utils.remove_points_in_boxes3d(points, large_sampled_gt_boxes) points = np.concatenate([obj_points[:, :points.shape[-1]], points], axis=0) gt_names = np.concatenate([gt_names, sampled_gt_names], axis=0) gt_boxes = np.concatenate([gt_boxes, sampled_gt_boxes], axis=0) data_dict['gt_boxes'] = gt_boxes data_dict['gt_names'] = gt_names data_dict['points'] = points if self.img_aug_type is not None: data_dict = self.copy_paste_to_image(img_aug_gt_dict, data_dict, points) return data_dict def __call__(self, data_dict): """ Args: data_dict: gt_boxes: (N, 7 + C) [x, y, z, dx, dy, dz, heading, ...] Returns: """ gt_boxes = data_dict['gt_boxes'] gt_names = data_dict['gt_names'].astype(str) existed_boxes = gt_boxes total_valid_sampled_dict = [] sampled_mv_height = [] sampled_gt_boxes2d = [] for class_name, sample_group in self.sample_groups.items(): if self.limit_whole_scene: num_gt = np.sum(class_name == gt_names) sample_group['sample_num'] = str(int(self.sample_class_num[class_name]) - num_gt) if int(sample_group['sample_num']) > 0: sampled_dict = self.sample_with_fixed_number(class_name, sample_group) sampled_boxes = np.stack([x['box3d_lidar'] for x in sampled_dict], axis=0).astype(np.float32) assert not self.sampler_cfg.get('DATABASE_WITH_FAKELIDAR', False), 'Please use latest codes to generate GT_DATABASE' iou1 = iou3d_nms_utils.boxes_bev_iou_cpu(sampled_boxes[:, 0:7], existed_boxes[:, 0:7]) iou2 = iou3d_nms_utils.boxes_bev_iou_cpu(sampled_boxes[:, 0:7], sampled_boxes[:, 0:7]) iou2[range(sampled_boxes.shape[0]), range(sampled_boxes.shape[0])] = 0 iou1 = iou1 if iou1.shape[1] > 0 else iou2 valid_mask = ((iou1.max(axis=1) + iou2.max(axis=1)) == 0) if self.img_aug_type is not None: sampled_boxes2d, mv_height, valid_mask = self.sample_gt_boxes_2d(data_dict, sampled_boxes, valid_mask) sampled_gt_boxes2d.append(sampled_boxes2d) if mv_height is not None: sampled_mv_height.append(mv_height) valid_mask = valid_mask.nonzero()[0] valid_sampled_dict = [sampled_dict[x] for x in valid_mask] valid_sampled_boxes = sampled_boxes[valid_mask] existed_boxes = np.concatenate((existed_boxes, valid_sampled_boxes[:, :existed_boxes.shape[-1]]), axis=0) total_valid_sampled_dict.extend(valid_sampled_dict) sampled_gt_boxes = existed_boxes[gt_boxes.shape[0]:, :] if total_valid_sampled_dict.__len__() > 0: sampled_gt_boxes2d = np.concatenate(sampled_gt_boxes2d, axis=0) if len(sampled_gt_boxes2d) > 0 else None sampled_mv_height = np.concatenate(sampled_mv_height, axis=0) if len(sampled_mv_height) > 0 else None data_dict = self.add_sampled_boxes_to_scene( data_dict, sampled_gt_boxes, total_valid_sampled_dict, sampled_mv_height, sampled_gt_boxes2d ) data_dict.pop('gt_boxes_mask') return data_dict