# Academic Software License: Copyright © 2024.
import os
import cv2
import json
import time
import flow_vis
import argparse
import traceback
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import colormaps as cm
from pyquaternion import Quaternion as Q
CAM_NAME = [
'CAM_FRONT_LEFT',
'CAM_FRONT',
'CAM_FRONT_RIGHT',
'CAM_BACK_LEFT',
'CAM_BACK',
'CAM_BACK_RIGHT'
]
RAD_NAME = ['RAD_LEFT', 'RAD_FRONT', 'RAD_RIGHT', 'RAD_BACK']
OBJECT_CLASSES = {
12: 'pedestrian',
13: 'rider',
14: 'car',
15: 'truck',
16: 'bus',
18: 'motorcycle',
19: 'bicycle'
}
SIMBEV_PALETTE = {
'road': (196, 80, 196),
'car': (0, 128, 240),
'truck': (128, 240, 64),
'bus': (0, 144, 0),
'motorcycle': (240, 240, 0),
'bicycle': (0, 240, 240),
'rider': (240, 144, 0),
'pedestrian': (240, 0, 0)
}
RANGE = np.linspace(0.0, 1.0, 256)
RAINBOW = np.array(cm.get_cmap('rainbow')(RANGE))[:, :3]
LABEL_COLORS = np.array([
(255, 255, 255), # None
(128, 64, 128), # Road
(244, 35, 232), # Sidewalk
(70, 70, 70), # Building
(102, 102, 156), # Wall
(190, 153, 153), # Fence
(153, 153, 153), # Pole
(250, 170, 30), # TrafficLight
(220, 220, 0), # TrafficSign
(107, 142, 35), # Vegetation
(152, 251, 152), # Terrain
(70, 130, 180), # Sky
(220, 20, 60), # Pedestrian
(255, 0, 0), # Rider
(0, 0, 142), # Car
(0, 0, 70), # Truck
(0, 60, 100), # Bus
(0, 80, 100), # Train
(0, 0, 230), # Motorcycle
(119, 11, 32), # Bicycle
(110, 190, 160), # Static
(170, 120, 50), # Dynamic
(55, 90, 80), # Other
(45, 60, 150), # Water
(227, 227, 227), # RoadLine
(81, 0, 81), # Ground
(150, 100, 100), # Bridge
(230, 150, 140), # RailTrack
(180, 165, 180), # GuardRail
]) / 255.0
argparser = argparse.ArgumentParser(description='SimBEV visualization tool.')
argparser.add_argument(
'mode',
nargs='+',
help='visualization mode (all, rgb, depth, flow, lidar, lidar3d, lidar-with-bbox, lidar3d-with-bbox, '
'semantic-lidar, semantic-lidar3d, radar, radar3d, radar-with-bbox, radar3d-with-bbox)'
)
argparser.add_argument(
'--path',
default='/dataset',
help='path for saving the dataset (default: /dataset)'
)
argparser.add_argument(
'-s', '--scene',
type=int,
default=[-1],
nargs='+',
help='scene number(s) (default: -1, i.e. all scenes)'
)
argparser.add_argument(
'-f', '--frame',
type=int,
default=[-1],
nargs='+',
help='frame number(s) (default: -1, i.e. all frames)'
)
args = argparser.parse_args()
def draw_bbox(canvas, corners, labels, bbox_color=None, thickness=1):
'''
Draw bounding boxes on the canvas.
Args:
canvas: canvas to draw on.
corners: bounding box corners.
labels: bounding box labels.
bbox_color: bounding box color.
thickness: bounding box line thickness.
'''
# Filter out bounding boxes that are behind the camera.
indices = np.all(corners[..., 2] > 0, axis=1)
corners = corners[indices]
labels = labels[indices]
# Sort bounding boxes by their distance to the camera.
indices = np.argsort(-np.min(corners[..., 2], axis=1))
corners = corners[indices]
labels = labels[indices]
# Find the pixels corresponding to bounding box corners.
corners = corners.reshape(-1, 3)
corners[:, 2] = np.clip(corners[:, 2], a_min=1e-5, a_max=1e5)
corners[:, 0] /= corners[:, 2]
corners[:, 1] /= corners[:, 2]
corners = corners[..., :2].reshape(-1, 8, 2)
# Draw bounding box lines.
if corners is not None:
for index in range(corners.shape[0]):
name = labels[index]
for start, end in [
(0, 1),
(0, 2),
(0, 4),
(1, 3),
(1, 5),
(2, 3),
(2, 6),
(3, 7),
(4, 5),
(4, 6),
(5, 7),
(6, 7)
]:
cv2.line(
canvas,
corners[index, start].astype(int),
corners[index, end].astype(int),
bbox_color or SIMBEV_PALETTE[name],
thickness,
cv2.LINE_AA,
)
canvas = canvas.astype(np.uint8)
return canvas
def visualize_image(fpath, image, corners=None, labels=None, bbox_color=None, thickness=2):
'''
Visualize image with bounding boxes.
Args:
fpath: file path to save the image.
image: image to visualize.
corners: bounding box corners.
labels: bounding box labels.
color: bounding box color.
thickness: bounding box line thickness.
'''
canvas = image.copy()
if corners is not None and corners.shape[0] > 0:
canvas = cv2.cvtColor(canvas, cv2.COLOR_RGB2BGR)
canvas = draw_bbox(canvas, corners, labels, bbox_color=bbox_color, thickness=thickness)
canvas = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)
# Save the image.
cv2.imwrite(fpath, canvas)
def visualize_point_cloud(
fpath,
point_cloud,
corners=None,
labels=None,
color=None,
xlim=(-72, 72),
ylim=(-72, 72),
radius=16,
thickness=16
):
'''
Visualize point cloud from above with bounding boxes.
Args:
fpath: file path to save the image.
point_cloud: point cloud to visualize.
corners: bounding box corners.
labels: bounding box labels.
color: point cloud color(s).
xlim: x-axis limits.
ylim: y-axis limits.
radius: display point radius.
thickness: bounding box line thickness.
'''
fig = plt.figure(figsize=(xlim[1] - xlim[0], ylim[1] - ylim[0]))
ax = plt.gca()
ax.set_xlim(*xlim)
ax.set_ylim(*ylim)
ax.set_aspect(1)
ax.set_axis_off()
# Filter out points that are outside the limits.
mask = np.logical_and.reduce([
point_cloud[:, 0] >= xlim[0],
point_cloud[:, 0] < xlim[1],
point_cloud[:, 1] >= ylim[0],
point_cloud[:, 1] < ylim[1]
])
point_cloud = point_cloud[mask]
# Plot the point cloud.
if point_cloud is not None:
if color is None:
plt.scatter(
point_cloud[:, 0],
point_cloud[:, 1],
s=radius,
c=np.log(np.linalg.norm(point_cloud, axis=1)),
cmap='rainbow'
)
else:
color = color[mask]
plt.scatter(point_cloud[:, 0], point_cloud[:, 1], s=radius, c=color)
# Draw bounding box lines.
if corners is not None:
coords = corners[:, [0, 2, 6, 4, 0], :2]
for index in range(coords.shape[0]):
name = labels[index]
plt.plot(
coords[index, :, 0],
coords[index, :, 1],
color=np.array(SIMBEV_PALETTE[name]) / 255.0,
linewidth=thickness
)
# Save the image.
fig.savefig(fpath, dpi=16, facecolor='white', format='jpg', bbox_inches='tight', pad_inches=0)
plt.close()
def visualize_point_cloud_3d(
fpath,
point_cloud,
canvas,
corners=None,
labels=None,
color=None,
bbox_color=None,
thickness=1
):
'''
Visualize point cloud in 3D with bounding boxes.
Args:
fpath: file path to save the image.
point_cloud: point cloud to visualize.
canvas: canvas to draw on.
corners: bounding box corners.
labels: bounding box labels.
color: point cloud color(s).
thickness: bounding box line thickness.
'''
if color is not None:
canvas[point_cloud[:, 1].astype(int), point_cloud[:, 0].astype(int), :] = color
if corners is not None and corners.shape[0] > 0:
canvas = draw_bbox(canvas, corners, labels, bbox_color=bbox_color, thickness=thickness)
canvas = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)
cv2.imwrite(fpath, canvas)
def transform_bbox(gt_det, transform):
'''
Transform bounding boxes from the global coordinate system to the desired
coordinate system.
Args:
gt_det: ground truth objects.
transform: transformation matrix.
Returns:
corners: transformed bounding box corners.
labels: bounding box labels.
'''
corners = []
labels = []
# Transform bounding boxes from the global coordinate system to the
# desired coordinate system.
for det_object in gt_det:
if det_object['valid_flag']:
for tag in det_object['semantic_tags']:
if tag in OBJECT_CLASSES.keys():
global_bbox_corners = np.append(det_object['bounding_box'], np.ones((8, 1)), 1)
bbox_corners = (transform @ global_bbox_corners.T)[:3].T
corners.append(bbox_corners)
labels.append(OBJECT_CLASSES[tag])
corners = np.array(corners)
labels = np.array(labels)
return corners, labels
def main(mode):
try:
print(f'Visualizing {mode}...')
if mode == 'rgb':
for camera in CAM_NAME:
os.makedirs(f'{args.path}/simbev/viz/RGB-{camera}', exist_ok=True)
if mode == 'depth':
for camera in CAM_NAME:
os.makedirs(f'{args.path}/simbev/viz/DPT-{camera}', exist_ok=True)
if mode == 'flow':
for camera in CAM_NAME:
os.makedirs(f'{args.path}/simbev/viz/FLW-{camera}', exist_ok=True)
if mode == 'lidar':
os.makedirs(f'{args.path}/simbev/viz/LIDAR', exist_ok=True)
if mode == 'lidar3d':
os.makedirs(f'{args.path}/simbev/viz/LIDAR3D', exist_ok=True)
if mode == 'lidar-with-bbox':
os.makedirs(f'{args.path}/simbev/viz/LIDARwBBOX', exist_ok=True)
if mode == 'lidar3d-with-bbox':
os.makedirs(f'{args.path}/simbev/viz/LIDAR3DwBBOX', exist_ok=True)
if mode == 'semantic-lidar':
os.makedirs(f'{args.path}/simbev/viz/SEG-LIDAR', exist_ok=True)
if mode == 'semantic-lidar3d':
os.makedirs(f'{args.path}/simbev/viz/SEG-LIDAR3D', exist_ok=True)
if mode == 'radar':
os.makedirs(f'{args.path}/simbev/viz/RADAR', exist_ok=True)
if mode == 'radar3d':
os.makedirs(f'{args.path}/simbev/viz/RADAR3D', exist_ok=True)
if mode == 'radar-with-bbox':
os.makedirs(f'{args.path}/simbev/viz/RADARwBBOX', exist_ok=True)
if mode == 'radar3d-with-bbox':
os.makedirs(f'{args.path}/simbev/viz/RADAR3DwBBOX', exist_ok=True)
for split in ['train', 'val', 'test']:
if os.path.exists(f'{args.path}/simbev/infos/simbev_infos_{split}.json'):
with open(f'{args.path}/simbev/infos/simbev_infos_{split}.json', 'r') as f:
infos = json.load(f)
metadata = infos['metadata']
# Get the list of scenes to visualize.
if args.scene == [-1]:
scene_list = []
for scene in infos['data']:
scene_list.append(int(scene[-4:]))
else:
scene_list = args.scene
for scene_number in scene_list:
if f'scene_{scene_number:04d}' in infos['data']:
print(f'Visualizing scene {scene_number}...')
scene_data = infos['data'][f'scene_{scene_number:04d}']['scene_data']
if args.frame == [-1]:
frame_list = list(range(len(scene_data)))
else:
frame_list = [frame for frame in args.frame if frame < len(scene_data) and frame >= 0]
for frame_number in frame_list:
frame_data = scene_data[frame_number]
if mode == 'rgb' or 'bbox' in mode:
# Load object bounding boxes.
gt_det_path = frame_data['GT_DET']
gt_det = np.load(gt_det_path, allow_pickle=True)
# Ego to global transformation.
ego2global = np.eye(4).astype(np.float32)
ego2global[:3, :3] = Q(frame_data['ego2global_rotation']).rotation_matrix
ego2global[:3, 3] = frame_data['ego2global_translation']
if 'bbox' in mode:
# Lidar to ego transformation.
lidar2ego = np.eye(4).astype(np.float32)
lidar2ego[:3, :3] = Q(metadata['LIDAR']['sensor2ego_rotation']).rotation_matrix
lidar2ego[:3, 3] = metadata['LIDAR']['sensor2ego_translation']
global2lidar = np.linalg.inv(ego2global @ lidar2ego)
if '3d' in mode:
view2lidar = np.eye(4).astype(np.float32)
view2lidar[:3, :3] = Q([0.415627, -0.572061, 0.572061, -0.415627]).rotation_matrix
view2lidar[:3, 3] = [-40.0, 0.0, 12.0]
lidar2view = np.linalg.inv(view2lidar)
camera_intrinsics = np.eye(4).astype(np.float32)
camera_intrinsics[:3, :3] = metadata['camera_intrinsics']
lidar2image = camera_intrinsics @ lidar2view
if mode == 'rgb':
for camera in CAM_NAME:
# Camera to ego transformation.
camera2ego = np.eye(4).astype(np.float32)
camera2ego[:3, :3] = Q(metadata[camera]['sensor2ego_rotation']).rotation_matrix
camera2ego[:3, 3] = metadata[camera]['sensor2ego_translation']
# Camera intrinsics.
camera_intrinsics = np.eye(4).astype(np.float32)
camera_intrinsics[:3, :3] = metadata['camera_intrinsics']
# Global to camera transformation.
global2camera = np.linalg.inv(ego2global @ camera2ego)
global2image = camera_intrinsics @ global2camera
corners, labels = transform_bbox(gt_det, global2image)
image = cv2.imread(frame_data['RGB-' + camera])
visualize_image(
f'{args.path}/simbev/viz/RGB-{camera}/SimBEV-scene-' \
f'{scene_number:04d}-frame-{frame_number:04d}-RGB-{camera}.jpg',
image,
corners=corners,
labels=labels
)
if mode == 'depth':
for camera in CAM_NAME:
image = cv2.imread(frame_data['DPT-' + camera])
normalized_distance = (
image[:, :, 2] + image[:, :, 1] * 256.0 + image[:, :, 0] * 256.0 * 256.0
) / (256.0 * 256.0 * 256.0 - 1)
log_distance = 255 * np.log(256.0 * normalized_distance + 1) / np.log(257.0)
cv2.imwrite(
f'{args.path}/simbev/viz/DPT-{camera}/SimBEV-scene-' \
f'{scene_number:04d}-frame-{frame_number:04d}-DPT-{camera}.jpg',
log_distance.astype(np.uint8)
)
if mode == 'flow':
for camera in CAM_NAME:
flow = np.load(frame_data['FLW-' + camera])['data']
image = flow_vis.flow_to_color(flow, convert_to_bgr=True)
cv2.imwrite(
f'{args.path}/simbev/viz/FLW-{camera}/SimBEV-scene-' \
f'{scene_number:04d}-frame-{frame_number:04d}-FLW-{camera}.jpg',
image
)
if mode in ['lidar', 'lidar3d', 'lidar-with-bbox', 'lidar3d-with-bbox']:
point_cloud = np.load(frame_data['LIDAR'])['data']
label_color = None
if mode in ['semantic-lidar', 'semantic-lidar3d']:
data = np.load(frame_data['SEG-LIDAR'])['data']
point_cloud = np.array([data['x'], data['y'], data['z']]).T
labels = np.array(data['ObjTag'])
label_color = LABEL_COLORS[labels]
if 'lidar3d' in mode:
distance = np.linalg.norm(point_cloud, axis=1)
point_cloud_3d = (
lidar2image @ np.append(point_cloud, np.ones((point_cloud.shape[0], 1)), 1).T
)[:3].T
# Filter out points that are behind the 3D view camera.
indices = point_cloud_3d[:, 2] > 0.0
point_cloud_3d = point_cloud_3d[indices]
distance = distance[indices]
if label_color is not None:
label_color = label_color[indices]
# Find the pixels corresponding to the point cloud.
point_cloud_3d[:, 2] = np.clip(point_cloud_3d[:, 2], a_min=1e-5, a_max=1e5)
point_cloud_3d[:, 0] /= point_cloud_3d[:, 2]
point_cloud_3d[:, 1] /= point_cloud_3d[:, 2]
# Set the 3D view camera dimensions.
width = camera_intrinsics[0, 2] * 2
height = camera_intrinsics[1, 2] * 2
canvas = np.zeros((int(height), int(width), 3), dtype=np.uint8)
canvas[:] = (255, 255, 255)
mask = np.logical_and.reduce([
point_cloud_3d[:, 0] >= 0,
point_cloud_3d[:, 0] < width,
point_cloud_3d[:, 1] >= 0,
point_cloud_3d[:, 1] < height
])
point_cloud_3d = point_cloud_3d[mask]
distance = distance[mask]
if label_color is not None:
label_color = label_color[mask]
# Calculate point cloud colors.
log_distance = np.log(distance)
log_distance_normalized = (
log_distance - log_distance.min()
) / (
log_distance.max() - log_distance.min() + 1e-6
)
color = np.c_[
np.interp(log_distance_normalized, RANGE, RAINBOW[:, 0]),
np.interp(log_distance_normalized, RANGE, RAINBOW[:, 1]),
np.interp(log_distance_normalized, RANGE, RAINBOW[:, 2])
] * 255.0
if mode == 'lidar':
visualize_point_cloud(
f'{args.path}/simbev/viz/LIDAR/' \
f'SimBEV-scene-{scene_number:04d}-frame-{frame_number:04d}-LIDAR.jpg',
point_cloud
)
if mode == 'lidar3d':
visualize_point_cloud_3d(
f'{args.path}/simbev/viz/LIDAR3D/' \
f'SimBEV-scene-{scene_number:04d}-frame-{frame_number:04d}-LIDAR3D.jpg',
point_cloud_3d,
canvas,
color=color
)
if mode == 'lidar-with-bbox':
corners, labels = transform_bbox(gt_det, global2lidar)
visualize_point_cloud(
f'{args.path}/simbev/viz/LIDARwBBOX/' \
f'SimBEV-scene-{scene_number:04d}-frame-{frame_number:04d}-LIDARwBBOX.jpg',
point_cloud,
corners=corners,
labels=labels
)
if mode == 'lidar3d-with-bbox':
global2image = lidar2image @ global2lidar
corners, labels = transform_bbox(gt_det, global2image)
visualize_point_cloud_3d(
f'{args.path}/simbev/viz/LIDAR3DwBBOX/' \
f'SimBEV-scene-{scene_number:04d}-frame-{frame_number:04d}-LIDAR3DwBBOX.jpg',
point_cloud_3d,
canvas,
corners=corners,
labels=labels,
color=color
)
if mode == 'semantic-lidar':
visualize_point_cloud(
f'{args.path}/simbev/viz/SEG-LIDAR/' \
f'SimBEV-scene-{scene_number:04d}-frame-{frame_number:04d}-SEG-LIDAR.jpg',
point_cloud,
color=label_color
)
if mode == 'semantic-lidar3d':
visualize_point_cloud_3d(
f'{args.path}/simbev/viz/SEG-LIDAR3D/' \
f'SimBEV-scene-{scene_number:04d}-frame-{frame_number:04d}-SEG-LIDAR3D.jpg',
point_cloud_3d,
canvas,
color=(label_color * 255.0).astype(np.uint8)
)
if 'radar' in mode:
point_cloud = []
velocity = []
for radar in RAD_NAME:
# Radar to lidar transformation.
radar2lidar = np.eye(4).astype(np.float32)
radar2lidar[:3, :3] = Q(metadata[radar]['sensor2lidar_rotation']).rotation_matrix
radar2lidar[:3, 3] = metadata[radar]['sensor2lidar_translation']
radar_points = np.load(frame_data[radar])['data']
velocity.append(radar_points[:, -1])
radar_points = radar_points[:, :-1]
# Convert the radar values of depth,
# altitude angle, and azimuth angle to
# x, y, and z coordinates.
x = radar_points[:, 0] * np.cos(radar_points[:, 1]) * np.cos(radar_points[:, 2])
y = radar_points[:, 0] * np.cos(radar_points[:, 1]) * np.sin(radar_points[:, 2])
z = radar_points[:, 0] * np.sin(radar_points[:, 1])
points = np.stack((x, y, z), axis=1)
points_transformed = (
radar2lidar @ np.append(points, np.ones((points.shape[0], 1)), 1).T
)[:3].T
point_cloud.append(points_transformed)
point_cloud = np.concatenate(point_cloud, axis=0)
velocity = np.concatenate(velocity, axis=0)
# Calculate point cloud colors.
log_velocity = np.log(1.0 + np.abs(velocity))
log_velocity_normalized = (
log_velocity - log_velocity.min()
) / (
log_velocity.max() - log_velocity.min() + 1e-6
)
color = np.c_[
np.interp(log_velocity_normalized, RANGE, RAINBOW[:, 0]),
np.interp(log_velocity_normalized, RANGE, RAINBOW[:, 1]),
np.interp(log_velocity_normalized, RANGE, RAINBOW[:, 2])
]
if 'radar3d' in mode:
point_cloud_3d = (
lidar2image @ np.append(point_cloud, np.ones((point_cloud.shape[0], 1)), 1).T
)[:3].T
# Filter out points that are behind the 3D view camera.
indices = point_cloud_3d[:, 2] > 0.0
point_cloud_3d = point_cloud_3d[indices]
color = color[indices]
# Find the pixels corresponding to the point cloud.
point_cloud_3d[:, 2] = np.clip(point_cloud_3d[:, 2], a_min=1e-5, a_max=1e5)
point_cloud_3d[:, 0] /= point_cloud_3d[:, 2]
point_cloud_3d[:, 1] /= point_cloud_3d[:, 2]
# Set the 3D view camera dimensions.
width = camera_intrinsics[0, 2] * 2
height = camera_intrinsics[1, 2] * 2
canvas = np.zeros((int(height), int(width), 3), dtype=np.uint8)
canvas[:] = (255, 255, 255)
mask = np.logical_and.reduce([
point_cloud_3d[:, 0] >= 0,
point_cloud_3d[:, 0] < width,
point_cloud_3d[:, 1] >= 0,
point_cloud_3d[:, 1] < height
])
point_cloud_3d = point_cloud_3d[mask]
color = color[mask]
if mode == 'radar':
visualize_point_cloud(
f'{args.path}/simbev/viz/RADAR/' \
f'SimBEV-scene-{scene_number:04d}-frame-{frame_number:04d}-RADAR.jpg',
point_cloud,
color=color,
radius=128
)
if mode == 'radar3d':
visualize_point_cloud_3d(
f'{args.path}/simbev/viz/RADAR3D/' \
f'SimBEV-scene-{scene_number:04d}-frame-{frame_number:04d}-RADAR3D.jpg',
point_cloud_3d,
canvas,
color=(color * 255.0).astype(np.uint8)
)
if mode == 'radar-with-bbox':
corners, labels = transform_bbox(gt_det, global2lidar)
visualize_point_cloud(
f'{args.path}/simbev/viz/RADARwBBOX/' \
f'SimBEV-scene-{scene_number:04d}-frame-{frame_number:04d}-RADARwBBOX.jpg',
point_cloud,
corners=corners,
labels=labels,
color=color,
radius=128
)
if mode == 'radar3d-with-bbox':
global2image = lidar2image @ global2lidar
corners, labels = transform_bbox(gt_det, global2image)
visualize_point_cloud_3d(
f'{args.path}/simbev/viz/RADAR3DwBBOX/' \
f'SimBEV-scene-{scene_number:04d}-frame-{frame_number:04d}-RADAR3DwBBOX.jpg',
point_cloud_3d,
canvas,
corners=corners,
labels=labels,
color=(color * 255.0).astype(np.uint8)
)
except Exception:
print(traceback.format_exc())
print('Killing the process...')
time.sleep(3.0)
if __name__ == '__main__':
try:
os.makedirs(f'{args.path}/simbev/viz', exist_ok=True)
if 'all' in args.mode:
mode_list = [
'rgb',
'depth',
'flow',
'lidar',
'lidar3d',
'lidar-with-bbox',
'lidar3d-with-bbox',
'semantic-lidar',
'semantic-lidar3d',
'radar',
'radar3d',
'radar-with-bbox',
'radar3d-with-bbox'
]
else:
mode_list = args.mode
for mode in mode_list:
main(mode)
except KeyboardInterrupt:
print('Killing the process...')
time.sleep(3.0)
finally:
print('Done.')