def prepare_test_img(self, idx): """Prepare an image for testing (multi-scale and flipping)""" sample_id = self.sample_ids[idx] # load image img = mmcv.imread(osp.join(self.img_prefix, '%06d.png' % sample_id)) img, img_shape, pad_shape, scale_factor = self.img_transform( img, 1, False) calib = Calibration(osp.join(self.calib_prefix, '%06d.txt' % sample_id)) if self.with_label: objects = read_label( osp.join(self.label_prefix, '%06d.txt' % sample_id)) gt_bboxes = [ object.box3d for object in objects if object.type in self.class_name ] if len(gt_bboxes) != 0: gt_bboxes = np.array(gt_bboxes, dtype=np.float32) gt_labels = np.ones(len(gt_bboxes), dtype=np.int64) # transfer from cam to lidar coordinates gt_bboxes[:, :3] = project_rect_to_velo( gt_bboxes[:, :3], calib) else: gt_bboxes = None gt_labels = None img_meta = dict(img_shape=img_shape, sample_idx=sample_id, calib=calib) data = dict(img=to_tensor(img), img_meta=DC(img_meta, cpu_only=True)) if self.anchors is not None: data['anchors'] = DC(to_tensor(self.anchors.astype(np.float32))) if self.with_mask: NotImplemented if self.with_point: points = read_lidar( osp.join(self.lidar_prefix, '%06d.bin' % sample_id)) if isinstance(self.generator, VoxelGenerator): #voxels, coordinates, num_points = self.generator.generate(points) voxel_size = self.generator.voxel_size pc_range = self.generator.point_cloud_range grid_size = self.generator.grid_size keep = points_op_cpu.points_bound_kernel(points, pc_range[:3], pc_range[3:]) voxels = points[keep, :] coordinates = ( (voxels[:, [2, 1, 0]] - np.array(pc_range[[2, 1, 0]], dtype=np.float32)) / np.array(voxel_size[::-1], dtype=np.float32)).astype(np.int32) num_points = np.ones(len(keep)).astype(np.int32) data['voxels'] = DC(to_tensor(voxels.astype(np.float32))) data['coordinates'] = DC(to_tensor(coordinates)) data['num_points'] = DC(to_tensor(num_points)) if self.anchor_area_threshold >= 0 and self.anchors is not None: dense_voxel_map = sparse_sum_for_anchors_mask( coordinates, tuple(grid_size[::-1][1:])) dense_voxel_map = dense_voxel_map.cumsum(0) dense_voxel_map = dense_voxel_map.cumsum(1) anchors_area = fused_get_anchors_area(dense_voxel_map, self.anchors_bv, voxel_size, pc_range, grid_size) anchors_mask = anchors_area > self.anchor_area_threshold data['anchors_mask'] = DC( to_tensor(anchors_mask.astype(np.uint8))) if self.with_label: data['gt_labels'] = DC(to_tensor(gt_labels), cpu_only=True) data['gt_bboxes'] = DC(to_tensor(gt_bboxes), cpu_only=True) else: data['gt_labels'] = DC(None, cpu_only=True) data['gt_bboxes'] = DC(None, cpu_only=True) return data
def prepare_train_img(self, idx): sample_id = self.sample_ids[idx] # load image img = mmcv.imread(osp.join(self.img_prefix, '%06d.png' % sample_id)) img, img_shape, pad_shape, scale_factor = self.img_transform( img, 1, False) objects = read_label( osp.join(self.label_prefix, '%06d.txt' % sample_id)) calib = Calibration(osp.join(self.calib_prefix, '%06d.txt' % sample_id)) gt_bboxes = [ object.box3d for object in objects if object.type not in ["DontCare"] ] gt_bboxes = np.array(gt_bboxes, dtype=np.float32) gt_types = [ object.type for object in objects if object.type not in ["DontCare"] ] # transfer from cam to lidar coordinates if len(gt_bboxes) != 0: gt_bboxes[:, :3] = project_rect_to_velo(gt_bboxes[:, :3], calib) img_meta = dict(img_shape=img_shape, sample_idx=sample_id, calib=calib) data = dict(img=to_tensor(img), img_meta=DC(img_meta, cpu_only=True)) if self.anchors is not None: data['anchors'] = DC(to_tensor(self.anchors.astype(np.float32))) if self.with_mask: NotImplemented if self.with_point: points = read_lidar( osp.join(self.lidar_prefix, '%06d.bin' % sample_id)) if self.augmentor is not None and self.test_mode is False: sampled_gt_boxes, sampled_gt_types, sampled_points = self.augmentor.sample_all( gt_bboxes, gt_types) assert sampled_points.dtype == np.float32 gt_bboxes = np.concatenate([gt_bboxes, sampled_gt_boxes]) gt_types = gt_types + sampled_gt_types assert len(gt_types) == len(gt_bboxes) # to avoid overlapping point (option) masks = points_in_rbbox(points, sampled_gt_boxes) #masks = points_op_cpu.points_in_bbox3d_np(points[:,:3], sampled_gt_boxes) points = points[np.logical_not(masks.any(-1))] # paste sampled points to the scene points = np.concatenate([sampled_points, points], axis=0) # select the interest classes selected = [ i for i in range(len(gt_types)) if gt_types[i] in self.class_names ] gt_bboxes = gt_bboxes[selected, :] gt_types = [ gt_types[i] for i in range(len(gt_types)) if gt_types[i] in self.class_names ] # force van to have same label as car gt_types = ['Car' if n == 'Van' else n for n in gt_types] gt_labels = np.array( [self.class_names.index(n) + 1 for n in gt_types], dtype=np.int64) self.augmentor.noise_per_object_(gt_bboxes, points, num_try=100) gt_bboxes, points = self.augmentor.random_flip(gt_bboxes, points) gt_bboxes, points = self.augmentor.global_rotation( gt_bboxes, points) gt_bboxes, points = self.augmentor.global_scaling( gt_bboxes, points) if isinstance(self.generator, VoxelGenerator): #voxels, coordinates, num_points = self.generator.generate(points) voxel_size = self.generator.voxel_size pc_range = self.generator.point_cloud_range grid_size = self.generator.grid_size keep = points_op_cpu.points_bound_kernel(points, pc_range[:3], pc_range[3:]) voxels = points[keep, :] coordinates = ( (voxels[:, [2, 1, 0]] - np.array(pc_range[[2, 1, 0]], dtype=np.float32)) / np.array(voxel_size[::-1], dtype=np.float32)).astype(np.int32) num_points = np.ones(len(keep)).astype(np.int32) data['voxels'] = DC(to_tensor(voxels.astype(np.float32))) data['coordinates'] = DC(to_tensor(coordinates)) data['num_points'] = DC(to_tensor(num_points)) if self.anchor_area_threshold >= 0 and self.anchors is not None: dense_voxel_map = sparse_sum_for_anchors_mask( coordinates, tuple(grid_size[::-1][1:])) dense_voxel_map = dense_voxel_map.cumsum(0) dense_voxel_map = dense_voxel_map.cumsum(1) anchors_area = fused_get_anchors_area(dense_voxel_map, self.anchors_bv, voxel_size, pc_range, grid_size) anchors_mask = anchors_area > self.anchor_area_threshold data['anchors_mask'] = DC( to_tensor(anchors_mask.astype(np.uint8))) # filter gt_bbox out of range bv_range = self.generator.point_cloud_range[[0, 1, 3, 4]] mask = filter_gt_box_outside_range(gt_bboxes, bv_range) gt_bboxes = gt_bboxes[mask] gt_labels = gt_labels[mask] else: NotImplementedError # skip the image if there is no valid gt bbox if len(gt_bboxes) == 0: return None # limit rad to [-pi, pi] gt_bboxes[:, 6] = limit_period(gt_bboxes[:, 6], offset=0.5, period=2 * np.pi) if self.with_label: data['gt_labels'] = DC(to_tensor(gt_labels)) data['gt_bboxes'] = DC(to_tensor(gt_bboxes)) return data