def render(self,
axis: Axes,
view: np.ndarray = np.eye(3),
normalize: bool = False,
colors: Tuple = ('b', 'r', 'k'),
linewidth: float = 2) -> None:
"""
Renders the box in the provided Matplotlib axis.
:param axis: Axis onto which the box should be drawn.
:param view: <np.array: 3, 3>. Define a projection in needed (e.g. for drawing projection in an image).
:param normalize: Whether to normalize the remaining coordinate.
:param colors: (<Matplotlib.colors>: 3). Valid Matplotlib colors (<str> or normalized RGB tuple) for front,
back and sides.
:param linewidth: Width in pixel of the box sides.
"""
corners = view_points(self.corners(), view, normalize=normalize)[:2, :]
def draw_rect(selected_corners, color):
prev = selected_corners[-1]
for corner in selected_corners:
axis.plot([prev[0], corner[0]], [prev[1], corner[1]],
color=color,
linewidth=linewidth)
prev = corner
# Draw the sides
for i in range(4):
axis.plot([corners.T[i][0], corners.T[i + 4][0]],
[corners.T[i][1], corners.T[i + 4][1]],
color=colors[2],
linewidth=linewidth)
# Draw front (first 4 corners) and rear (last 4 corners) rectangles(3d)/lines(2d)
draw_rect(corners.T[:4], colors[0])
draw_rect(corners.T[4:], colors[1])
# Draw line indicating the front
center_bottom_forward = np.mean(corners.T[2:4], axis=0)
center_bottom = np.mean(corners.T[[2, 3, 7, 6]], axis=0)
axis.plot([center_bottom[0], center_bottom_forward[0]],
[center_bottom[1], center_bottom_forward[1]],
color=colors[0],
linewidth=linewidth)
def _render_helper(self, color_channel: int, ax: Axes, view: np.ndarray,
x_lim: Tuple[float, float], y_lim: Tuple[float, float],
marker_size: float) -> None:
"""
Helper function for rendering.
:param color_channel: Point channel to use as color.
:param ax: Axes on which to render the points.
:param view: <np.float: n, n>. Defines an arbitrary projection (n <= 4).
:param x_lim: (min, max).
:param y_lim: (min, max).
:param marker_size: Marker size.
"""
points = view_points(self.points[:3, :], view, normalize=False)
ax.scatter(points[0, :],
points[1, :],
c=self.points[color_channel, :],
s=marker_size)
ax.set_xlim(x_lim[0], x_lim[1])
ax.set_ylim(y_lim[0], y_lim[1])
def render_cv2(self,
im: np.ndarray,
view: np.ndarray = np.eye(3),
normalize: bool = False,
colors: Tuple = ((0, 0, 255), (255, 0, 0), (155, 155, 155)),
linewidth: int = 2) -> None:
"""
Renders box using OpenCV2.
:param im: <np.array: width, height, 3>. Image array. Channels are in BGR order.
:param view: <np.array: 3, 3>. Define a projection if needed (e.g. for drawing projection in an image).
:param normalize: Whether to normalize the remaining coordinate.
:param colors: ((R, G, B), (R, G, B), (R, G, B)). Colors for front, side & rear.
:param linewidth: Linewidth for plot.
"""
corners = view_points(self.corners(), view, normalize=normalize)[:2, :]
def draw_rect(selected_corners, color):
prev = selected_corners[-1]
for corner in selected_corners:
cv2.line(im, (int(prev[0]), int(prev[1])),
(int(corner[0]), int(corner[1])), color, linewidth)
prev = corner
# Draw the sides
for i in range(4):
cv2.line(im, (int(corners.T[i][0]), int(corners.T[i][1])),
(int(corners.T[i + 4][0]), int(corners.T[i + 4][1])),
colors[2][::-1], linewidth)
# Draw front (first 4 corners) and rear (last 4 corners) rectangles(3d)/lines(2d)
draw_rect(corners.T[:4], colors[0][::-1])
draw_rect(corners.T[4:], colors[1][::-1])
# Draw line indicating the front
center_bottom_forward = np.mean(corners.T[2:4], axis=0)
center_bottom = np.mean(corners.T[[2, 3, 7, 6]], axis=0)
cv2.line(
im, (int(center_bottom[0]), int(center_bottom[1])),
(int(center_bottom_forward[0]), int(center_bottom_forward[1])),
colors[0][::-1], linewidth)
def project_kitti_box_to_image(box: Box, p_left: np.ndarray, imsize: Tuple[int, int]) \
-> Union[None, Tuple[int, int, int, int]]:
"""
Projects 3D box into KITTI image FOV.
:param box: 3D box in KITTI reference frame.
:param p_left: <np.float: 3, 4>. Projection matrix.
:param imsize: (width, height). Image size.
:return: (xmin, ymin, xmax, ymax). Bounding box in image plane or None if box is not in the image.
"""
# Create a new box.
box = box.copy()
# KITTI defines the box center as the bottom center of the object.
# We use the true center, so we need to adjust half height in negative y direction.
box.translate(np.array([0, -box.wlh[2] / 2, 0]))
# Check that some corners are inside the image.
corners = np.array([corner for corner in box.corners().T if corner[2] > 0]).T
if len(corners) == 0:
return None
# Project corners that are in front of the camera to 2d to get bbox in pixel coords.
imcorners = view_points(corners, p_left, normalize=True)[:2]
bbox = (np.min(imcorners[0]), np.min(imcorners[1]), np.max(imcorners[0]), np.max(imcorners[1]))
# Crop bbox to prevent it extending outside image.
bbox_crop = tuple(max(0, b) for b in bbox)
bbox_crop = (min(imsize[0], bbox_crop[0]),
min(imsize[0], bbox_crop[1]),
min(imsize[0], bbox_crop[2]),
min(imsize[1], bbox_crop[3]))
# Detect if a cropped box is empty.
if bbox_crop[0] >= bbox_crop[2] or bbox_crop[1] >= bbox_crop[3]:
return None
return bbox_crop
def project_pts_to_image(self, pointcloud: LidarPointCloud, token: str) -> np.ndarray:
"""
Project lidar points into image.
:param pointcloud: The LidarPointCloud in nuScenes lidar frame.
:param token: Unique KITTI token.
:return: <np.float: N, 3.> X, Y are points in image pixel coordinates. Z is depth in image.
"""
# Copy and convert pointcloud.
pc_image = LidarPointCloud(points=pointcloud.points.copy())
pc_image.rotate(self.kitti_to_nu_lidar_inv) # Rotate to KITTI lidar.
# Transform pointcloud to camera frame.
transforms = self.get_transforms(token, root=self.root)
pc_image.rotate(transforms['velo_to_cam']['R'])
pc_image.translate(transforms['velo_to_cam']['T'])
# Project to image.
depth = pc_image.points[2, :]
points_fov = view_points(pc_image.points[:3, :], transforms['p_combined'], normalize=True)
points_fov[2, :] = depth
return points_fov
def _nuscenes(path: str) -> Dataset: # noqa
images: Dataset.Images = []
classes: Dataset.Classes = DETECTION_NAMES
annotations: Dataset.Annotations = {}
nusc = NuScenes(version="v1.0-trainval", dataroot=path)
for sample in nusc.sample:
for cam in CAMERAS:
cam_token = sample['data'][cam]
#Returns the data path as well as all annotations related to that sample_data.
#Note that the boxes are transformed into the current sensor's coordinate frame.
data_path, boxes, camera_intrinsic = nusc.get_sample_data(
cam_token, box_vis_level=BoxVisibility.ALL)
images.append(data_path)
annotations[data_path] = []
for box in boxes:
img_corners = view_points(box.corners(),
camera_intrinsic,
normalize=True)[:2, :]
# Take an outer rect of the 3d projection
xmin = img_corners[0].min()
xmax = img_corners[0].max()
ymin = img_corners[1].min()
ymax = img_corners[1].max()
bounds = Bounds2D(xmin, ymin, xmax - xmin, ymax - ymin)
label = category_to_detection_name(box.name)
if label is not None:
class_index = classes.index(
category_to_detection_name(box.name))
annotations[data_path].append(Object2D(
bounds, class_index))
return Dataset(DATASET_NAME, images, classes, annotations)
def render_sample_data(self,
token: str,
sensor_modality: str = 'lidar',
with_anns: bool = True,
axes_limit: float = 30,
ax: Axes = None,
view_3d: np.ndarray = np.eye(4),
color_func: Any = None,
augment_previous: bool = False,
box_linewidth: int = 2,
filter_classes: List[str] = None,
max_dist: float = None,
out_path: str = None,
render_2d: bool = False) -> None:
"""
Render sample data onto axis. Visualizes lidar in nuScenes lidar frame and camera in camera frame.
:param token: KITTI token.
:param sensor_modality: The modality to visualize, e.g. lidar or camera.
:param with_anns: Whether to draw annotations.
:param axes_limit: Axes limit for lidar data (measured in meters).
:param ax: Axes onto which to render.
:param view_3d: 4x4 view matrix for 3d views.
:param color_func: Optional function that defines the render color given the class name.
:param augment_previous: Whether to augment an existing plot (does not redraw pointcloud/image).
:param box_linewidth: Width of the box lines.
:param filter_classes: Optionally filter the classes to render.
:param max_dist: Maximum distance in m to still draw a box.
:param out_path: Optional path to save the rendered figure to disk.
:param render_2d: Whether to render 2d boxes (only works for camera data).
"""
# Default settings.
if color_func is None:
color_func = NuScenesExplorer.get_color
boxes = self.get_boxes(token, filter_classes=filter_classes, max_dist=max_dist) # In nuScenes lidar frame.
if sensor_modality == 'lidar':
# Load pointcloud.
pc = self.get_pointcloud(token, self.root) # In KITTI lidar frame.
pc.rotate(self.kitti_to_nu_lidar.rotation_matrix) # In nuScenes lidar frame.
# Alternative options:
# depth = pc.points[1, :]
# height = pc.points[2, :]
intensity = pc.points[3, :]
# Project points to view.
points = view_points(pc.points[:3, :], view_3d, normalize=False)
coloring = intensity
if ax is None:
_, ax = plt.subplots(1, 1, figsize=(9, 9))
if not augment_previous:
ax.scatter(points[0, :], points[1, :], c=coloring, s=1)
ax.set_xlim(-axes_limit, axes_limit)
ax.set_ylim(-axes_limit, axes_limit)
if with_anns:
for box in boxes:
color = np.array(color_func(box.name)) / 255
box.render(ax, view=view_3d, colors=(color, color, 'k'), linewidth=box_linewidth)
elif sensor_modality == 'camera':
im_path = KittiDB.get_filepath(token, 'image_2', root=self.root)
im = Image.open(im_path)
if ax is None:
_, ax = plt.subplots(1, 1, figsize=(9, 16))
if not augment_previous:
ax.imshow(im)
ax.set_xlim(0, im.size[0])
ax.set_ylim(im.size[1], 0)
if with_anns:
if render_2d:
# Use KITTI's 2d boxes.
boxes_2d, names = self.get_boxes_2d(token, filter_classes=filter_classes)
for box, name in zip(boxes_2d, names):
color = np.array(color_func(name)) / 255
ax.plot([box[0], box[0]], [box[1], box[3]], color=color, linewidth=box_linewidth)
ax.plot([box[2], box[2]], [box[1], box[3]], color=color, linewidth=box_linewidth)
ax.plot([box[0], box[2]], [box[1], box[1]], color=color, linewidth=box_linewidth)
ax.plot([box[0], box[2]], [box[3], box[3]], color=color, linewidth=box_linewidth)
else:
# Project 3d boxes to 2d.
transforms = self.get_transforms(token, self.root)
for box in boxes:
# Undo the transformations in get_boxes() to get back to the camera frame.
box.rotate(self.kitti_to_nu_lidar_inv) # In KITTI lidar frame.
box.rotate(Quaternion(matrix=transforms['velo_to_cam']['R']))
box.translate(transforms['velo_to_cam']['T']) # In KITTI camera frame, un-rectified.
box.rotate(Quaternion(matrix=transforms['r0_rect'])) # In KITTI camera frame, rectified.
# Filter boxes outside the image (relevant when visualizing nuScenes data in KITTI format).
if not box_in_image(box, transforms['p_left'][:3, :3], im.size, vis_level=BoxVisibility.ANY):
continue
# Render.
color = np.array(color_func(box.name)) / 255
box.render(ax, view=transforms['p_left'][:3, :3], normalize=True, colors=(color, color, 'k'),
linewidth=box_linewidth)
else:
raise ValueError("Unrecognized modality {}.".format(sensor_modality))
ax.axis('off')
ax.set_title(token)
ax.set_aspect('equal')
# Render to disk.
plt.tight_layout()
if out_path is not None:
plt.savefig(out_path)
def get_2d_boxes(sample_data_token: str,
visibilities: List[str]) -> List[OrderedDict]:
"""
Get the 2D annotation records for a given `sample_data_token`.
:param sample_data_token: Sample data token belonging to a camera keyframe.
:param visibilities: Visibility filter.
:return: List of 2D annotation record that belongs to the input `sample_data_token`
"""
# Get the sample data and the sample corresponding to that sample data.
sd_rec = nusc.get('sample_data', sample_data_token)
assert sd_rec[
'sensor_modality'] == 'camera', 'Error: get_2d_boxes only works for camera sample_data!'
if not sd_rec['is_key_frame']:
raise ValueError(
'The 2D re-projections are available only for keyframes.')
s_rec = nusc.get('sample', sd_rec['sample_token'])
# Get the calibrated sensor and ego pose record to get the transformation matrices.
cs_rec = nusc.get('calibrated_sensor', sd_rec['calibrated_sensor_token'])
pose_rec = nusc.get('ego_pose', sd_rec['ego_pose_token'])
camera_intrinsic = np.array(cs_rec['camera_intrinsic'])
# Get all the annotation with the specified visibilties.
ann_recs = [
nusc.get('sample_annotation', token) for token in s_rec['anns']
]
ann_recs = [
ann_rec for ann_rec in ann_recs
if (ann_rec['visibility_token'] in visibilities)
]
repro_recs = []
for ann_rec in ann_recs:
# Augment sample_annotation with token information.
ann_rec['sample_annotation_token'] = ann_rec['token']
ann_rec['sample_data_token'] = sample_data_token
# Get the box in global coordinates.
box = nusc.get_box(ann_rec['token'])
# Move them to the ego-pose frame.
box.translate(-np.array(pose_rec['translation']))
box.rotate(Quaternion(pose_rec['rotation']).inverse)
# Move them to the calibrated sensor frame.
box.translate(-np.array(cs_rec['translation']))
box.rotate(Quaternion(cs_rec['rotation']).inverse)
# Filter out the corners that are not in front of the calibrated sensor.
corners_3d = box.corners()
in_front = np.argwhere(corners_3d[2, :] > 0).flatten()
corners_3d = corners_3d[:, in_front]
# Project 3d box to 2d.
corner_coords = view_points(corners_3d, camera_intrinsic,
True).T[:, :2].tolist()
# Keep only corners that fall within the image.
final_coords = post_process_coords(corner_coords)
# Skip if the convex hull of the re-projected corners does not intersect the image canvas.
if final_coords is None:
continue
else:
min_x, min_y, max_x, max_y = final_coords
# Generate dictionary record to be included in the .json file.
repro_rec = generate_record(ann_rec, min_x, min_y, max_x, max_y,
sample_data_token, sd_rec['filename'])
repro_recs.append(repro_rec)
return repro_recs
def visualize_sample(nusc: NuScenes,
sample_token: str,
gt_boxes: EvalBoxes,
pred_boxes: EvalBoxes,
nsweeps: int = 1,
conf_th: float = 0.15,
eval_range: float = 50,
verbose: bool = True,
savepath: str = None) -> None:
"""
Visualizes a sample from BEV with annotations and detection results.
:param nusc: NuScenes object.
:param sample_token: The nuScenes sample token.
:param gt_boxes: Ground truth boxes grouped by sample.
:param pred_boxes: Prediction grouped by sample.
:param nsweeps: Number of sweeps used for lidar visualization.
:param conf_th: The confidence threshold used to filter negatives.
:param eval_range: Range in meters beyond which boxes are ignored.
:param verbose: Whether to print to stdout.
:param savepath: If given, saves the the rendering here instead of displaying.
"""
# Retrieve sensor & pose records.
sample_rec = nusc.get('sample', sample_token)
sd_record = nusc.get('sample_data', sample_rec['data']['LIDAR_TOP'])
cs_record = nusc.get('calibrated_sensor', sd_record['calibrated_sensor_token'])
pose_record = nusc.get('ego_pose', sd_record['ego_pose_token'])
# Get boxes.
boxes_gt_global = gt_boxes[sample_token]
boxes_est_global = pred_boxes[sample_token]
# Map GT boxes to lidar.
boxes_gt = boxes_to_sensor(boxes_gt_global, pose_record, cs_record)
# Map EST boxes to lidar.
boxes_est = boxes_to_sensor(boxes_est_global, pose_record, cs_record)
# Add scores to EST boxes.
for box_est, box_est_global in zip(boxes_est, boxes_est_global):
box_est.score = box_est_global.detection_score
# Get point cloud in lidar frame.
pc, _ = LidarPointCloud.from_file_multisweep(nusc, sample_rec, 'LIDAR_TOP', 'LIDAR_TOP', nsweeps=nsweeps)
# Init axes.
_, ax = plt.subplots(1, 1, figsize=(9, 9))
# Show point cloud.
points = view_points(pc.points[:3, :], np.eye(4), normalize=False)
dists = np.sqrt(np.sum(pc.points[:2, :] ** 2, axis=0))
colors = np.minimum(1, dists / eval_range)
ax.scatter(points[0, :], points[1, :], c=colors, s=0.2)
# Show ego vehicle.
ax.plot(0, 0, 'x', color='black')
# Show GT boxes.
for box in boxes_gt:
box.render(ax, view=np.eye(4), colors=('g', 'g', 'g'), linewidth=2)
# Show EST boxes.
for box in boxes_est:
# Show only predictions with a high score.
assert not np.isnan(box.score), 'Error: Box score cannot be NaN!'
if box.score >= conf_th:
box.render(ax, view=np.eye(4), colors=('b', 'b', 'b'), linewidth=1)
# Limit visible range.
axes_limit = eval_range + 3 # Slightly bigger to include boxes that extend beyond the range.
ax.set_xlim(-axes_limit, axes_limit)
ax.set_ylim(-axes_limit, axes_limit)
# Show / save plot.
if verbose:
print('Rendering sample token %s' % sample_token)
plt.title(sample_token)
if savepath is not None:
plt.savefig(savepath)
plt.close()
else:
plt.show()
def pointcloud_color_from_image(
nusc: NuScenes, pointsensor_token: str,
camera_token: str) -> Tuple[np.array, np.array]:
"""
Given a point sensor (lidar/radar) token and camera sample_data token, load point-cloud and map it to the image
plane, then retrieve the colors of the closest image pixels.
:param nusc: NuScenes instance.
:param pointsensor_token: Lidar/radar sample_data token.
:param camera_token: Camera sample data token.
:return (coloring <np.float: 3, n>, mask <np.bool: m>). Returns the colors for n points that reproject into the
image out of m total points. The mask indicates which points are selected.
"""
cam = nusc.get('sample_data', camera_token)
pointsensor = nusc.get('sample_data', pointsensor_token)
pc = LidarPointCloud.from_file(
osp.join(nusc.dataroot, pointsensor['filename']))
im = Image.open(osp.join(nusc.dataroot, cam['filename']))
# Points live in the point sensor frame. So they need to be transformed via global to the image plane.
# First step: transform the point-cloud to the ego vehicle frame for the timestamp of the sweep.
cs_record = nusc.get('calibrated_sensor',
pointsensor['calibrated_sensor_token'])
pc.rotate(Quaternion(cs_record['rotation']).rotation_matrix)
pc.translate(np.array(cs_record['translation']))
# Second step: transform to the global frame.
poserecord = nusc.get('ego_pose', pointsensor['ego_pose_token'])
pc.rotate(Quaternion(poserecord['rotation']).rotation_matrix)
pc.translate(np.array(poserecord['translation']))
# Third step: transform into the ego vehicle frame for the timestamp of the image.
poserecord = nusc.get('ego_pose', cam['ego_pose_token'])
pc.translate(-np.array(poserecord['translation']))
pc.rotate(Quaternion(poserecord['rotation']).rotation_matrix.T)
# Fourth step: transform into the camera.
cs_record = nusc.get('calibrated_sensor', cam['calibrated_sensor_token'])
pc.translate(-np.array(cs_record['translation']))
pc.rotate(Quaternion(cs_record['rotation']).rotation_matrix.T)
# Fifth step: actually take a "picture" of the point cloud.
# Grab the depths (camera frame z axis points away from the camera).
depths = pc.points[2, :]
# Take the actual picture (matrix multiplication with camera-matrix + renormalization).
points = view_points(pc.points[:3, :],
np.array(cs_record['camera_intrinsic']),
normalize=True)
# Remove points that are either outside or behind the camera. Leave a margin of 1 pixel for aesthetic reasons.
mask = np.ones(depths.shape[0], dtype=bool)
mask = np.logical_and(mask, depths > 0)
mask = np.logical_and(mask, points[0, :] > 1)
mask = np.logical_and(mask, points[0, :] < im.size[0] - 1)
mask = np.logical_and(mask, points[1, :] > 1)
mask = np.logical_and(mask, points[1, :] < im.size[1] - 1)
points = points[:, mask]
# Pick the colors of the points
im_data = np.array(im)
coloring = np.zeros(points.shape)
for i, p in enumerate(points.transpose()):
point = p[:2].round().astype(np.int32)
coloring[:, i] = im_data[point[1], point[0], :]
return coloring, mask