# /// script
# requires-python = ">=3.12"
# dependencies = [
# "matplotlib>=3.10.3",
# "np>=1.0.2",
# "nuscenes-devkit>=1.1.9",
# "pillow>=11.3.0",
# "pydantic>=2.11.5",
# "pypcd4>=1.2.1",
# "requests>=2.32.3",
# "scipy>=1.15.3",
# "tqdm>=4.67.1",
# ]
# ///
from __future__ import annotations
import argparse
import json
import os
import pathlib
import re
import shutil
import tarfile
from dataclasses import dataclass
from enum import StrEnum, auto
from math import floor
from typing import Annotated, Any, Literal
import numpy as np
import pypcd4
import requests
import tqdm
from nuscenes import nuscenes
from pydantic import BaseModel, ConfigDict, Field
from scipy.spatial.transform import Rotation
"""
This script processes scenes from the nuScenes (https://www.nuscenes.org/) dataset and converts them into a
the Encord upload JSON format for visualization and annotation. It can handle
lidar, radar, and camera data, as well as 3D annotations and ego-vehicle poses.
The script downloads the nuScenes minisplit if not found locally, and processes it, including:
- Converting the point cloud data from .bin to .pcd
- Timestamps are normalized to start from 0 at the beginning of the scene
- Converting positions so that the vehicle's starting position is treated as the origin (0, 0, 0)
"""
def snake2camel(snake: str, start_lower: bool = True) -> str:
"""
Converts a snake_case string to camelCase.
The `start_lower` argument determines whether the first letter in the generated camelcase should
be lowercase (if `start_lower` is True), or capitalized (if `start_lower` is False).
"""
camel = snake.title()
camel = re.sub("([0-9A-Za-z])_(?=[0-9A-Z])", lambda m: m.group(1), camel)
if start_lower:
camel = re.sub("(^_*[A-Z])", lambda m: m.group(1).lower(), camel)
return camel
class CamelModel(BaseModel):
model_config = ConfigDict(alias_generator=snake2camel, populate_by_name=True)
@dataclass
class CameraIntrinsics:
fx: Annotated[float, Field(description="Focal length x")]
fy: Annotated[float, Field(description="Focal length y")]
ox: Annotated[float, Field(description="Principal point offset x")]
oy: Annotated[float, Field(description="Principal point offset y")]
s: Annotated[float, Field(description="Axis skew")]
@dataclass
class CameraExtrinsics:
rotation: Annotated[
tuple[float, float, float, float, float, float, float, float, float],
Field(description="Rotation matrix R"),
]
position: Annotated[
tuple[float, float, float], Field(description="Translation vector T")
]
@dataclass
class CameraParams:
width_px: int
height_px: int
intrinsics: Annotated[CameraIntrinsics, Field(description="The intrinsic matrix K")]
extrinsics: Annotated[
CameraExtrinsics, Field(description="The extrinsic 4x4 matrix R|T")
]
@dataclass
class FrameOfReference:
id: Annotated[str, Field(description="ID of this frame of reference")]
parent_FOR: Annotated[
str | None, Field(description="ID of a parent frame of reference")
]
rotation: tuple[float, float, float, float, float, float, float, float, float]
position: tuple[float, float, float]
Position = tuple[float, float, float]
EulerOrientation = tuple[float, float, float]
Size = tuple[float, float, float]
class Pose(CamelModel):
position: Position
orientation: EulerOrientation
class CuboidGeometry(CamelModel):
type: Literal["cuboid"] = "cuboid"
pose: Pose
size: Size
@dataclass
class _FORIdMixin:
frame_of_reference_id: Annotated[
str | None, Field(description="ID of the frame of reference the entity is in")
] = None
@dataclass
class _URIMixin:
uri: str
@dataclass
class _EventMixin:
timestamp: float | None = None
class URIEvent(CamelModel, _EventMixin, _URIMixin):
pass
class CameraParamsEvent(CamelModel, _EventMixin, CameraParams):
pass
class FOREvent(CamelModel, _EventMixin, FrameOfReference):
pass
class ModelEvent(CamelModel, _EventMixin):
geometries: list[CuboidGeometry]
class CompositeScene(CamelModel):
type: Literal["composite"] = "composite"
streams: dict[str, EventStream]
class EntityType(StrEnum):
POINT_CLOUD = auto()
FRAME_OF_REFERENCE = auto()
IMAGE = auto()
MODEL = auto()
CAMERA_PARAMETERS = auto()
class PCDStream(CamelModel, _FORIdMixin):
entity_type: Literal[EntityType.POINT_CLOUD] = EntityType.POINT_CLOUD
events: Annotated[list[URIEvent], Field(description="List of point cloud events")]
class CameraStream(CamelModel, _FORIdMixin):
entity_type: Literal[EntityType.CAMERA_PARAMETERS] = EntityType.CAMERA_PARAMETERS
events: list[CameraParamsEvent]
class ImageStream(CamelModel):
entity_type: Literal[EntityType.IMAGE] = EntityType.IMAGE
events: list[URIEvent]
camera_id: Annotated[
str | None,
Field(
description="ID of the camera associated with the image. Used to position the image in-scene"
),
]
class ModelStream(CamelModel):
entity_type: Literal[EntityType.MODEL] = EntityType.MODEL
events: list[URIEvent | ModelEvent]
camera_id: str | None
class FORStream(CamelModel):
entity_type: Literal[EntityType.FRAME_OF_REFERENCE] = EntityType.FRAME_OF_REFERENCE
events: Annotated[
list[FOREvent], Field(description="List of frame of reference events")
]
class EventStream(CamelModel):
type: Literal["event"] = "event"
id: str
stream: Annotated[
PCDStream | CameraStream | FORStream | ImageStream | ModelStream,
Field(discriminator="entity_type"),
]
DATASET_DIR = pathlib.Path("./dataset")
class Config:
env: str
output_dir: pathlib.Path
base_url: str
def __init__(self):
self.env = "remote"
self.output_dir = pathlib.Path("./scenes")
self.base_url = (
"https://storage.cloud.google.com/my-bucket-name/scenes/nuscenes" # Replace this with the file path in your bucket to the dataset
)
config = Config()
def ensure_scene_available(
root_dir: pathlib.Path, dataset_version: str, scene_name: str
) -> None:
"""
Ensure that the specified scene is available.
Downloads minisplit into root_dir if scene_name is part of it and root_dir is empty.
Raises ValueError if scene is not available and cannot be downloaded.
"""
try:
nusc = nuscenes.NuScenes(
version=dataset_version, dataroot=str(root_dir), verbose=False
)
except AssertionError: # dataset initialization failed
if dataset_version == "v1.0-mini":
download_minisplit(root_dir)
nusc = nuscenes.NuScenes(
version=dataset_version, dataroot=str(root_dir), verbose=False
)
else:
print(
f"Could not find dataset at {root_dir} and could not automatically download specified scene."
)
exit()
scene_names = [s["name"] for s in nusc.scene]
if scene_name not in scene_names:
raise ValueError(f"{scene_name=} not found in dataset")
def nuscene_sensor_names(nusc: nuscenes.NuScenes, scene_name: str) -> list[str]:
"""Return all sensor names in the scene."""
sensor_names = set()
scene = next(s for s in nusc.scene if s["name"] == scene_name)
first_sample = nusc.get("sample", scene["first_sample_token"])
for sample_data_token in first_sample["data"].values():
sample_data = nusc.get("sample_data", sample_data_token)
if sample_data["sensor_modality"] == "camera":
current_camera_token = sample_data_token
while current_camera_token != "":
sample_data = nusc.get("sample_data", current_camera_token)
sensor_name = sample_data["channel"]
sensor_names.add(sensor_name)
current_camera_token = sample_data["next"]
# For a known set of cameras, order the sensors in a circle.
ordering = {
"CAM_FRONT_LEFT": 0,
"CAM_FRONT": 1,
"CAM_FRONT_RIGHT": 2,
"CAM_BACK_RIGHT": 3,
"CAM_BACK": 4,
"CAM_BACK_LEFT": 5,
}
return sorted(
sensor_names, key=lambda sensor_name: ordering.get(sensor_name, float("inf"))
)
# Write all uri assets required for the scene to a separate output directory
def write_asset(path: pathlib.Path):
shutil.copyfile(path, pathlib.Path("./output") / path.name)
def write_nuscenes_json(scene: CompositeScene, name: str):
OUTPUT_FILE = config.output_dir / "nuscenes.json"
with open(OUTPUT_FILE, "w") as f:
dummy_json = scene.model_dump_json(by_alias=True, indent=2)
f.write(dummy_json)
print("Wrote to", OUTPUT_FILE)
def write_upload_json(scenes: list[tuple[CompositeScene, str]]):
scenes_final = []
for scene, name in scenes:
streams = list(scene.model_dump(by_alias=True)["streams"].values())
scenes_final.append(
{
"title": name,
"streams": streams,
}
)
final = {"scenes": scenes_final}
OUTPUT_FILE = config.output_dir / "upload.json"
with open(OUTPUT_FILE, "w") as f:
json.dump(final, f, indent=2)
print("Wrote to", OUTPUT_FILE)
first_timestamp = 0
first_position = [0, 0, 0]
hz = 0
def sub(a, b) -> tuple[float, float, float]:
return [a[i] - b[i] for i in range(len(a))]
def log_nuscenes(
nusc: nuscenes.NuScenes, scene_name: str, max_time_sec: float, sample_hz: float
) -> CompositeScene:
"""Log nuScenes scene."""
print(f"Logging scene {scene_name}")
result = CompositeScene(streams={})
scene = next(s for s in nusc.scene if s["name"] == scene_name)
location = nusc.get("log", scene["log_token"])["location"]
# Get the first sample
first_sample_token = scene["first_sample_token"]
first_sample = nusc.get("sample", scene["first_sample_token"])
# Get the timestamp (in seconds)
global first_timestamp
first_timestamp = first_sample["timestamp"] / 1e6
global first_position
first_position = (0, 0, 0)
global hz
hz = sample_hz
first_lidar_tokens = []
first_radar_tokens = []
first_camera_tokens = []
for sample_data_token in first_sample["data"].values():
sample_data = nusc.get("sample_data", sample_data_token)
log_sensor_calibration(result, sample_data, nusc)
if sample_data["sensor_modality"] == "lidar":
first_lidar_tokens.append(sample_data_token)
elif sample_data["sensor_modality"] == "radar":
first_radar_tokens.append(sample_data_token)
elif sample_data["sensor_modality"] == "camera":
first_camera_tokens.append(sample_data_token)
first_timestamp_us = nusc.get("sample_data", first_lidar_tokens[0])["timestamp"]
max_timestamp_us = first_timestamp_us + 1e6 * max_time_sec
log_lidar_and_ego_pose(result, location, first_lidar_tokens, nusc, max_timestamp_us)
log_cameras(result, first_camera_tokens, nusc, max_timestamp_us)
log_radars(result, first_radar_tokens, nusc, max_timestamp_us)
log_annotations(result, location, first_sample_token, nusc, max_timestamp_us)
return result
def log_cameras(
scene: CompositeScene,
first_camera_tokens: list[str],
nusc: nuscenes.NuScenes,
max_timestamp_us: float,
) -> None:
"""Log camera data."""
for first_camera_token in first_camera_tokens:
current_camera_token = first_camera_token
last_logged_timestamp = -10000
while current_camera_token != "":
sample_data = nusc.get("sample_data", current_camera_token)
if max_timestamp_us < sample_data["timestamp"]:
break
sensor_name = sample_data["channel"]
if sensor_name not in scene.streams:
scene.streams[sensor_name] = EventStream(
id=sensor_name,
stream=ImageStream(
events=[],
camera_id=sensor_name + "-camera",
frame_of_reference_id=sensor_name + "-calibration",
),
)
timestamp = sample_data["timestamp"] * 1e-6 - first_timestamp
if hz > 0:
timestamp *= hz
timestamp = floor(timestamp)
if hz > 0 and timestamp - last_logged_timestamp < 1.0:
current_camera_token = sample_data["next"]
continue
last_logged_timestamp = timestamp
data_file_path = nusc.dataroot / sample_data["filename"]
# write_asset(data_file_path)
event = URIEvent(
uri=config.base_url + "/" + str(data_file_path),
timestamp=timestamp,
)
scene.streams[sensor_name].stream.events.append(event)
current_camera_token = sample_data["next"]
def log_lidar_and_ego_pose(
scene: CompositeScene,
location: str,
first_lidar_token: list[str],
nusc: nuscenes.NuScenes,
max_timestamp_us: float,
) -> None:
"""Log lidar data and vehicle pose."""
scene.streams["ego_vehicle"] = EventStream(
id="ego_vehicle",
stream=FORStream(events=[]),
)
last_logged_timestamp = -10000
for current_lidar_token in first_lidar_token:
while current_lidar_token != "":
sample_data = nusc.get("sample_data", current_lidar_token)
sensor_name = sample_data["channel"]
if max_timestamp_us < sample_data["timestamp"]:
break
timestamp = sample_data["timestamp"] * 1e-6 - first_timestamp
if hz > 0:
timestamp *= hz
timestamp = floor(timestamp)
if hz > 0 and timestamp - last_logged_timestamp < 1.0:
current_lidar_token = sample_data["next"]
continue
last_logged_timestamp = timestamp
ego_pose = nusc.get("ego_pose", sample_data["ego_pose_token"])
rotation = (
Rotation.from_quat(ego_pose["rotation"], scalar_first=True)
.as_matrix()
.transpose()
.flatten()
)
position = ego_pose["translation"]
if timestamp == 0:
global first_position
first_position = position
event = FOREvent(
id="ego_vehicle",
parent_FOR="root",
position=sub(position, first_position),
rotation=rotation,
timestamp=timestamp,
)
scene.streams["ego_vehicle"].stream.events.append(event)
current_lidar_token = sample_data["next"]
data_file_path = nusc.dataroot / sample_data["filename"]
if sensor_name not in scene.streams:
scene.streams[sensor_name] = EventStream(
id=sensor_name,
stream=PCDStream(
events=[], frame_of_reference_id=sensor_name + "-calibration"
),
)
data_file_path = nusc.dataroot / sample_data["filename"]
pointcloud = nuscenes.LidarPointCloud.from_file(str(data_file_path))
points = pointcloud.points[:3].T
fields = ("x", "y", "z")
types = (
np.float32,
np.float32,
np.float32,
)
pc = pypcd4.PointCloud.from_points(points, fields, types)
# strip .bin extension from filename
new_path = str(data_file_path.parent / data_file_path.stem)
pc.save(new_path)
event = URIEvent(
uri=config.base_url + "/" + new_path,
timestamp=timestamp,
)
scene.streams[sensor_name].stream.events.append(event)
def log_radars(
scene: CompositeScene,
first_radar_tokens: list[str],
nusc: nuscenes.NuScenes,
max_timestamp_us: float,
) -> None:
"""Log radar data to the scene"""
for first_radar_token in first_radar_tokens:
current_camera_token = first_radar_token
last_logged_timestamp = -10000
while current_camera_token != "":
sample_data = nusc.get("sample_data", current_camera_token)
if max_timestamp_us < sample_data["timestamp"]:
break
sensor_name = sample_data["channel"]
if sensor_name not in scene.streams:
scene.streams[sensor_name] = EventStream(
id=sensor_name,
stream=PCDStream(
events=[], frame_of_reference_id=sensor_name + "-calibration"
),
)
timestamp = sample_data["timestamp"] * 1e-6 - first_timestamp
if hz > 0:
timestamp *= hz
timestamp = floor(timestamp)
if hz > 0 and timestamp - last_logged_timestamp < 1.0:
current_camera_token = sample_data["next"]
continue
last_logged_timestamp = timestamp
data_file_path = nusc.dataroot / sample_data["filename"]
current_camera_token = sample_data["next"]
# write_asset(data_file_path)
event = URIEvent(
uri=config.base_url + "/" + str(data_file_path),
timestamp=timestamp,
)
scene.streams[sensor_name].stream.events.append(event)
def log_sensor_calibration(
scene: CompositeScene, sample_data: dict[str, Any], nusc: nuscenes.NuScenes
) -> None:
"""Log sensor calibration (pinhole camera, sensor poses, etc.) to the scene"""
sensor_name = sample_data["channel"]
calibrated_sensor_token = sample_data["calibrated_sensor_token"]
calibrated_sensor = nusc.get("calibrated_sensor", calibrated_sensor_token)
rotation = (
Rotation.from_quat(calibrated_sensor["rotation"], scalar_first=True)
.as_matrix()
.transpose()
.flatten()
.tolist()
)
id = sensor_name + "-calibration"
scene.streams[id] = EventStream(
id=id,
stream=FORStream(events=[]),
)
position = sub(calibrated_sensor["translation"], first_position)
event = FOREvent(
id=id,
parent_FOR="ego_vehicle", # "ego_vehicle",
position=position,
rotation=rotation,
)
scene.streams[id].stream.events.append(event)
if len(calibrated_sensor["camera_intrinsic"]) != 0:
intrinsic = calibrated_sensor["camera_intrinsic"]
camera_id = sensor_name + "-camera"
scene.streams[camera_id] = EventStream(
id=camera_id,
stream=CameraStream(
events=[],
frame_of_reference_id=id, # might be "root"
),
)
event = CameraParamsEvent(
timestamp=0,
width_px=1600,
height_px=900,
intrinsics=CameraIntrinsics(
fx=intrinsic[0][0],
fy=intrinsic[1][1],
ox=intrinsic[0][2],
oy=intrinsic[1][2],
s=intrinsic[0][1],
),
extrinsics=CameraExtrinsics(
position=(0, 0, 0),
rotation=(0, 0, 1, -1, 0, 0, 0, -1, 0),
),
)
scene.streams[camera_id].stream.events.append(event)
def log_annotations(
scene: CompositeScene,
location: str,
first_sample_token: str,
nusc: nuscenes.NuScenes,
max_timestamp_us: float,
) -> None:
"""Log 3D cuboids to the scene"""
scene.streams["anns"] = EventStream(
id="anns",
stream=ModelStream(events=[], camera_id=None),
)
current_sample_token = first_sample_token
last_logged_timestamp = -10000
while current_sample_token != "":
sample_data = nusc.get("sample", current_sample_token)
if max_timestamp_us < sample_data["timestamp"]:
break
timestamp = sample_data["timestamp"] * 1e-6 - first_timestamp
if hz > 0:
timestamp *= hz
timestamp = floor(timestamp)
if hz > 0 and timestamp - last_logged_timestamp < 1.0:
current_sample_token = sample_data["next"]
continue
last_logged_timestamp = timestamp
ann_tokens = sample_data["anns"]
geometries = []
for ann_token in ann_tokens:
ann = nusc.get("sample_annotation", ann_token)
width, length, height = ann["size"]
# Convert rotation to euler angles
rotation = Rotation.from_quat(ann["rotation"], scalar_first=True).as_euler(
"XYZ"
)
geometries.append(
CuboidGeometry(
pose=Pose(
position=sub(ann["translation"], first_position),
orientation=rotation,
),
size=(length, width, height),
)
)
event = ModelEvent(
timestamp=timestamp,
geometries=geometries,
)
scene.streams["anns"].stream.events.append(event)
current_sample_token = sample_data["next"]
def download_file(url: str, dst_file_path: pathlib.Path) -> None:
"""Download file from url to dst_fpath."""
dst_file_path.parent.mkdir(parents=True, exist_ok=True)
print(f"Downloading {url} to {dst_file_path}")
response = requests.get(url, stream=True)
with tqdm.tqdm.wrapattr(
open(dst_file_path, "wb"),
"write",
miniters=1,
total=int(response.headers.get("content-length", 0)),
desc=f"Downloading {dst_file_path.name}",
) as f:
for chunk in response.iter_content(chunk_size=4096):
f.write(chunk)
def untar_file(
tar_file_path: pathlib.Path, dst_path: pathlib.Path, keep_tar: bool = True
) -> bool:
"""Untar tar file at tar_file_path to dst."""
print(f"Untar file {tar_file_path}")
try:
with tarfile.open(tar_file_path, "r") as tf:
tf.extractall(dst_path)
except Exception as error:
print(f"Error unzipping {tar_file_path}, error: {error}")
return False
if not keep_tar:
os.remove(tar_file_path)
return True
def download_minisplit(root_dir: pathlib.Path) -> None:
"""
Download nuScenes minisplit.
Adopted from <https://colab.research.google.com/github/nutonomy/nuscenes-devkit/blob/master/python-sdk/tutorials/nuscenes_tutorial.ipynb>
"""
MINISPLIT_URL = "https://www.nuscenes.org/data/v1.0-mini.tgz"
zip_file_path = pathlib.Path("./v1.0-mini.tgz")
if not zip_file_path.is_file():
download_file(MINISPLIT_URL, zip_file_path)
untar_file(zip_file_path, root_dir, keep_tar=True)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Visualizes the nuScenes dataset ")
parser.add_argument(
"--root-dir",
type=pathlib.Path,
default=DATASET_DIR,
help="Root directory of nuScenes dataset",
)
parser.add_argument(
"--scene-name",
type=str,
default="scene-0061",
help="Scene name to visualize (typically of form 'scene-xxxx')",
)
parser.add_argument(
"--dataset-version", type=str, default="v1.0-mini", help="Scene id to visualize"
)
parser.add_argument(
"--seconds",
type=float,
default=float("inf"),
help="If specified, limits the number of seconds logged",
)
parser.add_argument(
"--all",
"-A",
action="store_true",
help="If specified, logs all scenes",
)
parser.add_argument(
"--hz",
type=float,
default=0.0,
help="Limit the sample rate",
)
args = parser.parse_args()
# ensure_scene_available(
# root_dir=args.root_dir,
# dataset_version=args.dataset_version,
# scene_name=args.scene_name,
# )
nusc = nuscenes.NuScenes(
version=args.dataset_version, dataroot=args.root_dir, verbose=False
)
scene_names: list[str] = [args.scene_name]
if args.all:
scene_names = [s["name"] for s in nusc.scene]
scenes = [
(
log_nuscenes(
nusc, scene_name, max_time_sec=args.seconds, sample_hz=args.hz
),
scene_name,
)
for scene_name in scene_names
]
write_upload_json(scenes)
