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 show_lidar_with_boxes(pc_velo,
objects,
calib,
img_fov=False,
img_width=None,
img_height=None,
objects_pred=None):
''' Show all LiDAR points.
Draw 3d box in LiDAR point cloud (in velo coord system) '''
if 'mlab' not in sys.modules: import mayavi.mlab as mlab
# from viz_util import draw_lidar_simple, draw_lidar, draw_gt_boxes3d
print(('All point num: ', pc_velo.shape[0]))
fig = mlab.figure(figure=None,
bgcolor=(0, 0, 0),
fgcolor=None,
engine=None,
size=(1000, 500))
if img_fov:
pc_velo = get_lidar_in_image_fov(pc_velo, calib, 0, 0, img_width,
img_height)
print(('FOV point num: ', pc_velo.shape[0]))
draw_lidar(pc_velo, fig=fig)
color = (0, 1, 0)
for obj in objects:
if obj.type == 'DontCare': continue
# Draw 3d bounding box
box3d_pts_2d, box3d_pts_3d = utils.compute_box_3d(obj, calib.P2)
# wn update
box3d_pts_3d_velo = utils.project_rect_to_velo(box3d_pts_3d, calib)
# Draw heading arrow
ori3d_pts_2d, ori3d_pts_3d = utils.compute_orientation_3d(
obj, calib.P2)
# wn update
ori3d_pts_3d_velo = utils.project_rect_to_velo(ori3d_pts_3d, calib)
x1, y1, z1 = ori3d_pts_3d_velo[0, :]
x2, y2, z2 = ori3d_pts_3d_velo[1, :]
draw_gt_boxes3d([box3d_pts_3d_velo], fig=fig, color=color)
mlab.plot3d([x1, x2], [y1, y2], [z1, z2],
color=(0.5, 0.5, 0.5),
tube_radius=None,
line_width=1,
figure=fig)
# for prediction
if objects_pred is not None:
color = (1, 0, 0)
for obj in objects_pred:
if obj.type == "DontCare":
continue
# Draw 3d bounding box
box3d_pts_2d, box3d_pts_3d = utils.compute_box_3d(obj, calib.P2)
box3d_pts_3d_velo = utils.project_rect_to_velo(box3d_pts_3d, calib)
print("box3d_pts_3d_velo:")
print(box3d_pts_3d_velo)
draw_gt_boxes3d([box3d_pts_3d_velo], fig=fig, color=color)
# Draw heading arrow
ori3d_pts_2d, ori3d_pts_3d = utils.compute_orientation_3d(
obj, calib.P2)
ori3d_pts_3d_velo = utils.project_rect_to_velo(ori3d_pts_3d, calib)
x1, y1, z1 = ori3d_pts_3d_velo[0, :]
x2, y2, z2 = ori3d_pts_3d_velo[1, :]
mlab.plot3d([x1, x2], [y1, y2], [z1, z2],
color=color,
tube_radius=None,
line_width=1,
figure=fig)
mlab.show(1)
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
def sample_all(self, gt_boxes, gt_types, road_planes=None, calib=None):
avoid_coll_boxes = gt_boxes
sampled = []
sampled_gt_boxes = []
for i, class_name in enumerate(self._sample_classes):
sampled_num_per_class = int(
self._sample_max_num[i] -
np.sum([n == class_name for n in gt_types]))
if sampled_num_per_class > 0:
sampled_cls = self.sample(avoid_coll_boxes,
sampled_num_per_class, i)
else:
sampled_cls = []
sampled += sampled_cls
if len(sampled_cls) > 0:
if len(sampled_cls) == 1:
sampled_gt_box = sampled_cls[0]["box3d_lidar"][np.newaxis,
...]
else:
sampled_gt_box = np.stack(
[s["box3d_lidar"] for s in sampled_cls], axis=0)
sampled_gt_boxes += [sampled_gt_box]
avoid_coll_boxes = np.concatenate(
[avoid_coll_boxes, sampled_gt_box], axis=0)
if len(sampled) > 0:
sampled_gt_boxes = np.concatenate(sampled_gt_boxes, axis=0)
if road_planes is not None:
center = sampled_gt_boxes[:, 0:3]
a, b, c, d = road_planes
center_cam = project_velo_to_rect(center, calib)
cur_height_cam = (-d - a * center_cam[:, 0] -
c * center_cam[:, 2]) / b
center_cam[:, 1] = cur_height_cam
lidar_tmp_point = project_rect_to_velo(center_cam, calib)
cur_lidar_height = lidar_tmp_point[:, 2]
mv_height = sampled_gt_boxes[:, 2] - cur_lidar_height
sampled_gt_boxes[:, 2] -= mv_height # lidar view
s_points_list = []
sampled_gt_types = []
for i, info in enumerate(sampled):
s_points = np.fromfile(str(
pathlib.Path(self.root_path) / info["path"]),
dtype=np.float32)
s_points = s_points.reshape([-1, 4])
s_points[:, :3] += info["box3d_lidar"][:3]
if road_planes is not None:
s_points[:, 2] -= mv_height[i]
s_points_list.append(s_points)
sampled_gt_types.append(info['name'])
return sampled_gt_boxes.astype(
np.float32), sampled_gt_types, np.concatenate(s_points_list,
axis=0)
else:
return np.empty((0, 7), dtype=np.float32), [], np.empty(
(0, 4), dtype=np.float32)