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']
