class LyftLevel5Dataset(data.Dataset): def __init__(self, phase): df_name = phase if phase is not 'test' else 'sample_submission' dataset = 'train' if phase is not 'test' else 'test' self.phase = phase self.df = pd.read_csv( os.path.join(cfg.work_dir, 'data', df_name + '.csv')) self.lyft_dataset = LyftDataset( data_path=os.path.join(cfg.input_dir, dataset), json_path=os.path.join(cfg.input_dir, dataset, 'data'), verbose=False) def voxelize_pointclouds(self, pointclouds): # Shuffle the pointclouds. np.random.shuffle(pointclouds) # x, y and z correspond to vox_width, vox_height and vox_depth, respectively. voxel_coords = ( (pointclouds[:, :3] - np.array([cfg.xrange[0], cfg.yrange[0], cfg.zrange[0]])) / (cfg.vox_width, cfg.vox_height, cfg.vox_depth)).astype(np.int32) # Convert to (z, y, x) to sample unique voxel. voxel_coords = voxel_coords[:, [2, 1, 0]] voxel_coords, inv_ind, voxel_counts = np.unique(voxel_coords, axis=0, return_inverse=True, return_counts=True) # Fill each voxel with voxel feature. voxel_features = [] for i in range(len(voxel_coords)): voxel = np.zeros((cfg.pointclouds_per_vox, 7), dtype=np.float32) pts = pointclouds[inv_ind == i] if voxel_counts[i] > cfg.pointclouds_per_vox: pts = pts[:cfg.pointclouds_per_vox, :] voxel_counts[i] = cfg.pointclouds_per_vox # Augment the points: (x, y, z) -> (x, y, z, i, x-mean, y-mean, z-mean). voxel[:pts.shape[0], :] = np.concatenate( (pts, pts[:, :3] - np.mean(pts[:, :3], 0)), axis=1) voxel_features.append(voxel) return np.array(voxel_features), voxel_coords def cal_target(self, gt_boxes3d): anchors_d = np.sqrt(cfg.anchors[:, 3]**2 + cfg.anchors[:, 4]**2) pos_equal_one = np.zeros((*cfg.feature_map_shape, cfg.ac_rot_z)) neg_equal_one = np.zeros((*cfg.feature_map_shape, cfg.ac_rot_z)) targets = np.zeros((*cfg.feature_map_shape, 7 * cfg.ac_rot_z)) gt_boxes3d_xyzlwhr = np.array([[ gt_box3d.center[0], gt_box3d.center[1], gt_box3d.center[2], gt_box3d.wlh[1], gt_box3d.wlh[0], gt_box3d.wlh[2], gt_box3d.orientation.yaw_pitch_roll[0] ] for gt_box3d in gt_boxes3d]) # Some class is not included in dataset. try: # Only taken into account rotation around yaw axis. # It means bottom-corners equal to top-corner. gt_boxes2d_corners = utils.boxes3d_to_corners(gt_boxes3d_xyzlwhr) except Exception as e: return pos_equal_one, neg_equal_one, targets, False ac_boxes2d_corners = utils.boxes3d_to_corners(cfg.anchors) ac_boxes2d = utils.boxes2d_four_corners_to_two_corners( ac_boxes2d_corners) gt_boxes2d = utils.boxes2d_four_corners_to_two_corners( gt_boxes2d_corners) # iou: [num_ac_boxes2d, num_gt_boxes2d] iou = bbox_overlaps( np.ascontiguousarray(ac_boxes2d).astype(np.float32), np.ascontiguousarray(gt_boxes2d).astype(np.float32)) # id_highest: [num_gt_boxes2d] # - For each gt_boxes2d, index of max iou-scored anchor box. id_highest = np.argmax(iou.T, axis=1) # id_highest_gt: [num_gt_boxes2d] # - Ranged from 0 to num_gt_boxes2d - 1. id_highest_gt = np.arange(iou.T.shape[0]) # For gt_boxes3d, mask stands for filter of box with highest anchor which has iou > 0. mask = iou.T[id_highest_gt, id_highest] > 0 # Get rid of those gt_boxes3d with iou of 0 with each anchor. id_highest, id_highest_gt = id_highest[mask], id_highest_gt[mask] # In the table of iou, every (ac_boxes2d and gt_boxes2d) pair with iou > iou_threshold_p. # id_pos: [num_pos_ac_boxes2d] # id_pos_gt: [num_pos_ac_boxes2d] id_pos, id_pos_gt = np.where(iou > cfg.iou_pos_threshold) # In the table of iou, every anchor with iou < iou_threshold_n with all gt_boxes2d. id_neg = np.where( np.sum(iou < cfg.iou_neg_threshold, axis=1) == iou.shape[1])[0] id_pos = np.concatenate([id_pos, id_highest]) id_pos_gt = np.concatenate([id_pos_gt, id_highest_gt]) id_pos, index = np.unique(id_pos, return_index=True) # The index of id_pos of anchor appears first time. id_pos_gt = id_pos_gt[index] id_neg.sort() # Cal the target and set the equal one. index_x, index_y, index_z = np.unravel_index( id_pos, (*cfg.feature_map_shape, cfg.ac_rot_z)) pos_equal_one[index_x, index_y, index_z] = 1 targets[index_x, index_y, np.array(index_z) * 7] = \ (gt_boxes3d_xyzlwhr[id_pos_gt, 0] - cfg.anchors[id_pos, 0]) / anchors_d[id_pos] targets[index_x, index_y, np.array(index_z) * 7 + 1] = \ (gt_boxes3d_xyzlwhr[id_pos_gt, 1] - cfg.anchors[id_pos, 1]) / anchors_d[id_pos] targets[index_x, index_y, np.array(index_z) * 7 + 2] = \ (gt_boxes3d_xyzlwhr[id_pos_gt, 2] - cfg.anchors[id_pos, 2]) / cfg.anchors[id_pos, 3] targets[index_x, index_y, np.array(index_z) * 7 + 3] = \ np.log(gt_boxes3d_xyzlwhr[id_pos_gt, 3] / cfg.anchors[id_pos, 3]) targets[index_x, index_y, np.array(index_z) * 7 + 4] = \ np.log(gt_boxes3d_xyzlwhr[id_pos_gt, 4] / cfg.anchors[id_pos, 4]) targets[index_x, index_y, np.array(index_z) * 7 + 5] = \ np.log(gt_boxes3d_xyzlwhr[id_pos_gt, 5] / cfg.anchors[id_pos, 5]) targets[index_x, index_y, np.array(index_z) * 7 + 6] = \ np.sin(gt_boxes3d_xyzlwhr[id_pos_gt, 6] - cfg.anchors[id_pos, 6]) index_x, index_y, index_z = np.unravel_index( id_neg, (*cfg.feature_map_shape, cfg.ac_rot_z)) neg_equal_one[index_x, index_y, index_z] = 1 # To avoid a box be pos/neg in the same time index_x, index_y, index_z = np.unravel_index( id_highest, (*cfg.feature_map_shape, cfg.ac_rot_z)) neg_equal_one[index_x, index_y, index_z] = 0 return pos_equal_one, neg_equal_one, targets, True def __getitem__(self, idx): # Point clouds of lidar are positioned at the sensor frame, # while gt_boxes3d at the global frame. sample_token = 'Id' if self.phase is 'test' else 'sample_token' sample = self.lyft_dataset.get('sample', self.df[sample_token][idx]) lidar_info = self.lyft_dataset.get('sample_data', sample['data']['LIDAR_TOP']) lidar_data_path = self.lyft_dataset.get_sample_data_path( sample['data']['LIDAR_TOP']) gt_boxes3d = self.lyft_dataset.get_boxes(sample['data']['LIDAR_TOP']) ego_pose = self.lyft_dataset.get('ego_pose', lidar_info['ego_pose_token']) calibrated_sensor = self.lyft_dataset.get( 'calibrated_sensor', lidar_info['calibrated_sensor_token']) # pointclouds w.r.t. the sensor frame: [4 xyzi, num_points] pointclouds = LidarPointCloud.from_file(lidar_data_path) # pointclouds: [4 xyzi, num_points] -> [num_points, 4 xyzi] pointclouds = pointclouds.points.transpose(1, 0) if self.phase is not 'test': # Convert gt_boxes3d from the global frame to the sensor frame. gt_boxes3d = utils.convert_boxes3d_from_global_to_sensor_frame( gt_boxes3d, ego_pose, calibrated_sensor) # Data augmentation. #pointclouds, gt_box3d = aug_data(pointclouds, gt_box3d) # Filter point clouds and gt_boxes3d within a specific range and class name. pointclouds, gt_boxes3d = utils.filter_pointclouds_gt_boxes3d( pointclouds, gt_boxes3d, cfg.class_name) # Voxelize point clouds. voxel_features, voxel_coords = self.voxelize_pointclouds( pointclouds) # Encode bounding boxes. pos_equal_one, neg_equal_one, targets, exception = self.cal_target( gt_boxes3d) return voxel_features, voxel_coords, pos_equal_one, neg_equal_one, targets, exception else: pointclouds = utils.filter_pointclouds_gt_boxes3d(pointclouds) # Voxelize point clouds. voxel_features, voxel_coords = self.voxelize_pointclouds( pointclouds) return voxel_features, voxel_coords, self.df[sample_token][ idx], ego_pose, calibrated_sensor def __len__(self): return len(self.df)
z_offset = -2.0 bev_shape = (336, 336, 3) bev = create_voxel_pointcloud(lidar_pointcloud.points, bev_shape, voxel_size=voxel_size, z_offset=z_offset) # So that the values in the voxels range from 0,1 we set a maximum intensity. bev = normalize_voxel_intensities(bev) plt.figure(figsize=(16, 8)) plt.imshow(bev) plt.show() boxes = level5data.get_boxes(sample_lidar_token) print(boxes[1].bottom_corners()) print(bev.shape[:3]) target_im = np.zeros(bev.shape[:3], dtype=np.uint8) def move_boxes_to_car_space(boxes, ego_pose): """ Move boxes from world space to car space. Note: mutates input boxes. """ translation = -np.array(ego_pose['translation']) rotation = Quaternion(ego_pose['rotation']).inverse for box in boxes:
'prev_token': prev_token } token = sample['next'] continue # get the lidar, ego and sensor objects for the token lidar_data = l5d.get('sample_data', lidar_token) ego_pose = l5d.get('ego_pose', lidar_data['ego_pose_token']) cal_sensor = l5d.get('calibrated_sensor', lidar_data['calibrated_sensor_token']) car_from_sensor = transform_matrix(cal_sensor['translation'], Quaternion( cal_sensor['rotation']), inverse=False) lidar_pointcloud.transform(car_from_sensor) lidar_points = lidar_pointcloud.points[:3, :] boxes = l5d.get_boxes(lidar_token) # collect the ground truth boxes canv_boxes = move_boxes_to_canvas_space(boxes, ego_pose, lidar_points) boxes_fp = osp.join(box_dir, token + '_boxes.pkl') pickle.dump(canv_boxes, open(boxes_fp, 'wb')) prev_token = sample['prev'] data_dict[token] = { 'lidar_fp': str(lidar_filepath), 'ego_pose': ego_pose, 'cal_sensor': cal_sensor, 'boxes': boxes_fp, 'prev_token': prev_token } token_list.append(token) token = sample['next']