├── scene_reconstruction ├── __init__.py ├── cli │ ├── __init__.py │ ├── main.py │ ├── eval.py │ ├── config.py │ └── export.py ├── data │ ├── __init__.py │ └── nuscenes │ │ ├── __init__.py │ │ ├── scene_utils.py │ │ ├── io.py │ │ ├── typing.py │ │ ├── polars_helpers.py │ │ ├── parse.py │ │ ├── scene_flow.py │ │ └── dataset.py ├── math │ ├── __init__.py │ ├── spherical_coordinate_system.py │ └── dempster_shafer.py ├── occupancy │ ├── __init__.py │ ├── render_camera.py │ ├── temporal_accumulation.py │ ├── grid.py │ ├── transmission_reflection.py │ └── temporal_transmission_and_reflection.py ├── visualization │ ├── __init__.py │ ├── open3d │ │ ├── __init__.py │ │ ├── pointcloud.py │ │ └── volume.py │ ├── plotly │ │ ├── __init__.py │ │ └── volume.py │ ├── utils │ │ ├── __init__.py │ │ └── mesh.py │ ├── streamlit │ │ ├── __init__.py │ │ └── occupancy_gt.py │ ├── export │ │ └── image.py │ └── colormap.py ├── core │ ├── __init__.py │ ├── config.py │ ├── icp.py │ ├── transform.py │ └── volume.py └── eval │ └── lidar_depth.py ├── .gitattributes ├── .gitignore ├── pyproject.toml ├── conf ├── surround_occ.yaml ├── open_occupancy.yaml ├── bba.yaml ├── bba04.yaml ├── scene_as_occupancy.yaml ├── occ3d.yaml └── default.yaml ├── pixi.toml ├── README.md └── LICENSE /scene_reconstruction/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /scene_reconstruction/cli/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /scene_reconstruction/data/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /scene_reconstruction/math/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /scene_reconstruction/occupancy/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /scene_reconstruction/data/nuscenes/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /scene_reconstruction/visualization/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /scene_reconstruction/visualization/open3d/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /scene_reconstruction/visualization/plotly/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /scene_reconstruction/visualization/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /scene_reconstruction/visualization/streamlit/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /.gitattributes: -------------------------------------------------------------------------------- 1 | # GitHub syntax highlighting 2 | pixi.lock linguist-language=YAML linguist-generated=true 3 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | 2 | # pixi environments 3 | .pixi 4 | *.egg-info 5 | 6 | # data 7 | /data 8 | 9 | # workdir 10 | /workdir 11 | 12 | # python cache files 13 | __pycache__/ -------------------------------------------------------------------------------- /scene_reconstruction/core/__init__.py: -------------------------------------------------------------------------------- 1 | """Core functionality.""" 2 | from .transform import einsum_transform 3 | from .volume import Volume 4 | 5 | __all__ = ["Volume", "einsum_transform"] 6 | -------------------------------------------------------------------------------- /scene_reconstruction/data/nuscenes/scene_utils.py: -------------------------------------------------------------------------------- 1 | """Scene utils.""" 2 | import polars as pl 3 | 4 | from scene_reconstruction.data.nuscenes.polars_helpers import series_to_torch 5 | 6 | 7 | def scene_to_tensor_dict(scene: pl.DataFrame, mapping: dict[str, str], **kwargs): 8 | """Dataframe to dict of tensors.""" 9 | return {k: series_to_torch(scene[v]).to(**kwargs) for k, v in mapping.items()} 10 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "scene-reconstruction" 3 | description = 'Implementation of "Accurate Training Data for Occupancy Map Prediction in Automated Driving Using Evidence Theory" (CVPR 2024)' 4 | version = "0.1.0" 5 | readme = "README.md" 6 | license = {text = "AGPL-3.0"} 7 | requires-python = ">=3.9" 8 | 9 | [build-system] 10 | requires = ["setuptools"] 11 | build-backend = "setuptools.build_meta" 12 | 13 | [tool.setuptools] 14 | packages = ["scene_reconstruction"] -------------------------------------------------------------------------------- /scene_reconstruction/cli/main.py: -------------------------------------------------------------------------------- 1 | """CLI entry point.""" 2 | 3 | import typer 4 | 5 | from scene_reconstruction.cli import eval, export 6 | 7 | app = typer.Typer( 8 | name="cli", 9 | help="CLI entry point.", 10 | no_args_is_help=True, 11 | add_completion=False, 12 | pretty_exceptions_show_locals=False, 13 | ) 14 | app.add_typer(typer_instance=export.app) 15 | app.add_typer(typer_instance=eval.app) 16 | 17 | 18 | if __name__ == "__main__": 19 | app() 20 | -------------------------------------------------------------------------------- /scene_reconstruction/core/config.py: -------------------------------------------------------------------------------- 1 | """Configuration helpers.""" 2 | 3 | from pathlib import Path 4 | from typing import Optional 5 | 6 | from hydra import compose, initialize_config_dir 7 | from omegaconf import DictConfig 8 | 9 | 10 | def initialize_config(cfg_dir: Path, cfg_name: str, overrides: Optional[list[str]] = None) -> DictConfig: 11 | """Initialize and return hydra config.""" 12 | 13 | with initialize_config_dir(config_dir=str(cfg_dir.resolve()), version_base=None): 14 | return compose(config_name=cfg_name, overrides=overrides) 15 | -------------------------------------------------------------------------------- /conf/surround_occ.yaml: -------------------------------------------------------------------------------- 1 | eval: 2 | render_lidar_depth: 3 | _target_: scene_reconstruction.eval.lidar_depth.LidarDistanceEval 4 | ds: 5 | _target_: scene_reconstruction.data.nuscenes.dataset.NuscenesDataset 6 | data_root: data/nuscenes 7 | extra_data_root: data/nuscenes_extra 8 | version: v1.0-mini 9 | key_frames_only: true 10 | lower: [-50.0, -50.0, -5.0] 11 | upper: [50.0, 50.0, 3.0] 12 | volume_frame: lidar 13 | eval_ego_lower: [-40.0, -40.0, -1.0] 14 | eval_ego_upper: [40.0, 40.0, 5.4] 15 | min_distance: 2.5 16 | batch_size: 1 17 | method: surround_occ 18 | split: trainval 19 | save_path: workdir/surround_occ 20 | -------------------------------------------------------------------------------- /conf/open_occupancy.yaml: -------------------------------------------------------------------------------- 1 | eval: 2 | render_lidar_depth: 3 | _target_: scene_reconstruction.eval.lidar_depth.LidarDistanceEval 4 | ds: 5 | _target_: scene_reconstruction.data.nuscenes.dataset.NuscenesDataset 6 | data_root: data/nuscenes 7 | extra_data_root: data/nuscenes_extra 8 | version: v1.0-mini 9 | key_frames_only: true 10 | lower: [-51.2, -51.2, -5.0] 11 | upper: [51.2, 51.2, 3.0] 12 | volume_frame: lidar 13 | eval_ego_lower: [-40.0, -40.0, -1.0] 14 | eval_ego_upper: [40.0, 40.0, 5.4] 15 | min_distance: 2.5 16 | batch_size: 1 17 | method: open_occupancy 18 | split: trainval 19 | save_path: workdir/open_occupancy 20 | -------------------------------------------------------------------------------- /conf/bba.yaml: -------------------------------------------------------------------------------- 1 | eval: 2 | render_lidar_depth: 3 | _target_: scene_reconstruction.eval.lidar_depth.LidarDistanceEval 4 | ds: 5 | _target_: scene_reconstruction.data.nuscenes.dataset.NuscenesDataset 6 | data_root: data/nuscenes 7 | extra_data_root: data/nuscenes_extra 8 | version: v1.0-mini 9 | key_frames_only: true 10 | lower: [-40.0, -40.0, -1.0] 11 | upper: [40.0, 40.0, 5.4] 12 | volume_frame: ego 13 | eval_ego_lower: [-40.0, -40.0, -1.0] 14 | eval_ego_upper: [40.0, 40.0, 5.4] 15 | min_distance: 2.5 16 | batch_size: 1 17 | p_fn: 0.8 18 | p_fp: 0.2 19 | method: bba 20 | split: trainval 21 | save_path: workdir/bba 22 | -------------------------------------------------------------------------------- /conf/bba04.yaml: -------------------------------------------------------------------------------- 1 | eval: 2 | render_lidar_depth: 3 | _target_: scene_reconstruction.eval.lidar_depth.LidarDistanceEval 4 | ds: 5 | _target_: scene_reconstruction.data.nuscenes.dataset.NuscenesDataset 6 | data_root: data/nuscenes 7 | extra_data_root: data/nuscenes_extra 8 | version: v1.0-mini 9 | key_frames_only: true 10 | lower: [-40.0, -40.0, -1.0] 11 | upper: [40.0, 40.0, 5.4] 12 | volume_frame: ego 13 | eval_ego_lower: [-40.0, -40.0, -1.0] 14 | eval_ego_upper: [40.0, 40.0, 5.4] 15 | min_distance: 2.5 16 | batch_size: 1 17 | p_fn: 0.9 18 | p_fp: 0.1 19 | method: bba04 20 | split: trainval 21 | save_path: workdir/bba04 22 | -------------------------------------------------------------------------------- /conf/scene_as_occupancy.yaml: -------------------------------------------------------------------------------- 1 | eval: 2 | render_lidar_depth: 3 | _target_: scene_reconstruction.eval.lidar_depth.LidarDistanceEval 4 | ds: 5 | _target_: scene_reconstruction.data.nuscenes.dataset.NuscenesDataset 6 | data_root: data/nuscenes 7 | extra_data_root: data/nuscenes_extra 8 | version: v1.0-mini 9 | key_frames_only: true 10 | lower: [-50.0, -50.0, -5.0] 11 | upper: [50.0, 50.0, 3.0] 12 | volume_frame: lidar 13 | eval_ego_lower: [-40.0, -40.0, -1.0] 14 | eval_ego_upper: [40.0, 40.0, 5.4] 15 | min_distance: 2.5 16 | batch_size: 1 17 | method: scene_as_occupancy 18 | split: trainval 19 | save_path: workdir/scene_as_occupancy 20 | -------------------------------------------------------------------------------- /conf/occ3d.yaml: -------------------------------------------------------------------------------- 1 | eval: 2 | render_lidar_depth: 3 | _target_: scene_reconstruction.eval.lidar_depth.LidarDistanceEval 4 | ds: 5 | _target_: scene_reconstruction.data.nuscenes.dataset.NuscenesDataset 6 | data_root: data/nuscenes 7 | extra_data_root: data/nuscenes_extra 8 | version: v1.0-mini 9 | key_frames_only: true 10 | lower: [-40.0, -40.0, -1.0] 11 | upper: [40.0, 40.0, 5.4] 12 | volume_frame: ego 13 | eval_ego_lower: [-40.0, -40.0, -1.0] 14 | eval_ego_upper: [40.0, 40.0, 5.4] 15 | min_distance: 2.5 16 | batch_size: 1 17 | p_fn: 0.9 18 | p_fp: 0.05 19 | method: cvpr2023 # occ3d 20 | split: trainval 21 | save_path: workdir/cvpr2023 22 | -------------------------------------------------------------------------------- /scene_reconstruction/visualization/export/image.py: -------------------------------------------------------------------------------- 1 | """Export visualizations.""" 2 | import matplotlib.pyplot as plt 3 | 4 | from scene_reconstruction.visualization.colormap import turbo_black 5 | 6 | 7 | def image_with_colormap(data, filename, cmap=turbo_black, background=(1.0, 1.0, 1.0)): 8 | """Export heatmap encoded with colormap as image.""" 9 | px = 1 / plt.rcParams["figure.dpi"] # pixel in inches 10 | fig, ax = plt.subplots(1, figsize=(data.shape[1] * px, data.shape[0] * px), facecolor=background) 11 | fig.subplots_adjust(left=0, right=1, bottom=0, top=1) 12 | im = ax.imshow(data, cmap=cmap, vmin=0.0, vmax=1.0) 13 | ax.axis("off") 14 | fig.savefig(filename, bbox_inches="tight", pad_inches=0, transparent=True) 15 | plt.close(fig) 16 | -------------------------------------------------------------------------------- /scene_reconstruction/cli/eval.py: -------------------------------------------------------------------------------- 1 | """Commands for data export.""" 2 | 3 | 4 | import typer 5 | from hydra.utils import instantiate 6 | 7 | from scene_reconstruction.cli.config import make_cfg 8 | from scene_reconstruction.eval.lidar_depth import LidarDistanceEval 9 | 10 | app = typer.Typer(name="eval", callback=make_cfg, help="Various export commands.", no_args_is_help=True) 11 | 12 | 13 | @app.command(name="render-lidar-depth") 14 | def lidar_depth(ctx: typer.Context) -> None: 15 | """Eval rendered lidar depth.""" 16 | cfg = ctx.meta["cfg"] 17 | 18 | transmission_and_reflections: LidarDistanceEval = instantiate(cfg.eval.render_lidar_depth) 19 | results = transmission_and_reflections.eval() 20 | print("Config yaml:") 21 | print(cfg.eval) 22 | for k, v in results.items(): 23 | print(f"{k}: {v}") 24 | -------------------------------------------------------------------------------- /scene_reconstruction/cli/config.py: -------------------------------------------------------------------------------- 1 | """Common constants.""" 2 | 3 | from __future__ import annotations 4 | 5 | from pathlib import Path 6 | from typing import Annotated, Optional 7 | 8 | import typer 9 | 10 | from ..core.config import initialize_config 11 | 12 | CFG_NAME = typer.Option(help="Configuration name.") 13 | CFG_DIR = typer.Option(envvar="SR_CONFIG_DIR", dir_okay=True, readable=True, help="Configuration root directory.") 14 | CFG_OVERRIDES = typer.Argument(show_default=False, help="Additional configuration overrides.") 15 | 16 | OptionalOverrides = Optional[list[str]] 17 | 18 | 19 | def make_cfg( 20 | ctx: typer.Context, 21 | cfg_dir: Annotated[Path, CFG_DIR], 22 | cfg_name: Annotated[str, CFG_NAME], 23 | cfg_overrides: Annotated[OptionalOverrides, CFG_OVERRIDES] = None, 24 | ): 25 | """Create configuration.""" 26 | ctx.meta["cfg"] = initialize_config(cfg_dir=cfg_dir, cfg_name=cfg_name, overrides=cfg_overrides) 27 | -------------------------------------------------------------------------------- /scene_reconstruction/core/icp.py: -------------------------------------------------------------------------------- 1 | """ICP using open3d.""" 2 | 3 | import open3d as o3d 4 | import torch 5 | from torch import Tensor 6 | 7 | 8 | def torch_to_pcl(points: Tensor): 9 | """Torch tensor of shape [N, 3] to open3d pointcloud.""" 10 | pcl = o3d.geometry.PointCloud() 11 | pcl.points = o3d.utility.Vector3dVector(points.cpu().numpy()) 12 | return pcl 13 | 14 | 15 | def register_frame(ref_points: Tensor, new_points: Tensor): 16 | """ICP registration using open3d.""" 17 | ref_from_guess = torch.as_tensor( 18 | o3d.pipelines.registration.registration_icp( 19 | torch_to_pcl(new_points), 20 | torch_to_pcl(ref_points), 21 | max_correspondence_distance=0.2, 22 | criteria=o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=20), 23 | ).transformation, 24 | dtype=torch.float, 25 | ) 26 | return ref_from_guess.to(device=ref_points.device) 27 | -------------------------------------------------------------------------------- /scene_reconstruction/visualization/open3d/pointcloud.py: -------------------------------------------------------------------------------- 1 | """Open3d pointcloud visualization.""" 2 | 3 | from typing import Optional 4 | 5 | import numpy as np 6 | import open3d as o3d 7 | from torch import Tensor 8 | 9 | 10 | def open3d_pcl_from_torch(points: Tensor, color: Optional[Tensor] = None): 11 | """Colored open3d pointcloud form torch.""" 12 | vec = o3d.utility.Vector3dVector(points.cpu().numpy()) 13 | pcl = o3d.geometry.PointCloud(vec) 14 | if color is not None: 15 | pcl.colors = o3d.utility.Vector3dVector(color.cpu().numpy()) 16 | return pcl 17 | 18 | 19 | def show_pointcloud(points: np.ndarray, color: Optional[np.ndarray] = None): 20 | """Colored pointcloud visualization.""" 21 | vec = o3d.utility.Vector3dVector(points) 22 | pcl = o3d.geometry.PointCloud(vec) 23 | if color is not None: 24 | pcl.colors = o3d.utility.Vector3dVector(color) 25 | 26 | o3d.visualization.draw_geometries([pcl]) # pylint: disable=E1101 27 | -------------------------------------------------------------------------------- /scene_reconstruction/data/nuscenes/io.py: -------------------------------------------------------------------------------- 1 | """Data input / output from files.""" 2 | 3 | from typing import Optional 4 | 5 | import numpy as np 6 | import torch 7 | 8 | 9 | def load_lidar_from_file(filename: str, pad_length: Optional[int] = None): 10 | """Loads LIDAR data from binary numpy format. Data is stored as (x, y, z, intensity, ring index).""" 11 | dims_to_load = [0, 1, 2] # [x, y, z] 12 | scan = np.fromfile(filename, dtype=np.float32) 13 | points = scan.reshape((-1, 5))[:, dims_to_load] 14 | if pad_length is None: 15 | return points 16 | points_padded = np.empty((pad_length, 3), dtype=points.dtype) 17 | num_points = len(points) 18 | points_padded[:num_points] = points 19 | points_padded[num_points:] = float("nan") 20 | return points_padded 21 | 22 | 23 | def load_occupancy_from_file(filename: str) -> dict[str, torch.Tensor]: 24 | """Loads occupancy data from file.""" 25 | data = np.load(file=filename) 26 | data = {k: torch.as_tensor(v) for k, v in data.items()} 27 | return data 28 | -------------------------------------------------------------------------------- /scene_reconstruction/visualization/streamlit/occupancy_gt.py: -------------------------------------------------------------------------------- 1 | """Visualize occupancy.""" 2 | 3 | import streamlit as st 4 | 5 | from scene_reconstruction.core.volume import Volume 6 | from scene_reconstruction.data.nuscenes.io import load_occupancy_from_file 7 | from scene_reconstruction.visualization.plotly.volume import plotly_voxel_volume 8 | 9 | st.set_page_config(layout="wide") 10 | 11 | st.sidebar.title("Occupancy") 12 | 13 | file = st.sidebar.file_uploader("Choose a file", type="npz") 14 | 15 | use_camera_mask = st.sidebar.toggle("Camera mask") 16 | use_lidar_mask = st.sidebar.toggle("Lidar mask") 17 | show_free = st.sidebar.toggle("Show free") 18 | hide_visible_freespace = st.sidebar.toggle("Hide visible free") 19 | if file is not None: 20 | occ_data = load_occupancy_from_file(file) 21 | lower = [-40, -40.0, -1.0] 22 | upper = [40.0, 40.0, 5.4] 23 | volume = Volume.new_volume(lower, upper) 24 | semantics = occ_data["semantics"][None, None] 25 | mask_camera = occ_data["mask_camera"][None, None].bool() 26 | mask_lidar = occ_data["mask_lidar"][None, None].bool() 27 | if not show_free: 28 | semantics[semantics == 17] = -1 29 | if use_camera_mask: 30 | semantics[~mask_camera] = -1 31 | if use_lidar_mask: 32 | semantics[~mask_lidar] = -1 33 | if hide_visible_freespace: 34 | semantics[mask_camera & (semantics == 17)] = -1 35 | fig = plotly_voxel_volume(semantics, volume, colormap="turbo_r", scale=(0, 17)) 36 | fig.update_layout(height=900) 37 | st.plotly_chart(fig, use_container_width=True) 38 | -------------------------------------------------------------------------------- /scene_reconstruction/visualization/colormap.py: -------------------------------------------------------------------------------- 1 | """Colormaps.""" 2 | 3 | import numpy as np 4 | import torch 5 | from matplotlib import colormaps 6 | from matplotlib.colors import ListedColormap 7 | 8 | occupancy_to_color = { 9 | "ignore_class": (0, 0, 0), # Black. 10 | "barrier": (112, 128, 144), # Slategrey 11 | "bicycle": (220, 20, 60), # Crimson 12 | "bus": (255, 127, 80), # Coral 13 | "car": (255, 158, 0), # Orange 14 | "construction_vehicle": (233, 150, 70), # Darksalmon 15 | "motorcycle": (255, 61, 99), # Red 16 | "pedestrian": (0, 0, 230), # Blue 17 | "traffic_cone": (47, 79, 79), # Darkslategrey 18 | "trailer": (255, 140, 0), # Darkorange 19 | "truck": (255, 99, 71), # Tomato 20 | "driveable_surface": (0, 207, 191), # nuTonomy green 21 | "other_flat": (175, 0, 75), 22 | "sidewalk": (75, 0, 75), 23 | "terrain": (112, 180, 60), 24 | "manmade": (222, 184, 135), # Burlywood 25 | "vegetation": (0, 175, 0), 26 | "free": (52, 107, 235), 27 | } 28 | occupancy_color_map = torch.tensor(list(occupancy_to_color.values())) 29 | 30 | 31 | def _create_turbo_alpha(): 32 | newcolors = colormaps["turbo"](np.linspace(0, 1, 256)) 33 | newcolors[0, -1] = 0.0 34 | turbo_alpha = ListedColormap(newcolors) 35 | return turbo_alpha 36 | 37 | 38 | turbo_alpha = _create_turbo_alpha() 39 | 40 | 41 | def _create_turbo_black(): 42 | newcolors = colormaps["turbo"](np.linspace(0, 1, 256)) 43 | newcolors[0, :3] = 0.0 44 | turbo_alpha = ListedColormap(newcolors) 45 | return turbo_alpha 46 | 47 | 48 | turbo_black = _create_turbo_black() 49 | -------------------------------------------------------------------------------- /scene_reconstruction/data/nuscenes/typing.py: -------------------------------------------------------------------------------- 1 | """NuScenes typing.""" 2 | 3 | from typing import Literal, Union 4 | 5 | # pylint: disable=C0103 6 | 7 | CAMERA_CHANNELS = Literal[ 8 | "CAM_FRONT", 9 | "CAM_FRONT_RIGHT", 10 | "CAM_BACK_RIGHT", 11 | "CAM_BACK", 12 | "CAM_BACK_LEFT", 13 | "CAM_FRONT_LEFT", 14 | ] 15 | LIDAR_CHANNELS = Literal["LIDAR_TOP"] 16 | RADAR_CHANNELS = Literal[ 17 | "RADAR_FRONT", 18 | "RADAR_FRONT_LEFT", 19 | "RADAR_FRONT_RIGHT", 20 | "RADAR_BACK_LEFT", 21 | "RADAR_BACK_RIGHT", 22 | ] 23 | SENSOR_CHANNELS = Union[CAMERA_CHANNELS, LIDAR_CHANNELS, RADAR_CHANNELS] 24 | 25 | SENSOR_MODALITY = Literal["camera", "lidar", "radar"] 26 | CATEGOTY_NAMES = Literal[ 27 | "noise", 28 | "animal", 29 | "human.pedestrian.adult", 30 | "human.pedestrian.child", 31 | "human.pedestrian.construction_worker", 32 | "human.pedestrian.personal_mobility", 33 | "human.pedestrian.police_officer", 34 | "human.pedestrian.stroller", 35 | "human.pedestrian.wheelchair", 36 | "movable_object.barrier", 37 | "movable_object.debris", 38 | "movable_object.pushable_pullable", 39 | "movable_object.trafficcone", 40 | "static_object.bicycle_rack", 41 | "vehicle.bicycle", 42 | "vehicle.bus.bendy", 43 | "vehicle.bus.rigid", 44 | "vehicle.car", 45 | "vehicle.construction", 46 | "vehicle.emergency.ambulance", 47 | "vehicle.emergency.police", 48 | "vehicle.motorcycle", 49 | "vehicle.trailer", 50 | "vehicle.truck", 51 | "flat.driveable_surface", 52 | "flat.other", 53 | "flat.sidewalk", 54 | "flat.terrain", 55 | "static.manmade", 56 | "static.other", 57 | "static.vegetation", 58 | "vehicle.ego", 59 | ] 60 | -------------------------------------------------------------------------------- /scene_reconstruction/cli/export.py: -------------------------------------------------------------------------------- 1 | """Commands for data export.""" 2 | 3 | 4 | import typer 5 | from hydra.utils import instantiate 6 | 7 | from scene_reconstruction.cli.config import make_cfg 8 | from scene_reconstruction.data.nuscenes.scene_flow import SceneFlow 9 | from scene_reconstruction.occupancy.temporal_transmission_and_reflection import TemporalTransmissionAndReflection 10 | from scene_reconstruction.occupancy.transmission_reflection import ReflectionTransmissionSpherical 11 | 12 | app = typer.Typer(name="export", callback=make_cfg, help="Various export commands.", no_args_is_help=True) 13 | 14 | SAVE_DIR = typer.Option(help="Directory to save data to.", dir_okay=True) 15 | BATCH_SIZE = typer.Option(help="Batch size for data processing.") 16 | 17 | 18 | @app.command(name="transmissions-reflections") 19 | def transmissions_reflections(ctx: typer.Context) -> None: 20 | """Export sensor count maps to specified path.""" 21 | cfg = ctx.meta["cfg"] 22 | 23 | transmission_and_reflections: ReflectionTransmissionSpherical = instantiate(cfg.export.transmissions_reflections) 24 | transmission_and_reflections.process_data() 25 | 26 | 27 | REF_KEYFRAME_ONLY = typer.Option(help="Only accumulate for reference keyframes.") 28 | 29 | 30 | @app.command(name="temporal-accumulation") 31 | def temporal_transmissions_reflections( 32 | ctx: typer.Context, 33 | ) -> None: 34 | """Accumulate sensor count maps over time to specified path.""" 35 | cfg = ctx.meta["cfg"] 36 | 37 | temporal_accumulation: TemporalTransmissionAndReflection = instantiate(cfg.export.temporal_accumulation) 38 | temporal_accumulation.process_data() 39 | 40 | 41 | @app.command(name="scene-flow") 42 | def scene_flow(ctx: typer.Context) -> None: 43 | """Accumulate sensor count maps over time to specified path.""" 44 | cfg = ctx.meta["cfg"] 45 | 46 | scene_flow: SceneFlow = instantiate(cfg.export.scene_flow) 47 | scene_flow.process_data() 48 | 49 | 50 | @app.command(name="sensor-belief-maps", no_args_is_help=True) 51 | def sensor_belief_maps(ctx: typer.Context) -> None: 52 | """Export sensor belief maps to specified path.""" 53 | -------------------------------------------------------------------------------- /scene_reconstruction/occupancy/render_camera.py: -------------------------------------------------------------------------------- 1 | """Camera rendering.""" 2 | 3 | from collections.abc import Sequence 4 | 5 | import torch 6 | from torch import Tensor 7 | 8 | from scene_reconstruction.core import Volume 9 | from scene_reconstruction.core.transform import einsum_transform 10 | 11 | 12 | def depth_centers(min_depth: float, max_depth: float, num_bins: int): 13 | """Depth bins center points.""" 14 | return ((torch.arange(num_bins) + 0.5) / num_bins) * (max_depth - min_depth) + min_depth 15 | 16 | 17 | def camera_sample_grid( 18 | camera_volume_shape: Sequence[int], camera_from_image: Tensor, min_depth, max_depth, img_width, img_height 19 | ): 20 | """Create sample grid from camera intrinsics.""" 21 | cam_lower = torch.stack([torch.zeros_like(img_width), torch.zeros_like(img_height), torch.ones_like(img_height)], 1) 22 | cam_upper = torch.stack([img_width, img_height, torch.ones_like(img_height)], 1) 23 | img_volume = Volume(lower=cam_lower, upper=cam_upper) 24 | 25 | img_grid = img_volume.new_coord_grid([camera_volume_shape[0], camera_volume_shape[1], 1]) 26 | img_grid_camera_no_depth = einsum_transform("bci,bxyzi->bxyzc", camera_from_image, points=img_grid) 27 | img_grid_direction = img_grid_camera_no_depth / img_grid_camera_no_depth.norm(dim=-1, keepdim=True) 28 | img_grid_camera = img_grid_direction * depth_centers(min_depth, max_depth, camera_volume_shape[2]).view( 29 | 1, 1, 1, -1, 1 30 | ) 31 | return img_grid_camera 32 | 33 | 34 | def sample_camera_rays( 35 | ego_features: Tensor, 36 | volume_ego: Volume, 37 | ego_from_camera: Tensor, 38 | camera_from_image: Tensor, 39 | camera_volume_shape: Sequence[int], 40 | min_depth, 41 | max_depth, 42 | img_width, 43 | img_height, 44 | ): 45 | """Sample camera rays in ego volume.""" 46 | img_grid_camera = camera_sample_grid( 47 | camera_volume_shape=camera_volume_shape, 48 | camera_from_image=camera_from_image, 49 | min_depth=min_depth, 50 | max_depth=max_depth, 51 | img_width=img_width, 52 | img_height=img_height, 53 | ) 54 | img_grid_ego = einsum_transform("bec,bxyzc->bxyze", ego_from_camera, points=img_grid_camera) 55 | # TODO: rescaling based on volume 56 | samped_features = volume_ego.sample_volume(ego_features, img_grid_ego) 57 | return samped_features 58 | -------------------------------------------------------------------------------- /scene_reconstruction/math/spherical_coordinate_system.py: -------------------------------------------------------------------------------- 1 | """Math related to spherical coordinates.""" 2 | 3 | import torch 4 | from torch import Tensor 5 | 6 | # https://en.wikipedia.org/wiki/Spherical_coordinate_system#Coordinate_system_conversions 7 | 8 | 9 | def cartesian_to_spherical(xyz: Tensor): 10 | """Transforms cartestian coordinates into spherical coordinates.""" 11 | # pylint: disable=C0103 12 | x, y, z = xyz.unbind(-1) 13 | r = xyz.norm(dim=-1) # radius >= 0 14 | theta = torch.arccos(z / r) # 0 <= polar angle <= pi 15 | phi = torch.atan2(y, x) # -pi <= azimuth <= pi 16 | spherical_coords = torch.stack([r, theta, phi], -1) 17 | return spherical_coords 18 | 19 | 20 | def spherical_to_cartesian(spherical_coords: Tensor): 21 | """Transforms spherical coordinates into cartestian coordinates.""" 22 | # pylint: disable=C0103 23 | r, theta, phi = spherical_coords.unbind(-1) 24 | sin_theta = theta.sin() 25 | cos_theta = theta.cos() 26 | sin_phi = phi.sin() 27 | cos_phi = phi.cos() 28 | r_sin_theta = r * sin_theta 29 | x = r_sin_theta * cos_phi 30 | y = r_sin_theta * sin_phi 31 | z = r * cos_theta 32 | xyz = torch.stack([x, y, z], -1) 33 | return xyz 34 | 35 | 36 | def spherical_volume_element(spherical_lower: Tensor, spherical_upper: Tensor): 37 | """Calculates to volume of a spherical volume element specified by its lower and upper bound.""" 38 | r_lower, theta_lower, phi_lower = spherical_lower.unbind(-1) 39 | r_upper, theta_upper, phi_upper = spherical_upper.unbind(-1) 40 | # https://en.wikipedia.org/wiki/Multiple_integral#Spherical_coordinates 41 | volume = (phi_upper - phi_lower) * (-theta_upper.cos() + theta_lower.cos()) * (r_upper**3 - r_lower**3) / 3.0 42 | return volume 43 | 44 | 45 | def spherical_volume_element_center_and_voxel_size(spherical_center: Tensor, spherical_voxel_size: Tensor): 46 | """Calculates to volume of a spherical volume element specified by its center and voxel size.""" 47 | r, theta, phi = spherical_center.unbind(-1) 48 | dr, dtheta, dphi = spherical_voxel_size.unbind(-1) 49 | 50 | # https://en.wikipedia.org/wiki/Multiple_integral#Spherical_coordinates 51 | volume = ( 52 | dphi 53 | * (-(theta + 0.5 * dtheta).cos() + (theta - 0.5 * dtheta).cos()) 54 | * ((r + 0.5 * dr) ** 3 - (r - 0.5 * dr) ** 3) 55 | / 3.0 56 | ) 57 | return volume 58 | -------------------------------------------------------------------------------- /scene_reconstruction/core/transform.py: -------------------------------------------------------------------------------- 1 | """Homogenoues transformation related functions.""" 2 | 3 | import logging 4 | 5 | import torch 6 | from torch import Tensor 7 | 8 | 9 | def transform_volume_bounds( 10 | lower: Tensor, 11 | upper: Tensor, 12 | new_lower: Tensor, 13 | new_upper: Tensor, 14 | verbose: bool = False, 15 | ): 16 | """Homogenous transformation matrix between volumes specified by lower and upper bounds.""" 17 | # xyz = ((xyz - lower) / (upper - lower)) * (new_upper - new_lower) + new_lower 18 | scale = (new_upper - new_lower) / (upper - lower) 19 | 20 | if verbose and not scale.allclose(scale[..., 0:1]): 21 | logging.getLogger(__file__).warning(f"Scale is not equal across all axis: {scale}") 22 | scale_homogenous = torch.nn.functional.pad(scale, (0, 1), "constant", 1.0) # pylint: disable=E1102 23 | offset = new_lower - lower * scale 24 | transform = torch.diag_embed(scale_homogenous) 25 | transform[..., :3, 3] = offset 26 | return transform 27 | 28 | 29 | def to_homogenous(xyz: Tensor): 30 | """Converts coordinates into homogenous coordinates.""" 31 | return torch.nn.functional.pad(xyz, (0, 1), "constant", 1.0) # pylint: disable=E1102 32 | 33 | 34 | def from_homogenous(xyzh: Tensor): 35 | """Converts homogenous coordinates into normal coordinates.""" 36 | return xyzh[..., :3] 37 | 38 | 39 | def transform_to_grid_sample_coords(lower, upper): 40 | """Converts from bounded volume to coordintates which can be used with grid_sample.""" 41 | # Grid sample uses normalized coordinates between -1.0 and 1.0 42 | # Addionally xyz needs to be reversed to zyx 43 | normalization_transform = transform_volume_bounds(lower, upper, -torch.ones_like(lower), torch.ones_like(upper)) 44 | swap_transform = torch.eye(4, device=lower.device, dtype=lower.dtype)[[2, 1, 0, 3]] 45 | transform = swap_transform @ normalization_transform 46 | return transform 47 | 48 | 49 | def einsum_transform(einsum_str: str, *transform: Tensor, points: Tensor): 50 | """Applies a transformation to points using the einsum notation.""" 51 | is_homogenous = points.shape[-1] == 4 52 | if not is_homogenous: 53 | points = to_homogenous(points) 54 | assert points.shape[-1] == 4 55 | points = torch.einsum(einsum_str, *transform, points) 56 | if not is_homogenous: 57 | points = from_homogenous(points) 58 | return points 59 | -------------------------------------------------------------------------------- /scene_reconstruction/occupancy/temporal_accumulation.py: -------------------------------------------------------------------------------- 1 | """Accumulation of BBA.""" 2 | import torch 3 | from torch import Tensor 4 | 5 | from scene_reconstruction.core import einsum_transform 6 | from scene_reconstruction.core.volume import Volume 7 | from scene_reconstruction.math.dempster_shafer import yager_rule_of_combination_stacked 8 | 9 | 10 | def recursive_accumulation(m: Tensor, volume: Volume, ego_from_global: torch.Tensor, global_from_ego: torch.Tensor): 11 | """Recursive yager rule.""" 12 | volume = volume.cuda() 13 | grid_current_frame = volume.new_coord_grid(volume.volume_shape(m)) 14 | m_out = torch.empty_like(m) 15 | m_out[0] = m[0] 16 | last_from_current = (ego_from_global[:-1] @ global_from_ego[1:]).cuda() 17 | for i in range(1, len(m)): 18 | grid_last_frame = einsum_transform("blc,bxyzc->bxyzl", last_from_current[i - 1 : i], points=grid_current_frame) 19 | m_sampled = volume.sample_volume(m_out[i - 1 : i].cuda(), grid_last_frame.cuda()) 20 | m_out[i : i + 1] = yager_rule_of_combination_stacked(m_sampled, m[i : i + 1].cuda()).cpu() 21 | return m_out 22 | 23 | 24 | def recursive_accumulation_backward(m, volume: Volume, ego_from_global: torch.Tensor, global_from_ego: torch.Tensor): 25 | """Recursive yager rule backwards.""" 26 | volume = volume.cuda() 27 | grid_current_frame = volume.new_coord_grid(volume.volume_shape(m)) 28 | m_out = torch.empty_like(m) 29 | m_out[-1] = m[-1] 30 | next_from_curr = (ego_from_global[1:] @ global_from_ego[:-1]).cuda() 31 | for i in range(len(m) - 1, 0, -1): 32 | grid_next_frame = einsum_transform("blc,bxyzc->bxyzl", next_from_curr[i - 1 : i], points=grid_current_frame) 33 | m_sampled = volume.sample_volume(m_out[i : i + 1].cuda(), grid_next_frame.cuda()) 34 | m_out[i - 1 : i] = yager_rule_of_combination_stacked(m_sampled, m[i - 1 : i].cuda()).cpu() 35 | return m_out 36 | 37 | 38 | def forward_backward_accumulation( 39 | m: Tensor, volume: Volume, ego_from_global: torch.Tensor, global_from_ego: torch.Tensor 40 | ): 41 | """Recursive yager rule forward and backwards.""" 42 | m_forward = recursive_accumulation(m, volume, ego_from_global=ego_from_global, global_from_ego=global_from_ego) 43 | m_backward = recursive_accumulation_backward( 44 | m, volume, ego_from_global=ego_from_global, global_from_ego=global_from_ego 45 | ) 46 | m_forward_backward = yager_rule_of_combination_stacked(m_forward, m_backward) 47 | return m_forward_backward 48 | -------------------------------------------------------------------------------- /conf/default.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | nuscenes: 3 | _target_: scene_reconstruction.data.nuscenes.dataset.NuscenesDataset 4 | data_root: data/nuscenes 5 | version: v1.0-mini 6 | key_frames_only: false 7 | 8 | volume: 9 | ego_lower: [-40.0, -40.0, -1.0] 10 | ego_upper: [40.0, 40.0, 5.4] 11 | ego_shape: [400, 400, 32] 12 | spherical_lower: [2.5, 1.3089969389957472, -3.141592653589793] #[0.0, (90 - 15) / 180 * math.pi, -math.pi] 13 | spherical_upper: [60.0, 2.1816615649929116, 3.141592653589793] #[0.0, (90 + 35) / 180 * math.pi, -math.pi] 14 | spherical_shape: [600, 100, 720] # voxel size [0.1m, 0.5°, 0.5°] 15 | lidar_min_distance: 2.5 16 | 17 | export: 18 | transmissions_reflections: 19 | _target_: scene_reconstruction.occupancy.transmission_reflection.ReflectionTransmissionSpherical 20 | ds: 21 | _target_: scene_reconstruction.data.nuscenes.dataset.NuscenesDataset 22 | data_root: data/nuscenes 23 | extra_data_root: data/nuscenes_extra 24 | version: v1.0-mini 25 | key_frames_only: false 26 | 27 | extra_data_root: data/nuscenes_extra 28 | spherical_lower: [2.5, 1.3089969389957472, -3.141592653589793] #[0.0, (90 - 15) / 180 * math.pi, -math.pi] 29 | spherical_upper: [60.0, 2.1816615649929116, 3.141592653589793] #[0.0, (90 + 35) / 180 * math.pi, -math.pi] 30 | spherical_shape: [600, 100, 720] # voxel size [0.1m, 0.5°, 0.5°] 31 | lidar_min_distance: 2.5 # meters 32 | # voxel_size cartesian = 0.2m 33 | # voxel_size sperical : 0.1m, 0.5°, 0.5° 34 | batch_size: 4 35 | 36 | scene_flow: 37 | _target_: scene_reconstruction.data.nuscenes.scene_flow.SceneFlow 38 | ds: 39 | _target_: scene_reconstruction.data.nuscenes.dataset.NuscenesDataset 40 | data_root: data/nuscenes 41 | extra_data_root: data/nuscenes_extra 42 | version: v1.0-mini 43 | key_frames_only: false 44 | 45 | extra_data_root: data/nuscenes_extra 46 | cartesian_lower: [-40.0, -40.0, -1.0] 47 | cartesian_upper: [40.0, 40.0, 5.4] 48 | cartesian_shape: [400, 400, 32] 49 | 50 | temporal_accumulation: 51 | _target_: scene_reconstruction.occupancy.temporal_transmission_and_reflection.TemporalTransmissionAndReflection 52 | ds: 53 | _target_: scene_reconstruction.data.nuscenes.dataset.NuscenesDataset 54 | data_root: data/nuscenes 55 | extra_data_root: data/nuscenes_extra 56 | version: v1.0-mini 57 | key_frames_only: false 58 | 59 | extra_data_root: data/nuscenes_extra 60 | frame_accumulation_kwargs: 61 | icp_alignment: false 62 | batch_size: 1 63 | max_num_frames: 50 64 | max_ego_pose_difference: 20.0 65 | device: cuda 66 | num_threads: 8 67 | -------------------------------------------------------------------------------- /pixi.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "scene-reconstruction" 3 | authors = ["Jonas Kälble "] 4 | description = 'Implementation for "Accurate Training Data for Occupancy Map Prediction in Automated Driving Using Evidence Theory" (CVPR 2024)' 5 | platforms = ["linux-64"] 6 | version = "0.1.0" 7 | channels = [ 8 | "nvidia/label/cuda-11.6.2", 9 | "nvidia/label/cuda-11.6.1", 10 | "nvidia/label/cuda-11.6.0", 11 | "conda-forge", 12 | "pytorch", 13 | ] 14 | readme = "README.md" 15 | license = "AGPL-3.0" 16 | 17 | [tasks] 18 | # Preprocessing 19 | transmissions-reflections = "python -m scene_reconstruction.cli.main export ./conf default transmissions-reflections" 20 | scene-flow = "python -m scene_reconstruction.cli.main export ./conf default scene-flow" 21 | temporal-accumulation = "python -m scene_reconstruction.cli.main export ./conf default temporal-accumulation" 22 | data-processing = { depends-on = [ 23 | "transmissions-reflections", 24 | "scene-flow", 25 | "temporal-accumulation", 26 | ] } 27 | # Evaluation 28 | eval-bba04 = "python -m scene_reconstruction.cli.main eval ./conf bba04 render-lidar-depth" 29 | eval-bba = "python -m scene_reconstruction.cli.main eval ./conf bba render-lidar-depth" 30 | eval-occ3d = "python -m scene_reconstruction.cli.main eval ./conf occ3d render-lidar-depth" 31 | eval-open-occupancy = "python -m scene_reconstruction.cli.main eval ./conf open_occupancy render-lidar-depth" 32 | eval-surround-occ = "python -m scene_reconstruction.cli.main eval ./conf surround_occ render-lidar-depth" 33 | eval-scene-as-occupancy = "python -m scene_reconstruction.cli.main eval ./conf scene_as_occupancy render-lidar-depth" 34 | eval-all = { depends-on = [ 35 | "eval-bba04", 36 | "eval-bba", 37 | "eval-occ3d", 38 | "eval-open-occupancy", 39 | "eval-surround-occ", 40 | "eval-scene-as-occupancy", 41 | ] } 42 | 43 | [dependencies] 44 | python = "3.9.*" 45 | binutils = "*" 46 | cxx-compiler = "*" 47 | git = "*" 48 | gxx = "<12" 49 | libcusolver-dev = "*" 50 | ninja = "*" 51 | numba = "*" 52 | numpy = "<1.24" 53 | opencv = "*" 54 | pip = "*" 55 | cuda = { version = "11.6.*", channel = "nvidia/label/cuda-11.6.0" } 56 | pytorch = { version = "1.12.*", channel = "pytorch" } 57 | pytorch-cuda = { version = "11.6.*", channel = "pytorch" } 58 | torchvision = { version = "*", channel = "pytorch" } 59 | mkl = "2020.0.*" 60 | scipy = "*" 61 | 62 | [pypi-dependencies] 63 | scene-reconstruction = { path = ".", editable = true } 64 | einops = "*" 65 | hydra-core = "*" 66 | ipympl = "*" 67 | jupyter = "*" 68 | notebook = "*" 69 | open3d = "*" 70 | polars = "*" 71 | rich = "*" 72 | spconv-cu116 = "*" 73 | typer = "*" 74 | pyarrow = "*" 75 | streamlit = "*" 76 | seaborn = "*" 77 | torchmetrics = "*" 78 | nuscenes-devkit = "*" 79 | kaolin = "~=0.14" 80 | 81 | [pypi-options] 82 | index-url = "https://pypi.org/simple" 83 | find-links = [ 84 | { url = "https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-1.12.0_cu116.html" }, 85 | ] 86 | -------------------------------------------------------------------------------- /scene_reconstruction/visualization/utils/mesh.py: -------------------------------------------------------------------------------- 1 | """Mesh.""" 2 | 3 | import torch 4 | from torch import Tensor 5 | 6 | CUBE_CORNERS = ( 7 | torch.tensor( 8 | [ 9 | [0.0, 0.0, 0.0], 10 | [1.0, 0.0, 0.0], 11 | [0.0, 0.0, 1.0], 12 | [1.0, 0.0, 1.0], 13 | [0.0, 1.0, 0.0], 14 | [1.0, 1.0, 0.0], 15 | [0.0, 1.0, 1.0], 16 | [1.0, 1.0, 1.0], 17 | ], 18 | dtype=torch.float, 19 | ) 20 | - 0.5 21 | ) # [8, 3] 22 | NUM_CUBE_CORNERS = len(CUBE_CORNERS) # 8 23 | CUBE_TRIANGLES = torch.tensor( 24 | [ 25 | [4, 7, 5], 26 | [4, 6, 7], 27 | [0, 2, 4], 28 | [2, 6, 4], 29 | [0, 1, 2], 30 | [1, 3, 2], 31 | [1, 5, 7], 32 | [1, 7, 3], 33 | [2, 3, 7], 34 | [2, 7, 6], 35 | [0, 4, 1], 36 | [1, 4, 5], 37 | ], 38 | dtype=torch.int32, 39 | ) # [12, 3] 40 | CUBE_TRIANGLE_NORMALS = torch.tensor( 41 | [ 42 | [0.0, 1.0, 0.0], 43 | [0.0, 1.0, 0.0], 44 | [-1.0, 0.0, 0.0], 45 | [-1.0, 0.0, 0.0], 46 | [0.0, -1.0, 0.0], 47 | [0.0, -1.0, 0.0], 48 | [1.0, 0.0, 0.0], 49 | [1.0, 0.0, 0.0], 50 | [0.0, 0.0, 1.0], 51 | [0.0, 0.0, 1.0], 52 | [0.0, 0.0, -1.0], 53 | [0.0, 0.0, -1.0], 54 | ], 55 | dtype=torch.float, 56 | ) # [12, 3] 57 | NUM_CUBE_TRIANGLES = len(CUBE_TRIANGLES) # 12 58 | 59 | 60 | def cube_mesh( 61 | cube_centers: Tensor, 62 | cube_sizes: Tensor, 63 | ): 64 | """Generate vertices and triangles from cube centers and sizes.""" 65 | # cube_centers: [N, 3] 66 | # cube_size: [N, 3], [N, 1] 67 | assert cube_centers.ndim == 2 and cube_centers.shape[1] == 3 68 | cube_sizes = cube_sizes.expand_as(cube_centers) 69 | num_cubes = cube_centers.shape[0] 70 | cube_corners = CUBE_CORNERS.clone() 71 | if cube_sizes is not None: 72 | cube_corners = cube_sizes[:, None, :] * cube_corners 73 | vertices = cube_centers[:, None, :] + cube_corners # [N, CUBE_CORNERS, 3] 74 | 75 | triangles = NUM_CUBE_CORNERS * torch.arange(num_cubes)[:, None, None] + CUBE_TRIANGLES # [N, CUBE_TRIANGLES, 3] 76 | 77 | vertices_flat = vertices.reshape((-1, 3)) # [N * CUBE_CORNERS, 3] 78 | triangles_flat = triangles.reshape((-1, 3)) # [N * NUM_CUBE_CORNERS, 3] 79 | 80 | return vertices_flat, triangles_flat 81 | 82 | 83 | def cube_mesh_colored( 84 | cube_centers: Tensor, 85 | cube_sizes: Tensor, 86 | cube_colors: Tensor, 87 | ): 88 | """Generate vertices, triangles and vertex colors from cube centers, sizes and colors.""" 89 | # cube_centers: [N, 3] 90 | # cube_size: [N, 3], [N, 1] 91 | # cube_colors: [N, 3], [N, 1] 92 | vertices_flat, triangles_flat = cube_mesh(cube_centers, cube_sizes) 93 | color_channels = cube_colors.shape[-1] 94 | vertex_colors = cube_colors[:, None, :].expand(-1, NUM_CUBE_CORNERS, -1).reshape((-1, color_channels)) 95 | assert vertices_flat.shape[0] == vertex_colors.shape[0] 96 | return vertices_flat, triangles_flat, vertex_colors 97 | -------------------------------------------------------------------------------- /scene_reconstruction/visualization/open3d/volume.py: -------------------------------------------------------------------------------- 1 | """Open3d dense grids.""" 2 | 3 | import matplotlib as mpl 4 | import numpy as np 5 | import open3d as o3d 6 | from torch import Tensor 7 | 8 | from scene_reconstruction.core import Volume 9 | 10 | from ..utils.mesh import cube_mesh_colored 11 | 12 | 13 | def open3d_mesh_from_volume( 14 | features: Tensor, 15 | volume: Volume, 16 | batch_index: int = 0, 17 | channel_index: int = 0, 18 | colormap="turbo", 19 | scale: tuple[float, float] = (0.0, 1.0), 20 | legacy: bool = True, 21 | ): 22 | """Generate open3d mesh from volume.""" 23 | centers = volume.coord_grid(features)[batch_index] # [X, Y, Z, 3] 24 | features = features[batch_index, channel_index] # [X, Y, Z] 25 | if features.is_floating_point(): 26 | mask = ~features.isnan() 27 | else: 28 | mask = features != -1 29 | features = (features - scale[0]) / (scale[1] - scale[0]) 30 | centers_flat = centers[mask] # [N, 3] 31 | features_flat = features[mask] 32 | voxel_size = volume.voxel_size(features) 33 | if voxel_size.shape[0] != 1: 34 | voxel_size = voxel_size[batch_index] 35 | vertices_flat, triangles_flat, vertex_color = cube_mesh_colored( 36 | cube_centers=centers_flat, cube_sizes=voxel_size, cube_colors=features_flat[:, None] 37 | ) 38 | vertex_colors = mpl.colormaps[colormap](vertex_color.numpy()[:, 0])[:, :3].astype(np.float32) 39 | if legacy: 40 | vertices = o3d.utility.Vector3dVector(vertices_flat.numpy()) 41 | triangles = o3d.utility.Vector3iVector(triangles_flat.numpy()) 42 | 43 | colors = o3d.utility.Vector3dVector(vertex_colors) 44 | mesh = o3d.geometry.TriangleMesh(vertices, triangles) 45 | mesh.vertex_colors = colors 46 | mesh.compute_triangle_normals() 47 | 48 | else: 49 | mesh = o3d.t.geometry.TriangleMesh() 50 | mesh.vertex.positions = o3d.core.Tensor(vertices_flat.numpy()) 51 | mesh.triangle.indices = o3d.core.Tensor(triangles_flat.numpy()) 52 | mesh.vertex.colors = o3d.core.Tensor(vertex_colors) 53 | mesh.compute_triangle_normals() 54 | mesh.material = o3d.visualization.Material("defaultLit") 55 | # mesh.material.scalar_properties["roughness"] = 0.8 56 | # mesh.material.scalar_properties["reflectance"] = 0.3 57 | # mesh.material.scalar_properties["transmission"] = 0.0 58 | # mesh.material.scalar_properties["thickness"] = 0.3 59 | # mesh.material.scalar_properties["absorption_distance"] = 0.1 60 | # mesh.material.vector_properties['absorption_color'] = np.array([1.0, 1.0, 1.0, 1.0]) 61 | return mesh 62 | 63 | 64 | def open3d_voxel_volume( 65 | features: Tensor, 66 | volume: Volume, 67 | batch_index: int = 0, 68 | channel_index: int = 0, 69 | colormap="turbo", 70 | scale: tuple[float, float] = (0.0, 1.0), 71 | legacy: bool = True, 72 | ): 73 | """Plot volumes using open3d.""" 74 | mesh = open3d_mesh_from_volume( 75 | features, 76 | volume, 77 | batch_index, 78 | channel_index, 79 | colormap, 80 | scale, 81 | legacy=legacy, 82 | ) 83 | if legacy: 84 | o3d.visualization.draw_geometries([mesh]) # pylint: disable=E1101 85 | else: 86 | geometries = [{"name": "volume", "geometry": mesh}] 87 | o3d.visualization.draw(geometries, up=[0.0, 0.0, 1.0]) 88 | -------------------------------------------------------------------------------- /scene_reconstruction/visualization/plotly/volume.py: -------------------------------------------------------------------------------- 1 | """Plotly volumes.""" 2 | 3 | import plotly.graph_objects as go 4 | from torch import Tensor 5 | 6 | from scene_reconstruction.core.volume import Volume 7 | 8 | from ..utils.mesh import cube_mesh_colored 9 | 10 | 11 | def plotly_mesh_from_volume( 12 | features: Tensor, 13 | volume: Volume, 14 | batch_index: int = 0, 15 | channel_index: int = 0, 16 | colormap="turbo", 17 | scale: tuple[float, float] = (0.0, 1.0), 18 | ): 19 | """Generate plotly mesh from volume.""" 20 | centers = volume.coord_grid(features)[batch_index] # [X, Y, Z, 3] 21 | features = features[batch_index, channel_index] # [X, Y, Z] 22 | if features.is_floating_point(): 23 | mask = ~features.isnan() 24 | else: 25 | mask = features != -1 26 | features = (features - scale[0]) / (scale[1] - scale[0]) 27 | centers_flat = centers[mask] # [N, 3] 28 | features_flat = features[mask] 29 | voxel_size = volume.voxel_size(features) 30 | vertices_flat, triangles_flat, vertex_color = cube_mesh_colored( 31 | cube_centers=centers_flat, cube_sizes=voxel_size, cube_colors=features_flat[:, None] 32 | ) 33 | upper = volume.upper.expand(features.shape[0], -1) 34 | mesh = go.Mesh3d( 35 | # vertices 36 | x=vertices_flat[:, 0].numpy(), 37 | y=vertices_flat[:, 1].numpy(), 38 | z=vertices_flat[:, 2].numpy(), 39 | # triangles 40 | i=triangles_flat[:, 0].numpy(), 41 | j=triangles_flat[:, 1].numpy(), 42 | k=triangles_flat[:, 2].numpy(), 43 | # color 44 | cmin=0.0, 45 | cmax=1.0, 46 | # vertexcolor=vertex_color[:, 0].numpy() 47 | intensity=vertex_color[:, 0].numpy(), 48 | colorscale=colormap, 49 | # lightning 50 | flatshading=True, 51 | lighting={ 52 | "ambient": 0.4, 53 | "diffuse": 1, 54 | "fresnel": 0.1, 55 | "specular": 0.1, 56 | "roughness": 0.9, 57 | "vertexnormalsepsilon": 0, 58 | "facenormalsepsilon": 0, 59 | }, 60 | lightposition={ 61 | "x": upper[batch_index, 0].item(), 62 | "y": upper[batch_index, 1].item(), 63 | "z": upper[batch_index, 2].item(), 64 | }, 65 | # 66 | hoverinfo="skip", 67 | ) 68 | return mesh 69 | 70 | 71 | def plotly_voxel_volume( 72 | features: Tensor, 73 | volume: Volume, 74 | batch_index: int = 0, 75 | channel_index: int = 0, 76 | colormap="turbo", 77 | scale: tuple[float, float] = (0.0, 1.0), 78 | show_grid: bool = False, 79 | equal_axis: bool = True, 80 | ): 81 | """Plot volumes using plotly.""" 82 | fig = go.Figure( 83 | data=[ 84 | plotly_mesh_from_volume( 85 | features, 86 | volume, 87 | batch_index=batch_index, 88 | channel_index=channel_index, 89 | colormap=colormap, 90 | scale=scale, 91 | ) 92 | ], 93 | layout={ 94 | "scene": { 95 | "aspectmode": "data" if equal_axis else "auto", 96 | "xaxis": {"visible": show_grid}, 97 | "yaxis": {"visible": show_grid}, 98 | "zaxis": {"visible": show_grid}, 99 | } 100 | }, 101 | ) 102 | return fig 103 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Evidential Occupancy 2 | Official implementation for "Accurate Training Data for Occupancy Map Prediction in Automated Driving Using Evidence Theory" (CVPR 2024, [arXiv](https://arxiv.org/pdf/2405.10575)). 3 | 4 | 5 | # Installation 6 | We use [Pixi](https://pixi.sh/latest/) to manage the dependencies. To install the dependencies, run the following command: 7 | ``` 8 | pixi install 9 | ``` 10 | This will create the Python environment with all necessary dependencies.. 11 | To run a shell with the activated environment run `pixi shell` or to start a single task, use `pixi run python ...`. 12 | 13 | 14 | # Data preparation 15 | 16 | To use the default configuration, sturcture the data as follows: 17 | ``` 18 | ./data/ 19 | ├── nuscenes/ # (original nuScenes dataset goes here (required)) 20 | │ ├── samples/ 21 | │ ├── scenes/ 22 | │ ├── sweeps/ 23 | │ └── ... 24 | └── nuscenes_extra/ # (additional data) 25 | ├── nuscenes_occ3d/ # (Occ3D (optional)) 26 | │ └── gts/ 27 | │ ├── scene-* 28 | │ └── ... 29 | ├── nuScenes-Occupancy-v0.1/ # (OpenOccupancy (optional)) 30 | │ ├── scene_* 31 | │ └── ... 32 | ├── occ_gt_release_v1_0/ # (Scene as Occupancy (optional)) 33 | │ └── trainval/ 34 | │ ├── scene-* 35 | │ └── ... 36 | └── surround_occ_occupancy/ # (SurroundOcc (optional)) 37 | └── samples/ 38 | └── ... 39 | ``` 40 | 41 | 42 | # Data Processing 43 | 44 | We use the YAML file `./conf/default.yaml` to configure the processing steps. 45 | For testing purposes, the nuScenes mini dataset is selected by default. 46 | To use the entire nuScenes dataset, set `version` to `v1.0-train`, `v1.0-val` or `v1.0-trainval` (entire dataset) and set the appropriate split (`train`,`val` or `trainval`). 47 | 48 | The data processing requires multiple steps. 49 | The necessary task are defined in the `pixi.toml`. 50 | You can run the entire pipeline with the following command: 51 | ``` 52 | pixi run data-processing 53 | ``` 54 | This will create the following file structure: 55 | ``` 56 | ./data/ 57 | └── nuscenes_extra/ 58 | ├── reflection_and_transmission_spherical/ 59 | │ ├── scene-* 60 | │ └── ... 61 | ├── scene_flow/ 62 | │ ├── scene-* 63 | │ └── ... 64 | ├── sample_annotation_cache/ 65 | │ └── ... 66 | └── reflection_and_transmission_multi_frame/ 67 | ├── scene-* 68 | └── ... 69 | ``` 70 | 71 | 72 | **OR** 73 | 74 | run the following commands step by step: 75 | 76 | ## Relfections and Transmissions 77 | Run the following command to calculate the reflections and transmissions: 78 | ``` 79 | pixi run transmissions-reflections 80 | ``` 81 | 82 | ## Scene Flow 83 | Run the following command to calculate the scene flow information: 84 | ``` 85 | pixi run scene-flow 86 | ``` 87 | 88 | ## Temporal Accumulation 89 | Run the following command to generate the accumulated reflections and transmissions: 90 | ``` 91 | pixi run temporal-accumulation 92 | ``` 93 | 94 | # Evaluation 95 | The different occupancy methods can be evaluated separately using the following commands: 96 | ``` 97 | pixi run eval-bba 98 | pixi run eval-bba04 99 | pixi run eval-occ3d 100 | pixi run eval-open-occupancy 101 | pixi run eval-scene-as-occupancy 102 | pixi run eval-surround-occ 103 | ``` 104 | Note: This requires the corresponding data to be processed / setup up properly. 105 | 106 | To evaluate all methods at once run: 107 | ``` 108 | pixi run eval-all 109 | ``` 110 | 111 | The results are stored in the `./workdir` directory. -------------------------------------------------------------------------------- /scene_reconstruction/math/dempster_shafer.py: -------------------------------------------------------------------------------- 1 | """Dempster Shafer theory.""" 2 | import itertools 3 | import math 4 | 5 | import einops 6 | import torch 7 | from torch import Tensor 8 | 9 | 10 | def yager_rule_of_combination_stacked(m1: Tensor, m2: Tensor, dim=1): 11 | """Yager's rule of combination with stacked BBA.""" 12 | m_omega1 = 1.0 - m1.sum(dim, keepdim=True) 13 | m_omega2 = 1.0 - m2.sum(dim, keepdim=True) 14 | m = m1 * m2 + m1 * m_omega2 + m_omega1 * m2 15 | return m 16 | 17 | 18 | def yager_rule_of_combination_single_element(m_of: Tensor): 19 | """Yager's rule of combination with stacked BBA.""" 20 | # m_of: 2 B 21 | m_o, m_f = m_of.unbind(0) # B 22 | m_omega = 1.0 - m_o - m_f 23 | m_o = torch.cartesian_prod(*torch.stack([m_o, m_omega]).unbind(1)).prod(-1).sum() - m_omega.prod() 24 | m_f = torch.cartesian_prod(*torch.stack([m_f, m_omega]).unbind(1)).prod(-1).sum() - m_omega.prod() 25 | m_of = torch.stack([m_o, m_f], -1) 26 | return m_of 27 | 28 | 29 | def yager_rule_of_combination_across_batch(m: Tensor): 30 | """Yager's rule of combination over first dimension with stacked BBA.""" 31 | # [B, 2, X, Y, Z] 32 | _, _, X, Y, Z = m.shape 33 | m_flat = einops.rearrange(m, "b m x y z -> (x y z) m b") 34 | m_combined_flat = torch.vmap(yager_rule_of_combination_single_element)(m_flat) 35 | m_combined = einops.repeat(m_combined_flat, "(x y z) m -> 1 m x y z", x=X, y=Y, z=Z) 36 | return m_combined 37 | 38 | 39 | def yager_rule_of_combination_across_batch_iterative(m: Tensor, channel_dim=1, combination_dim: int = 0): 40 | """Yager's rule of combination over first dimension with stacked BBA.""" 41 | # [B, 2, X, Y, Z] 42 | if m.shape[combination_dim] == 1: 43 | return m 44 | omega = 1.0 - m.sum(channel_dim, keepdim=True) 45 | out = torch.zeros_like(m.select(combination_dim, 0)) # [2, X, Y, Z] 46 | 47 | # iterate over all combinations of m and omega 48 | for first, *combination in itertools.product(*zip(m.unbind(combination_dim), omega.unbind(combination_dim))): 49 | # and add them up 50 | out += math.prod(combination, start=first) 51 | # correct for all omega combination 52 | omega_first, *omega_other = omega.unbind(combination_dim) 53 | out -= math.prod(omega_other, start=omega_first) 54 | return out 55 | 56 | 57 | def dempster_rule_of_combination(m_o1, m_o2, m_f1, m_f2): 58 | """Dempster's rule of combination.""" 59 | m_omega1 = 1 - (m_o1 + m_f1) 60 | m_omega2 = 1 - (m_o2 + m_f2) 61 | conflicts = m_o1 * m_f2 + m_f1 * m_o2 62 | m_o = (m_o1 * m_o2 + m_o1 * m_omega2 + m_omega1 * m_o2) / (1 - conflicts).clamp_min(1e-8) 63 | m_f = (m_f1 * m_f2 + m_f1 * m_omega2 + m_omega1 * m_f2) / (1 - conflicts).clamp_min(1e-8) 64 | return m_o, m_f 65 | 66 | 67 | def yager_rule_of_combination(m_o1, m_o2, m_f1, m_f2): 68 | """Yager's rule of combination.""" 69 | m_omega1 = 1 - (m_o1 + m_f1) 70 | m_omega2 = 1 - (m_o2 + m_f2) 71 | m_o = m_o1 * m_o2 + m_o1 * m_omega2 + m_omega1 * m_o2 72 | m_f = m_f1 * m_f2 + m_f1 * m_omega2 + m_omega1 * m_f2 73 | return m_o, m_f 74 | 75 | 76 | def belief_from_reflection_and_transmission(num_reflections, num_transmissions, p_fn=0.8, p_fp=0.1): 77 | """Basic belief assignement from number of reflections and transmissions.""" 78 | # p_fn = 0.9 # occupied with transmission 79 | # p_fp = 0.6 # empty with reflection 80 | # https://arxiv.org/pdf/1801.05297.pdf 81 | 82 | m_o = p_fn**num_transmissions * (1.0 - p_fp**num_reflections) 83 | m_f = p_fp**num_reflections * (1.0 - p_fn**num_transmissions) 84 | return torch.cat([m_o, m_f], 1) 85 | 86 | 87 | def belief_from_reflection_and_transmission_stacked(num_rt: Tensor, p_fn=0.8, p_fp=0.05, with_omega: bool = False): 88 | """Basic belief assignement from number of reflections and transmissions.""" 89 | # p_fn = 0.9 # occupied with transmission 90 | # p_fp = 0.6 # empty with reflection 91 | # https://arxiv.org/pdf/1801.05297.pdf 92 | num_reflections, num_transmissions = num_rt.split(1, 1) 93 | m_o = p_fn**num_transmissions * (1.0 - p_fp**num_reflections) 94 | m_f = p_fp**num_reflections * (1.0 - p_fn**num_transmissions) 95 | if with_omega: 96 | m_omega = 1.0 - m_o - m_f 97 | return torch.cat([m_o, m_f, m_omega], 1) 98 | return torch.cat([m_o, m_f], 1) 99 | -------------------------------------------------------------------------------- /scene_reconstruction/core/volume.py: -------------------------------------------------------------------------------- 1 | """Dense 3D Volumes.""" 2 | 3 | from collections.abc import Sequence 4 | from dataclasses import dataclass 5 | 6 | import torch 7 | from torch import Tensor 8 | 9 | from .transform import einsum_transform, transform_to_grid_sample_coords, transform_volume_bounds 10 | 11 | 12 | @dataclass 13 | class Volume: 14 | """Indexing for 3D volume of shape [B, C, X, Y, Z].""" 15 | 16 | lower: Tensor # [B, 3] 17 | upper: Tensor # [B, 3] 18 | 19 | @classmethod 20 | def new_volume(cls, lower: Sequence[float], upper: Sequence[float], **kwargs): 21 | """New volume from lower and upper bound.""" 22 | assert len(lower) == 3 and len(upper) == 3 23 | return cls( 24 | lower=torch.tensor([lower], **kwargs), 25 | upper=torch.tensor([upper], **kwargs), 26 | ) 27 | 28 | @classmethod 29 | def new_normalized(cls, **kwargs): 30 | """New volume from -1.0, to 1.0.""" 31 | return cls.new_volume(lower=[-1.0, -1.0, -1.0], upper=[1.0, 1.0, 1.0], **kwargs) 32 | 33 | @classmethod 34 | def new_index(cls, volume_shape: Sequence[int], **kwargs): 35 | """New volume from 0 to volume_shape.""" 36 | return cls.new_volume([0, 0, 0], volume_shape, **kwargs) 37 | 38 | @staticmethod 39 | def volume_shape(features: Tensor): 40 | """Volume shape of features.""" 41 | return features.shape[-3:] 42 | 43 | @property 44 | def device(self) -> torch.device: 45 | """Device of data.""" 46 | assert self.lower.device == self.upper.device 47 | return self.lower.device 48 | 49 | def to(self, *args, **kwargs) -> "Volume": 50 | """Convert dtype and/or device.""" 51 | return Volume(lower=self.lower.to(*args, **kwargs), upper=self.upper.to(*args, **kwargs)) 52 | 53 | def cuda(self, *args, **kwargs) -> "Volume": 54 | """Move to cuda device.""" 55 | return Volume(lower=self.lower.cuda(*args, **kwargs), upper=self.upper.cuda(*args, **kwargs)) 56 | 57 | def voxel_size(self, features: Tensor) -> Tensor: 58 | """Voxel size.""" 59 | return (self.upper - self.lower).abs() / torch.tensor(self.volume_shape(features), device=self.device) 60 | 61 | def voxel_size_from_shape(self, volume_shape: Sequence[int]) -> Tensor: 62 | """Voxel size for given volume shape.""" 63 | return (self.upper - self.lower).abs() / torch.tensor(volume_shape, device=self.device) 64 | 65 | def other_from_self(self, other: "Volume") -> Tensor: 66 | """Homogenous transformation from own volume to other volume.""" 67 | return transform_volume_bounds(self.lower, self.upper, other.lower, other.upper) 68 | 69 | def self_from_other(self, other: "Volume") -> Tensor: 70 | """Homogenous transformation from other volume to own volume.""" 71 | return other.other_from_self(self) 72 | 73 | def _grid_sample_from_self(self) -> Tensor: 74 | """Homogenous transformation to grid_sample compatible coordinates.""" 75 | return transform_to_grid_sample_coords(self.lower, self.upper) 76 | 77 | def new_coord_grid(self, volume_shape: Sequence[int]) -> Tensor: 78 | """Uniform grid coordinates over volume.""" 79 | # coordinates correspond to centers of each voxel 80 | assert len(volume_shape) == 3 81 | coordinates = [(torch.arange(num_steps, device=self.device) + 0.5) / num_steps for num_steps in volume_shape] 82 | grid = torch.stack(torch.meshgrid(*coordinates, indexing="ij"), -1) 83 | grid = (self.upper - self.lower)[:, None, None, None, :] * grid + self.lower[:, None, None, None, :] 84 | return grid 85 | 86 | def coord_grid(self, features: Tensor, expand: bool = True) -> Tensor: 87 | """Grid coordinates for the given features.""" 88 | # coordinates correspond to centers of each voxel 89 | return self.new_coord_grid(self.volume_shape(features)).expand( 90 | features.shape[0] if expand else 1, -1, -1, -1, -1 91 | ) 92 | 93 | def sample_volume( 94 | self, volume_features: Tensor, sample_coords: Tensor, fill_invalid: float = 0.0, mode: str = "bilinear" 95 | ) -> Tensor: 96 | """Samples the provided features at given sample points.""" 97 | # volume_features: [B, C, X, Y, Z] 98 | # sample_coords: [B, X, Y, Z, 3] 99 | normalized_from_grid = self._grid_sample_from_self() 100 | normalized_coords = einsum_transform("bng,bxyzg->bxyzn", normalized_from_grid, points=sample_coords) 101 | normalized_coords = normalized_coords.expand(volume_features.shape[0], -1, -1, -1, -1) 102 | return torch.nn.functional.grid_sample( 103 | volume_features, normalized_coords, mode=mode, align_corners=False 104 | ).nan_to_num(fill_invalid) 105 | 106 | def clamp_points_along_line(self, points: Tensor, line_origin: Tensor): 107 | """Clamps points to volume bounds along given ray direction.""" 108 | # points [B, N, 3] 109 | # line_origin [B, N, 3] 110 | direction = points - line_origin 111 | direction = direction / direction.norm(dim=-1, keepdim=True) 112 | below_lower = (points - self.lower.unsqueeze(1)).clamp_max(0) 113 | above_upper = (points - self.upper.unsqueeze(1)).clamp_min(0) 114 | s = torch.maximum( 115 | (below_lower / direction).nan_to_num(0).amax(-1, keepdim=True), (above_upper / direction).nan_to_num(0).amax(-1, keepdim=True) 116 | ) 117 | boxed = points - s * direction 118 | return boxed 119 | -------------------------------------------------------------------------------- /scene_reconstruction/data/nuscenes/polars_helpers.py: -------------------------------------------------------------------------------- 1 | """Polars utils.""" 2 | 3 | import re 4 | from typing import Literal, Optional 5 | 6 | import numpy as np 7 | import polars as pl 8 | import pyarrow as pa 9 | import torch 10 | from polars.type_aliases import JoinStrategy, JoinValidation 11 | from scipy.spatial.transform import Rotation 12 | 13 | # pylint: disable=C0103 14 | 15 | 16 | def numpy_to_arrow(x: np.ndarray): 17 | """Numpy array into a a nested FixedSizeListArray.""" 18 | shape = x.shape 19 | arr = pa.array(x.ravel()) 20 | for last_dim in shape[-1:0:-1]: 21 | arr = pa.FixedSizeListArray.from_arrays(arr, last_dim) 22 | return arr 23 | 24 | 25 | def numpy_to_series(name: str, values: np.ndarray): 26 | """Numpy array to series.""" 27 | return pl.Series(name=name, values=numpy_to_arrow(values)) 28 | 29 | 30 | def torch_to_series(name: str, values: torch.Tensor): 31 | """Torch tensor to series.""" 32 | return numpy_to_series(name, values.numpy()) 33 | 34 | 35 | def series_to_numpy( 36 | series: pl.Series, 37 | zero_copy_only: bool = False, 38 | writable: bool = False, 39 | use_pyarrow: bool = True, 40 | ): 41 | """Polars series to numpy array.""" 42 | shape = [len(series)] 43 | while series.dtype == pl.Array: 44 | shape.append(series.dtype.width) 45 | series = series.explode() 46 | return series.to_numpy(zero_copy_only=zero_copy_only, writable=writable, use_pyarrow=use_pyarrow).reshape(shape) 47 | 48 | 49 | def series_to_torch( 50 | series: pl.Series, 51 | zero_copy_only: bool = False, 52 | use_pyarrow: bool = True, 53 | ): 54 | """Polars series to torch array.""" 55 | return torch.from_numpy( 56 | series_to_numpy(series=series, zero_copy_only=zero_copy_only, writable=True, use_pyarrow=use_pyarrow) 57 | ) 58 | 59 | 60 | TOKEN_REGEX = r"^.*\.token$" 61 | 62 | 63 | def col_is_token(col: str): 64 | """Check if column matches token regex.""" 65 | return re.search(TOKEN_REGEX, col, re.DOTALL) is not None 66 | 67 | 68 | def common_tokens(*dfs: pl.DataFrame): 69 | """Common tokens of multiple DataFrames.""" 70 | assert dfs 71 | all_tokens = [set(filter(col_is_token, df.columns)) for df in dfs] 72 | common = set.intersection(*all_tokens) 73 | return sorted(list(common)) 74 | 75 | 76 | def common_non_tokens(*dfs: pl.DataFrame): 77 | """Common non token columns of multiple DataFrames.""" 78 | assert dfs, "No DataFrames are given." 79 | all_non_tokens = [set(filter(lambda x: not col_is_token(x), df.columns)) for df in dfs] 80 | common = set.intersection(*all_non_tokens) 81 | return sorted(list(common)) 82 | 83 | 84 | def join_on_token( 85 | df: pl.DataFrame, 86 | other: pl.DataFrame, 87 | how: JoinStrategy = "inner", 88 | *, 89 | suffix: str = "_right", 90 | allow_duplicates: bool = False, 91 | validate: JoinValidation = "m:m", 92 | verbose: bool = False, 93 | ): 94 | """Join DataFrames on based on tokens.""" 95 | tokens = common_tokens(df, other) 96 | assert ( 97 | allow_duplicates or len(common_non_tokens(df, other)) == 0 98 | ), f"There are duplicate columns: {common_non_tokens(df, other)}" 99 | if verbose: 100 | print(f"Joining on {tokens}") 101 | return df.join(other, on=tokens, how=how, suffix=suffix, validate=validate) 102 | 103 | 104 | def quaternion_to_matrix(quat: np.ndarray, quat_order: Literal["wxyz", "xyzw"] = "wxyz"): 105 | """Convert quaternion to rotation matrix.""" 106 | if quat_order == "wxyz": 107 | quat = quat[:, [1, 2, 3, 0]] 108 | rot_mat = Rotation.from_quat(quat).as_matrix().astype(quat.dtype) 109 | return rot_mat 110 | 111 | 112 | def homogenous_transform(*, rotation_matrix: np.ndarray, translation_vector: np.ndarray): 113 | """Build homgenous transform from rotation and translation.""" 114 | *batch_size_translation, _3 = translation_vector.shape 115 | *batch_size_rotation, _3, _3 = rotation_matrix.shape 116 | assert batch_size_translation == batch_size_rotation 117 | homo = np.zeros(batch_size_rotation + [4, 4], dtype=rotation_matrix.dtype) 118 | homo[:, :3, :3] = rotation_matrix 119 | homo[:, :3, 3] = translation_vector 120 | homo[:, 3, 3] = 1.0 121 | return homo 122 | 123 | 124 | def inverse_transform(transform: np.ndarray): 125 | """Inverse homogenous transform.""" 126 | transform_inv = np.empty_like(transform) 127 | rot = transform[..., :3, :3] 128 | trans = transform[..., :3, 3] 129 | rot_inv = np.einsum("...ij->...ji", rot) 130 | trans_inv = -np.einsum("...ji,...i->...j", rot_inv, trans) 131 | transform_inv[..., :3, :3] = rot_inv 132 | transform_inv[..., :3, 3] = trans_inv 133 | transform_inv[..., 3, :3] = 0 134 | transform_inv[..., 3, 3] = 1 135 | return transform_inv 136 | 137 | 138 | def transform_from_columns( 139 | df: pl.DataFrame, 140 | rotation: str, 141 | translation: str, 142 | *, 143 | transform: Optional[str] = None, 144 | transform_inv: Optional[str] = None, 145 | ): 146 | """Parse homogenous transform from columns.""" 147 | if transform is None and transform_inv is None: 148 | return df 149 | rotation_quat = series_to_numpy(df[rotation]) 150 | translation_vec = series_to_numpy(df[translation]) 151 | rotation_matrix = quaternion_to_matrix(rotation_quat) 152 | transform_matrix = homogenous_transform(rotation_matrix=rotation_matrix, translation_vector=translation_vec) 153 | new_cols = {} 154 | if transform is not None: 155 | new_cols[transform] = transform_matrix 156 | if transform_inv is not None: 157 | new_cols[transform_inv] = inverse_transform(transform_matrix) 158 | return df.with_columns(*[numpy_to_series(k, v) for k, v in new_cols.items()]) 159 | 160 | 161 | def pad_intrinsics_to_4x4(intrinsics: np.ndarray): 162 | """Pad intrinsics to homogenous transform.""" 163 | *batch_size_rotation, _3, _3 = intrinsics.shape 164 | homo = np.zeros(batch_size_rotation + [4, 4], dtype=intrinsics.dtype) 165 | homo[:, :3, :3] = intrinsics 166 | homo[:, 3, 3] = 1.0 167 | return homo 168 | 169 | 170 | def pad_intrinsics_from_colums(df: pl.DataFrame, intrinsics: str, transform: str): 171 | """Pad intrinsics columns to homogenous transform.""" 172 | matrix = series_to_numpy(df[intrinsics]) 173 | transform_np = pad_intrinsics_to_4x4(matrix) 174 | return df.with_columns(numpy_to_series(transform, transform_np)) 175 | 176 | 177 | def nested_to_numpy(x, fill_value=None): 178 | """Nested numpy arrays to single numpy tensor.""" 179 | if isinstance(x, np.ndarray) and x.dtype == object: 180 | to_stack = [nested_to_numpy(i, fill_value=fill_value) if i is not None else fill_value for i in x] 181 | return np.stack(to_stack) 182 | return x 183 | 184 | 185 | def nested_to_torch(x, fill_value=None): 186 | """Nested numpy arrays to torch tensor.""" 187 | return torch.from_numpy(nested_to_numpy(x, fill_value=fill_value)) 188 | -------------------------------------------------------------------------------- /scene_reconstruction/data/nuscenes/parse.py: -------------------------------------------------------------------------------- 1 | """Polars NuScenes parsing.""" 2 | 3 | from collections.abc import Sequence 4 | from typing import Literal, Optional 5 | 6 | import polars as pl 7 | 8 | 9 | def token(input_name, output_name): 10 | """Token column.""" 11 | return pl.when(pl.col(input_name) != "").then(pl.col(input_name)).otherwise(None).alias(output_name) 12 | 13 | 14 | def rename(input_name, output_name, dtype: Optional[pl.DataType] = None): 15 | """Rename and cast..""" 16 | col = pl.col(input_name).alias(output_name) 17 | if dtype is not None: 18 | col = col.cast(dtype) 19 | return col 20 | 21 | 22 | def array(input_name, output_name, shape: Sequence[int], dtype=pl.Float32): 23 | """Rename and cast to array.""" 24 | array_type = dtype 25 | for dim in shape[::-1]: 26 | array_type = pl.Array(width=dim, inner=array_type) 27 | return pl.col(input_name).alias(output_name).cast(array_type) 28 | 29 | 30 | def timestamp(input_name, output_name): 31 | """Rename and cast to timestamp.""" 32 | return rename(input_name, output_name, dtype=pl.Datetime("us")) 33 | 34 | 35 | def camera_intrinsic(input_name, output_name): 36 | """Rename and cast to camera intrinsic.""" 37 | return ( 38 | pl.when(pl.col(input_name).list.len() > 0) 39 | .then(pl.col(input_name)) 40 | .otherwise(None) 41 | # identity matrix if no intrinsics are given to allow cast to array 42 | .fill_null([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]) 43 | .cast(pl.Array(width=3, inner=pl.Array(width=3, inner=pl.Float32))) 44 | .alias(output_name) 45 | ) 46 | 47 | 48 | def attribute_token(input_name, output_name): 49 | """Rename and cast to attribute token.""" 50 | return pl.col(input_name).list.first().alias(output_name) 51 | 52 | 53 | ATTRIBUTE_SELECTION = ( 54 | token("token", "attribute.token"), 55 | rename("name", "attribute.name"), 56 | rename("description", "attribute.description"), 57 | ) 58 | 59 | 60 | CALIBRATED_SENSOR_SELECTION = ( 61 | token("token", "calibrated_sensor.token"), 62 | token("sensor_token", "sensor.token"), 63 | array("translation", "calibrated_sensor.translation", [3]), 64 | array("rotation", "calibrated_sensor.rotation", [4]), 65 | camera_intrinsic("camera_intrinsic", "calibrated_sensor.camera_intrinsic"), 66 | ) 67 | 68 | 69 | CATEGORY_SELECTION = ( 70 | token("token", "category.token"), 71 | rename("name", "category.name"), 72 | rename("description", "category.description"), 73 | rename("index", "category.index"), 74 | ) 75 | 76 | EGO_POSE_SELECTION = ( 77 | token("token", "ego_pose.token"), 78 | array("translation", "ego_pose.translation", [3]), 79 | array("rotation", "ego_pose.rotation", [4]), 80 | timestamp("timestamp", "ego_pose.timestamp"), 81 | ) 82 | 83 | INSTANCE_SELECTION = ( 84 | token("token", "instance.token"), 85 | token("category_token", "category.token"), 86 | token("first_annotation_token", "first_annotation.token"), 87 | token("last_annotation_token", "last_annotation.token"), 88 | rename("nbr_annotations", "instance.nbr_annotations"), 89 | ) 90 | 91 | LIDARSEG_SELECTION = ( 92 | token("token", "LIDAR_TOP.lidarseg.token"), 93 | token("sample_data_token", "LIDAR_TOP.sample_data.token"), 94 | rename("filename", "LIDAR_TOP.lidarseg.filename"), 95 | ) 96 | 97 | 98 | LOG_SELECTION = ( 99 | token("token", "log.token"), 100 | rename("logfile", "log.logfile"), 101 | rename("vehicle", "log.vehicle"), 102 | rename("date_captured", "log.date_captured"), 103 | rename("location", "log.location"), 104 | ) 105 | 106 | 107 | MAP_SELECTION = ( 108 | token("token", "map.token"), 109 | rename("category", "map.category"), 110 | rename("filename", "map.filename"), 111 | rename("log_tokens", "map.log_tokens"), 112 | ) 113 | 114 | 115 | SAMPLE_SELECTION = ( 116 | token("token", "sample.token"), 117 | token("next", "sample.next"), 118 | token("prev", "sample.prev"), 119 | token("scene_token", "scene.token"), 120 | timestamp("timestamp", "sample.timestamp"), 121 | ) 122 | 123 | 124 | SAMPLE_ANNOTATION_SELECTION = ( 125 | token("token", "sample_annotation.token"), 126 | token("next", "sample_annotation.next"), 127 | token("prev", "sample_annotation.prev_token"), 128 | token("sample_token", "sample.token"), 129 | token("instance_token", "instance.token"), 130 | token("visibility_token", "visibility.token"), 131 | attribute_token( 132 | "attribute_tokens", "attribute.token" 133 | ), # there is always ONE attribute in this list -> single token 134 | array("translation", "sample_annotation.translation", [3]), 135 | array("rotation", "sample_annotation.rotation", [4]), 136 | array("size", "sample_annotation.size", [3]), 137 | rename("num_lidar_pts", "sample_annotation.num_lidar_pts"), 138 | rename("num_radar_pts", "sample_annotation.num_radar_pts"), 139 | ) 140 | 141 | SAMPLE_DATA_SELECTION = ( 142 | token("token", "sample_data.token"), 143 | token("next", "sample_data.next"), 144 | token("prev", "sample_data.prev"), 145 | token("sample_token", "sample.token"), 146 | token("ego_pose_token", "ego_pose.token"), 147 | token("calibrated_sensor_token", "calibrated_sensor.token"), 148 | rename("filename", "sample_data.filename"), 149 | rename("fileformat", "sample_data.fileformat"), 150 | rename("width", "sample_data.width"), 151 | rename("height", "sample_data.height"), 152 | timestamp("timestamp", "sample_data.timestamp"), 153 | rename("is_key_frame", "sample_data.is_key_frame"), 154 | ) 155 | 156 | SCENE_SELECTION = ( 157 | token("token", "scene.token"), 158 | token("log_token", "log.token"), 159 | token("first_sample_token", "first_sample.token"), 160 | token("last_sample_token", "last_sample.token"), 161 | rename("name", "scene.name"), 162 | rename("description", "scene.description"), 163 | rename("nbr_samples", "scene.nbr_samples"), 164 | ) 165 | 166 | SENSOR_SELECTION = ( 167 | token("token", "sensor.token"), 168 | rename("channel", "sensor.channel"), 169 | rename("modality", "sensor.modality"), 170 | ) 171 | 172 | 173 | VISIBILITY_SELECTION = ( 174 | token("token", "visibility.token"), 175 | rename("level", "visibility.level"), 176 | rename("description", "visibility.description"), 177 | ) 178 | PANOPTIC_SELECTION = ( 179 | token("token", "LIDAR_TOP.panoptic.token"), 180 | token("sample_data_token", "LIDAR_TOP.sample_data.token"), 181 | rename("filename", "LIDAR_TOP.panoptic.filename"), 182 | ) 183 | 184 | 185 | ANNOTATION_FILES = Literal[ 186 | "log", 187 | "scene", 188 | "instance", 189 | "category", 190 | "map", 191 | "sample", 192 | "lidarseg", 193 | "calibrated_sensor", 194 | "sample_data", 195 | "sample_annotation", 196 | "attribute", 197 | "sensor", 198 | "ego_pose", 199 | "visibility", 200 | "panoptic", 201 | ] 202 | 203 | SELECTION: dict[ANNOTATION_FILES, Sequence[pl.Expr]] = { 204 | "log": LOG_SELECTION, 205 | "scene": SCENE_SELECTION, 206 | "instance": INSTANCE_SELECTION, 207 | "category": CATEGORY_SELECTION, 208 | "map": MAP_SELECTION, 209 | "sample": SAMPLE_SELECTION, 210 | "lidarseg": LIDARSEG_SELECTION, 211 | "calibrated_sensor": CALIBRATED_SENSOR_SELECTION, 212 | "sample_data": SAMPLE_DATA_SELECTION, 213 | "sample_annotation": SAMPLE_ANNOTATION_SELECTION, 214 | "attribute": ATTRIBUTE_SELECTION, 215 | "sensor": SENSOR_SELECTION, 216 | "ego_pose": EGO_POSE_SELECTION, 217 | "visibility": VISIBILITY_SELECTION, 218 | "panoptic": PANOPTIC_SELECTION, 219 | } 220 | -------------------------------------------------------------------------------- /scene_reconstruction/occupancy/grid.py: -------------------------------------------------------------------------------- 1 | """Dense spherical occupancy grid.""" 2 | 3 | import math 4 | from collections.abc import Sequence 5 | from typing import Literal, Optional, Union 6 | 7 | import einops 8 | import torch 9 | from torch import Tensor 10 | 11 | from scene_reconstruction.core import Volume 12 | from scene_reconstruction.core.transform import einsum_transform 13 | from scene_reconstruction.math.spherical_coordinate_system import cartesian_to_spherical, spherical_volume_element 14 | 15 | MODES = Literal["clamp", "warp", "drop"] 16 | 17 | 18 | def volume_density_from_points( 19 | volume_shape: Sequence[int], 20 | volume: Volume, 21 | points: Tensor, 22 | points_weights: Optional[Tensor] = None, 23 | modes: Sequence[Union[tuple[MODES, MODES], MODES]] = ("clamp", "clamp", "clamp"), 24 | ): 25 | """Scatters points into nearby voxels by considering surrounding voxels.""" 26 | # points: [B, N, 3] 27 | # points_weights: [B, N, C] 28 | batch_size = points.shape[0] 29 | index_volume = Volume.new_index(volume_shape, device=points.device) 30 | index_from_volume = index_volume.self_from_other(volume) 31 | float_index = einsum_transform("biv,bnv->bni", index_from_volume, points=points) # [B, N, 3] 32 | # next point on the voxel grid 33 | next_grid_index = float_index.round() 34 | # cube with side length 1.0 and center 0.5 35 | corner_offsets = 0.5 * torch.tensor( 36 | [ 37 | [1, 1, 1], 38 | [1, 1, -1], 39 | [1, -1, 1], 40 | [1, -1, -1], 41 | [-1, 1, 1], 42 | [-1, 1, -1], 43 | [-1, -1, 1], 44 | [-1, -1, -1], 45 | ], 46 | dtype=torch.float, 47 | device=volume.device, 48 | ) 49 | cube_corners_float_index = float_index[:, :, None, :] + corner_offsets # B, N, 8, 3 50 | # volume in each surrounding volume element 51 | corner_volume = (cube_corners_float_index - next_grid_index[:, :, None, :]).abs().prod(-1) # B, N, 8 52 | cube_corners_index = cube_corners_float_index.floor().long() # B, N, 8, 3 53 | assert len(modes) == cube_corners_index.shape[-1] 54 | cube_corners_index_processed = [] 55 | for idx, mode, shape in zip(cube_corners_index.unbind(-1), modes, volume_shape): 56 | if not isinstance(mode, (list, tuple)): 57 | m_lower, m_upper = (mode, mode) 58 | else: 59 | m_lower, m_upper = mode 60 | 61 | lower_mask = idx < 0 62 | if m_lower == "clamp": 63 | idx = torch.where(lower_mask, 0, idx) 64 | elif m_lower == "warp": 65 | idx = torch.where(lower_mask, idx % shape, idx) 66 | elif m_lower == "drop": 67 | idx = torch.where(lower_mask, 0, idx) 68 | corner_volume = torch.where(lower_mask, 0.0, corner_volume) 69 | 70 | upper_mask = idx >= shape 71 | if m_upper == "clamp": 72 | idx = torch.where(upper_mask, shape - 1, idx) 73 | elif m_upper == "warp": 74 | idx = torch.where(upper_mask, idx % shape, idx) 75 | elif m_upper == "drop": 76 | idx = torch.where(upper_mask, shape - 1, idx) 77 | corner_volume = torch.where(upper_mask, 0.0, corner_volume) 78 | 79 | cube_corners_index_processed.append(idx) 80 | cube_corners_index = torch.stack(cube_corners_index_processed, -1) 81 | 82 | num_channels = points_weights.shape[-1] if points_weights is not None else 1 83 | features_flat = torch.zeros(batch_size, num_channels, math.prod(volume_shape), device=points.device) 84 | if points_weights is not None: 85 | feat_to_scatter = ( 86 | corner_volume[:, :, :, None] * points_weights[:, :, None, :] 87 | ) # [B, N, 8, 1] * [B, N, 1, C] = [B, N, 8, C] 88 | 89 | else: 90 | feat_to_scatter = corner_volume[:, :, :, None] 91 | feat_to_scatter = einops.rearrange(feat_to_scatter, "b n m c -> b c (n m)") 92 | stride = torch.tensor( 93 | [volume_shape[1] * volume_shape[2], volume_shape[2], 1], 94 | device=cube_corners_float_index.device, 95 | dtype=torch.long, 96 | ) 97 | cube_corners_index_flat = einops.rearrange(cube_corners_index, "b n m c -> b (n m) c") 98 | volume_index_flat = (stride * cube_corners_index_flat).sum(-1) # [B, N * 8] 99 | 100 | features_flat.scatter_add_(2, volume_index_flat.unsqueeze(1).expand_as(feat_to_scatter), feat_to_scatter) 101 | features = features_flat.unflatten(-1, volume_shape) 102 | return features 103 | 104 | 105 | def spherical_reflection_and_transmission_from_lidar( 106 | points_lidar: Tensor, 107 | points_weight: Tensor, 108 | spherical_volume: Volume, 109 | spherical_shape: Sequence[int], 110 | normalize: Optional[Literal["volume"]] = None, 111 | ): 112 | """Occupancy estimation from pointclouds.""" 113 | 114 | points_lidar_spherical = cartesian_to_spherical(points_lidar) 115 | 116 | # add one bin along r 117 | spherical_shape_with_inf = [spherical_shape[0] + 1, spherical_shape[1], spherical_shape[2]] 118 | range_voxel_size = spherical_volume.voxel_size_from_shape(spherical_shape)[0, 0].item() 119 | 120 | spherical_volume_with_inf = Volume( 121 | lower=spherical_volume.lower, 122 | upper=spherical_volume.upper + torch.tensor([range_voxel_size, 0.0, 0.0], device=spherical_volume.upper.device), 123 | ) 124 | 125 | reflection_and_transmission = volume_density_from_points( 126 | volume_shape=spherical_shape_with_inf, 127 | volume=spherical_volume_with_inf, 128 | points=points_lidar_spherical, 129 | points_weights=points_weight, 130 | modes=(("drop", "clamp"), "clamp", "warp"), # clamp upper range, drop lower range 131 | ) 132 | num_reflections = reflection_and_transmission 133 | 134 | # exclusive cumulative sum from end to start equal to 135 | # num_transmissions = ( 136 | # reflection_and_transmission.flip(2).cumsum(2).flip(2) - reflection_and_transmission 137 | # ) 138 | num_transmissions = reflection_and_transmission.sum(2, keepdim=True) - reflection_and_transmission.cumsum(2) 139 | # drop infinity bin 140 | reflection_and_transmission = torch.cat( 141 | [num_reflections[:, :, :-1, :, :], num_transmissions[:, :, :-1, :, :]], 1 142 | ) # [B, C, R, PHI, THETA] 143 | if normalize == "volume": 144 | grid = spherical_volume.coord_grid(reflection_and_transmission, expand=False) 145 | voxel_size = spherical_volume.voxel_size(reflection_and_transmission) 146 | elementwise_volume = spherical_volume_element( 147 | grid - 0.5 * voxel_size[:, None, None, None, :], grid + 0.5 * voxel_size[:, None, None, None, :] 148 | ) 149 | 150 | reflection_and_transmission = reflection_and_transmission / elementwise_volume.unsqueeze(1) 151 | 152 | return reflection_and_transmission 153 | 154 | 155 | def occupancy_from_points( 156 | points_lidar: Tensor, 157 | points_weight: Tensor, 158 | lidar_from_ego: Tensor, 159 | cartesian_volume: Volume, 160 | spherical_volume: Volume, 161 | cartesian_shape: Sequence[int], 162 | spherical_shape: Sequence[int], 163 | ): 164 | """Occupancy estimation from pointclouds.""" 165 | 166 | points_lidar_spherical = cartesian_to_spherical(points_lidar) 167 | 168 | density_spherical = volume_density_from_points( 169 | volume_shape=spherical_shape, 170 | volume=spherical_volume, 171 | points=points_lidar_spherical, 172 | points_weights=points_weight, 173 | modes=(("drop", "clamp"), "clamp", "warp"), # clamp upper range, drop lower range 174 | ) 175 | num_reflections = density_spherical 176 | num_transmissions = density_spherical.flip(2).cumsum(2).flip(2) - density_spherical # exclusive cumsum 177 | density_spherical = torch.cat([num_reflections, num_transmissions], 1) 178 | 179 | grid_ego = cartesian_volume.new_coord_grid(cartesian_shape) 180 | grid_lidar = einsum_transform("ble,bxyze->bxyzl", lidar_from_ego, points=grid_ego) 181 | grid_lidar_spherical = cartesian_to_spherical(grid_lidar) 182 | density_sampled = spherical_volume.sample_volume(density_spherical, grid_lidar_spherical) 183 | 184 | scale = cartesian_volume.voxel_size_from_shape(cartesian_shape).prod(-1) / spherical_volume_element( 185 | grid_lidar_spherical - 0.5 * spherical_volume.voxel_size_from_shape(spherical_shape)[:, None, None, None, :], 186 | grid_lidar_spherical + 0.5 * spherical_volume.voxel_size_from_shape(spherical_shape)[:, None, None, None, :], 187 | ).clamp_min(1e-6) 188 | density_sampled = density_sampled * scale[:, None] 189 | return density_sampled 190 | -------------------------------------------------------------------------------- /scene_reconstruction/occupancy/transmission_reflection.py: -------------------------------------------------------------------------------- 1 | """Transmission and reflections from lidar.""" 2 | 3 | from __future__ import annotations 4 | 5 | import math 6 | from dataclasses import dataclass 7 | from pathlib import Path 8 | 9 | import numpy as np 10 | import polars as pl 11 | import tqdm 12 | from torch import Tensor 13 | 14 | from scene_reconstruction.core.volume import Volume 15 | from scene_reconstruction.data.nuscenes.dataset import NuscenesDataset 16 | from scene_reconstruction.data.nuscenes.polars_helpers import series_to_torch, torch_to_series 17 | from scene_reconstruction.occupancy.grid import occupancy_from_points, spherical_reflection_and_transmission_from_lidar 18 | 19 | 20 | def _batch(iterable, n=1): 21 | lenght = len(iterable) 22 | for ndx in range(0, lenght, n): 23 | yield iterable[ndx : min(ndx + n, lenght)] 24 | 25 | 26 | @dataclass 27 | class TransmissionAndReflections: 28 | """Calculates number of transmission and reflectiones per voxel cell from lidar pointcloud.""" 29 | 30 | ds: NuscenesDataset 31 | extra_data_root: Path 32 | cartesian_lower: tuple[float, float, float] = (-40.0, -40.0, -1.0) 33 | cartesian_upper: tuple[float, float, float] = (40.0, 40.0, 5.4) 34 | cartesian_shape: tuple[int, int, int] = (400, 400, 32) 35 | spherical_lower: tuple[float, float, float] = (0.0, (90 - 15) / 180 * math.pi, -math.pi) 36 | spherical_upper: tuple[float, float, float] = (60.0, (90 + 35) / 180 * math.pi, math.pi) 37 | spherical_shape: tuple[int, int, int] = (600, 100, 720) 38 | lidar_min_distance: float = 3.0 39 | # voxel_size cartesian = 0.2m 40 | # voxel_size sperical : 0.1m, 0.5°, 0.5° 41 | 42 | batch_size: int = 4 43 | 44 | def process_data(self) -> None: 45 | """Process dataset.""" 46 | for scene in tqdm.tqdm(self.ds.scene.partition_by("scene.token"), position=0): 47 | scene = self.ds.join(scene, self.ds.sample) 48 | scene = self.ds.load_sample_data(scene, "LIDAR_TOP") 49 | self.process_scene(scene) 50 | 51 | def save_path(self, scene_name: str, lidar_top_token: str) -> Path: 52 | """Save from scene name and lidar token.""" 53 | path = ( 54 | Path(self.extra_data_root) 55 | / scene_name 56 | / "reflection_and_transmission" 57 | / "LIDAR_TOP" 58 | / f"{lidar_top_token}.npz" 59 | ) 60 | path.parent.mkdir(exist_ok=True, parents=True) 61 | return path 62 | 63 | def save_data(self, save_filename: Path, reflections_and_transmissions: Tensor): 64 | """Save data.""" 65 | num_reflections, num_transmissions = reflections_and_transmissions.unbind(0) 66 | np.savez_compressed( 67 | save_filename, 68 | num_reflections=num_reflections.numpy(), 69 | num_transmissions=num_transmissions.numpy(), 70 | ) 71 | 72 | def process_scene(self, scene: pl.DataFrame) -> None: 73 | """Process single scene.""" 74 | points_lidar = series_to_torch(scene["LIDAR_TOP.sample_data.points_lidar"]) 75 | ego_from_lidar = series_to_torch(scene["LIDAR_TOP.transform.ego_from_sensor"]) 76 | lidar_from_ego = ego_from_lidar.inverse() 77 | points_range_mask = points_lidar.norm(dim=-1) >= self.lidar_min_distance 78 | points_mask = series_to_torch(scene["LIDAR_TOP.sample_data.points_mask"]) & points_range_mask 79 | filename = [ 80 | self.save_path(s_name, l_token) 81 | for s_name, l_token in zip(scene["scene.name"], scene["LIDAR_TOP.sample_data.token"]) 82 | ] 83 | for points_lidar_batched, lidar_from_ego_batched, points_mask_batched, filename_batched in tqdm.tqdm( 84 | zip( 85 | points_lidar.split(self.batch_size), 86 | lidar_from_ego.split(self.batch_size), 87 | points_mask.split(self.batch_size), 88 | _batch(filename, self.batch_size), 89 | ), 90 | total=(points_lidar.shape[0] + self.batch_size - 1) // self.batch_size, 91 | position=1, 92 | ): 93 | reflections_and_transmissions_batched = self.process_batch( 94 | points_lidar_batched.cuda(), 95 | lidar_from_ego_batched.cuda(), 96 | points_mask_batched.cuda(), 97 | ) 98 | 99 | for filename, reflections_and_transmissions in zip( 100 | filename_batched, reflections_and_transmissions_batched.cpu().unbind(0) 101 | ): 102 | self.save_data(filename, reflections_and_transmissions) # type: ignore 103 | 104 | def process_batch(self, points_lidar: Tensor, lidar_from_ego: Tensor, points_mask: Tensor) -> Tensor: 105 | """Process single batch of data.""" 106 | # points_lidar: B, N, 3 107 | # points_mask: B, N 108 | # lidar_from_ego: B, 4, 4 109 | points_weight = points_mask.unsqueeze(-1).float() 110 | reflections_and_transmissions = occupancy_from_points( 111 | points_lidar, 112 | lidar_from_ego=lidar_from_ego, 113 | points_weight=points_weight, 114 | cartesian_volume=Volume.new_volume(self.cartesian_lower, self.cartesian_upper, device=points_lidar.device), 115 | spherical_volume=Volume.new_volume(self.spherical_lower, self.spherical_upper, device=points_lidar.device), 116 | cartesian_shape=self.cartesian_shape, 117 | spherical_shape=self.spherical_shape, 118 | ) 119 | return reflections_and_transmissions # [B, 2, X, Y, Z] 120 | 121 | 122 | @dataclass 123 | class ReflectionTransmissionSpherical: 124 | """Calculates number of transmission and reflectiones per voxel cell from lidar pointcloud in spherical coordinates in the lidar frame.""" 125 | 126 | ds: NuscenesDataset 127 | extra_data_root: Path 128 | spherical_lower: tuple[float, float, float] = (2.0, (90 - 15) / 180 * math.pi, -math.pi) 129 | spherical_upper: tuple[float, float, float] = (60.0, (90 + 35) / 180 * math.pi, math.pi) 130 | spherical_shape: tuple[int, int, int] = (600, 100, 720) 131 | lidar_min_distance: float = 2.0 132 | # voxel_size sperical : 0.1m, 0.5°, 0.5° 133 | 134 | batch_size: int = 1 135 | device: str = "cuda" 136 | name: str = "reflection_and_transmission_spherical" 137 | 138 | def process_data(self) -> None: 139 | """Process dataset.""" 140 | for scene in tqdm.tqdm( 141 | self.ds.scene.iter_slices(1), total=len(self.ds.scene), position=0, desc="Processing scenes" 142 | ): 143 | self.process_scene(scene) 144 | 145 | def save_path(self, scene_name: str, lidar_top_token: str) -> Path: 146 | """Save from scene name and lidar token.""" 147 | path = Path(self.extra_data_root) / self.name / scene_name / "LIDAR_TOP" / f"{lidar_top_token}.arrow" 148 | path.parent.mkdir(exist_ok=True, parents=True) 149 | return path 150 | 151 | def process_scene(self, scene: pl.DataFrame) -> None: 152 | """Process single scene.""" 153 | # load all sample data with lidar token 154 | scene = self.ds.join(scene, self.ds.sample) 155 | scene = self.ds.load_sample_data(scene, "LIDAR_TOP", with_data=False) 156 | spherical_volume = Volume.new_volume(self.spherical_lower, self.spherical_upper) 157 | volume_lower = torch_to_series( 158 | "LIDAR_TOP.reflection_and_transmission_spherical.volume.lower", spherical_volume.lower 159 | ) 160 | volume_upper = torch_to_series( 161 | "LIDAR_TOP.reflection_and_transmission_spherical.volume.upper", spherical_volume.upper 162 | ) 163 | for chunked_scene in tqdm.tqdm( 164 | scene.iter_slices(self.batch_size), 165 | total=(len(scene) + self.batch_size - 1) // self.batch_size, 166 | desc="Processing sample", 167 | position=1, 168 | ): 169 | chunked_scene = self.ds._load_sensor_data(chunked_scene, "LIDAR_TOP") 170 | points_lidar = series_to_torch(chunked_scene["LIDAR_TOP.sample_data.points_lidar"]).to( 171 | self.device, non_blocking=True 172 | ) 173 | points_range_mask = points_lidar.norm(dim=-1) >= self.lidar_min_distance 174 | points_mask = ( 175 | series_to_torch(chunked_scene["LIDAR_TOP.sample_data.points_mask"]).to(self.device, non_blocking=True) 176 | & points_range_mask 177 | ) 178 | reflections_and_transmissions = self.process_batch( 179 | points_lidar, points_mask, spherical_volume.to(points_mask.device, non_blocking=True) 180 | ) 181 | 182 | for sample, rt in zip( 183 | chunked_scene.select( 184 | "scene.token", "sample.token", "scene.name", "LIDAR_TOP.sample_data.token" 185 | ).iter_slices(1), 186 | reflections_and_transmissions.split(1), 187 | ): 188 | filename = self.save_path(sample["scene.name"].item(), sample["LIDAR_TOP.sample_data.token"].item()) 189 | sample = sample.select("LIDAR_TOP.sample_data.token").with_columns( 190 | torch_to_series(f"LIDAR_TOP.{self.name}", rt), 191 | volume_lower, 192 | volume_upper, 193 | ) 194 | 195 | sample.write_ipc(filename, compression="zstd") 196 | 197 | def process_batch(self, points_lidar: Tensor, points_mask: Tensor, spherical_volume: Volume) -> Tensor: 198 | """Process single batch of data.""" 199 | # points_lidar: B, N, 3 200 | # points_mask: B, N 201 | # lidar_from_ego: B, 4, 4 202 | points_weight = points_mask.unsqueeze(-1).float() 203 | reflections_and_transmissions = spherical_reflection_and_transmission_from_lidar( 204 | points_lidar, 205 | points_weight=points_weight, 206 | spherical_volume=spherical_volume, 207 | spherical_shape=self.spherical_shape, 208 | ) 209 | return reflections_and_transmissions.to(device="cpu", non_blocking=True) # .cpu() # [B, 2, X, Y, Z] 210 | -------------------------------------------------------------------------------- /scene_reconstruction/data/nuscenes/scene_flow.py: -------------------------------------------------------------------------------- 1 | """Scene flow.""" 2 | 3 | import math 4 | from collections.abc import Sequence 5 | from dataclasses import dataclass 6 | from pathlib import Path 7 | 8 | import numpy as np 9 | import polars as pl 10 | import torch 11 | import tqdm 12 | from torch import Tensor 13 | 14 | from scene_reconstruction.core.transform import einsum_transform 15 | from scene_reconstruction.core.volume import Volume 16 | from scene_reconstruction.data.nuscenes.dataset import NuscenesDataset 17 | from scene_reconstruction.data.nuscenes.polars_helpers import nested_to_torch, series_to_torch, torch_to_series 18 | 19 | 20 | def build_scene_instance_motion( 21 | ds: NuscenesDataset, 22 | scene: pl.DataFrame, 23 | scene_with_annos: pl.DataFrame, 24 | scene_instance_index: pl.DataFrame, 25 | ): 26 | """Aggregate motion information per instance.""" 27 | scene_matrix = scene_with_annos.with_columns( 28 | pl.col("transform.obj_from_global").cast(pl.List(pl.List(pl.Float32))) 29 | ).pivot( 30 | index="LIDAR_TOP.sample_data.token", 31 | columns="scene_instance.index", 32 | values="transform.obj_from_global", 33 | ) 34 | scene_matrix_token = scene_matrix.select("LIDAR_TOP.sample_data.token") 35 | scene_matrix_data = scene_matrix.select([str(i) for i in scene_instance_index["scene_instance.index"].to_list()]) 36 | scene_obj_from_global = nested_to_torch( 37 | scene_matrix_data.to_numpy(), fill_value=np.full((4, 4), float("nan"), dtype=np.float32) 38 | ) # [time index, scene instance index, 4, 4] 39 | scene_obj_from_global_valid_mask = nested_to_torch( 40 | scene_matrix_data.select(pl.all().is_not_null()).to_numpy() 41 | ) # [time index, scene instance index] 42 | 43 | # static transform for reserved instance id 0 44 | 45 | scene_obj_from_global = torch.cat( 46 | [torch.eye(4)[None, None].expand_as(scene_obj_from_global[:, :1]), scene_obj_from_global], 1 47 | ) 48 | scene_obj_from_global_valid_mask = torch.cat( 49 | [torch.ones_like(scene_obj_from_global_valid_mask[:, :1]), scene_obj_from_global_valid_mask], 1 50 | ) 51 | 52 | scene_matrix = scene_matrix_token.with_columns( 53 | torch_to_series("LIDAR_TOP.scene_flow.instance_from_global", scene_obj_from_global), 54 | torch_to_series("LIDAR_TOP.scene_flow.instance_from_global_mask", scene_obj_from_global_valid_mask), 55 | ) 56 | scene_matrix_token_without_annos = scene.select("LIDAR_TOP.sample_data.token").filter( 57 | pl.col("LIDAR_TOP.sample_data.token").is_in(scene_matrix["LIDAR_TOP.sample_data.token"]).not_() 58 | ) 59 | num_samples_without_annos = len(scene_matrix_token_without_annos) 60 | 61 | if num_samples_without_annos > 0: 62 | num_instances = scene_obj_from_global.shape[1] 63 | scene_obj_from_global_without_annos = torch.cat( 64 | [ 65 | torch.eye(4)[None, None].expand(num_samples_without_annos, -1, -1, -1), 66 | torch.full((num_samples_without_annos, num_instances - 1, 4, 4), fill_value=float("nan")), 67 | ], 68 | 1, 69 | ) 70 | scene_obj_from_global_valid_mask_without_annos = torch.cat( 71 | [ 72 | torch.ones(num_samples_without_annos, 1, dtype=torch.bool), 73 | torch.zeros(num_samples_without_annos, num_instances - 1, dtype=torch.bool), 74 | ], 75 | 1, 76 | ) 77 | scene_matrix_without_annos = scene_matrix_token_without_annos.with_columns( 78 | torch_to_series("LIDAR_TOP.scene_flow.instance_from_global", scene_obj_from_global_without_annos), 79 | torch_to_series( 80 | "LIDAR_TOP.scene_flow.instance_from_global_mask", scene_obj_from_global_valid_mask_without_annos 81 | ), 82 | ) 83 | if len(scene_matrix) > 0: 84 | scene_matrix = pl.concat([scene_matrix, scene_matrix_without_annos], how="vertical") 85 | else: 86 | scene_matrix = scene_matrix_without_annos 87 | 88 | scene_instance_motion = ds.join(scene, scene_matrix) 89 | assert len(scene) == len(scene_instance_motion) 90 | 91 | return scene_instance_motion 92 | 93 | 94 | def extract_instance_ids_chunked( 95 | lidar_token: pl.DataFrame, 96 | sample: pl.DataFrame, 97 | ego_volume: Volume, 98 | volume_shape: Sequence[int], 99 | max_points: int = 200000, 100 | **kwargs, 101 | ): 102 | """Extract instance ids for each grid point.""" 103 | obj_from_global = series_to_torch(sample["transform.obj_from_global"]).to(**kwargs) 104 | obj_bbox = series_to_torch(sample["sample_annotation.size_aligned"]).to(**kwargs) 105 | global_from_ego = series_to_torch(sample["LIDAR_TOP.transform.global_from_ego"]).to(**kwargs) 106 | instance_index = series_to_torch(sample["scene_instance.index"]).to(**kwargs) 107 | obj_from_ego = obj_from_global @ global_from_ego 108 | num_points = math.prod(volume_shape) 109 | num_chunks = (num_points + max_points - 1) // max_points 110 | grid_points = ego_volume.to(**kwargs).new_coord_grid(volume_shape) 111 | # chunk along x dim 112 | instance_ids_chunked = [] 113 | for grid_points_chunk in grid_points.chunk(num_chunks, dim=1): 114 | grid_coordinates_obj = einsum_transform("boe,bxyze->bxyzo", obj_from_ego, points=grid_points_chunk.to(**kwargs)) 115 | bbox_half = 0.5 * obj_bbox[:, None, None, None, :] 116 | distance_to_center = grid_coordinates_obj.norm(dim=-1) 117 | inside_box = (grid_coordinates_obj.abs() <= bbox_half).all(-1) 118 | soft_box_distance = torch.where(inside_box, distance_to_center, float("inf")) 119 | assigned_bbox = soft_box_distance.argmin(0) 120 | in_any_bbox = inside_box.sum(0) > 0 121 | instance_ids = torch.where(in_any_bbox, instance_index[assigned_bbox], 0) # 0 = static background 122 | instance_ids_chunked.append(instance_ids.cpu()) 123 | instance_ids_combined = torch.cat(instance_ids_chunked) 124 | assert instance_ids_combined.shape == volume_shape 125 | lidar_token = lidar_token.with_columns( 126 | torch_to_series("LIDAR_TOP.scene_flow.scene_instance_index", instance_ids_combined[None]) 127 | ) 128 | return lidar_token 129 | 130 | 131 | def extract_instance_ids( 132 | lidar_token: pl.DataFrame, sample: pl.DataFrame, ego_volume: Volume, volume_shape: Sequence[int] 133 | ): 134 | """Extract instance ids for each grid point.""" 135 | obj_from_global = series_to_torch(sample["transform.obj_from_global"]).cuda() 136 | obj_bbox = series_to_torch(sample["sample_annotation.size_aligned"]).cuda() 137 | global_from_ego = series_to_torch(sample["LIDAR_TOP.transform.global_from_ego"]).cuda() 138 | instance_index = series_to_torch(sample["scene_instance.index"]).cuda() 139 | obj_from_ego = obj_from_global @ global_from_ego 140 | grid_coordinates_obj = einsum_transform( 141 | "boe,bxyze->bxyzo", obj_from_ego, points=ego_volume.new_coord_grid(volume_shape).cuda() 142 | ) 143 | bbox_half = 0.5 * obj_bbox[:, None, None, None, :] 144 | distance_to_center = grid_coordinates_obj.norm(dim=-1) 145 | inside_box = (grid_coordinates_obj.abs() <= bbox_half).all(-1) 146 | soft_box_distance = torch.where(inside_box, distance_to_center, float("inf")) 147 | assigned_bbox = soft_box_distance.argmin(0) 148 | in_any_bbox = inside_box.sum(0) > 0 149 | instance_ids = torch.where(in_any_bbox, instance_index[assigned_bbox], 0) # 0 = static background 150 | lidar_token = lidar_token.with_columns(torch_to_series("scene_flow.scene_instance_index", instance_ids[None])) 151 | return lidar_token 152 | 153 | 154 | def lookup_instance_transform( 155 | instance_ids: Tensor, 156 | instance_from_tgt: Tensor, 157 | instance_from_src: Tensor, 158 | instance_from_tgt_mask: Tensor, 159 | instance_from_src_mask: Tensor, 160 | ): 161 | """Lookup transforms for each grid cell.""" 162 | # instance_ids: B, X, Y, Z 163 | # instance_from_tgt, instance_from_src: B, X, Y, Z, 4, 4 164 | # instance_from_tgt_mask, instance_from_src_mask: B, X, Y, Z 165 | # instance needs to be part of both frames transform 166 | # fall back to static transform 167 | instance_ids_flat = instance_ids.flatten(1) 168 | tgt_mask = instance_from_tgt_mask.gather(1, instance_ids_flat) 169 | src_mask = instance_from_src_mask.gather(1, instance_ids_flat) 170 | valid_instance_ids_flat = torch.where( 171 | tgt_mask & src_mask, 172 | instance_ids_flat, 173 | 0, 174 | ) 175 | tgt_from_src = instance_from_tgt.inverse() @ instance_from_src 176 | tgt_from_src_gathered_flat = tgt_from_src.gather(1, valid_instance_ids_flat[..., None, None].expand(-1, -1, 4, 4)) 177 | tgt_from_src_transform = tgt_from_src_gathered_flat.view(*instance_ids.shape, 4, 4) 178 | 179 | return tgt_from_src_transform 180 | 181 | 182 | @dataclass 183 | class SceneFlow: 184 | """Scene flow from boxes.""" 185 | 186 | ds: NuscenesDataset 187 | extra_data_root: Path 188 | 189 | cartesian_lower: tuple[float, float, float] = (-40.0, -40.0, -1.0) 190 | cartesian_upper: tuple[float, float, float] = (40.0, 40.0, 5.4) 191 | cartesian_shape: tuple[int, int, int] = (400, 400, 32) 192 | 193 | device: str = "cuda" 194 | missing_only: bool = False 195 | 196 | def load_scene(self, scene_name: str): 197 | """Loads a single scene for processing.""" 198 | # load single scene 199 | scene = self.ds.scene_by_name(scene_name) 200 | scene = self.ds.load_sample_data(scene, "LIDAR_TOP", with_data=False) 201 | scene_with_annos = self.ds.join(scene, self.ds.sample_annotation_dict["LIDAR_TOP"]) 202 | scene_with_annos = self.ds.join(scene_with_annos, self.ds.instance) 203 | # unique instance id per scene starting from 1 (0 reserved for static world) 204 | scene_instance_index = ( 205 | scene_with_annos.select("instance.token") 206 | .unique() 207 | .sort("instance.token") 208 | .with_row_count("scene_instance.index", offset=1) 209 | .with_columns(pl.col("scene_instance.index").cast(pl.Int64)) 210 | ) 211 | scene_with_annos = self.ds.join(scene_with_annos, scene_instance_index) 212 | return scene, scene_with_annos, scene_instance_index 213 | 214 | def save_scene_flow_polars(self, scene_name: str, sample: pl.DataFrame): 215 | """Saves a single sample.""" 216 | assert len(sample) == 1 217 | filename = self.save_path(scene_name, sample) 218 | filename.parent.mkdir(exist_ok=True, parents=True) 219 | # remove batch dim 220 | sample.write_ipc(filename, compression="zstd") 221 | 222 | def save_path(self, scene_name, sample: pl.DataFrame): 223 | """Save from scene name and lidar token.""" 224 | assert len(sample) == 1 225 | filename = ( 226 | Path(self.extra_data_root) 227 | / "scene_flow" 228 | / scene_name 229 | / "LIDAR_TOP" 230 | / f"{sample.item(0, 'LIDAR_TOP.sample_data.token')}.arrow" 231 | ) 232 | return filename 233 | 234 | def process_scene(self, scene_name: str): 235 | """Process a single scene.""" 236 | scene, scene_with_annos, scene_instance_index = self.load_scene(scene_name) 237 | scene = build_scene_instance_motion(self.ds, scene, scene_with_annos, scene_instance_index) 238 | ego_volume = Volume.new_volume(lower=self.cartesian_lower, upper=self.cartesian_upper) 239 | volume_lower = torch_to_series("LIDAR_TOP.scene_flow.volume.lower", ego_volume.lower) 240 | volume_upper = torch_to_series("LIDAR_TOP.scene_flow.volume.upper", ego_volume.upper) 241 | 242 | for sample in tqdm.tqdm(scene.iter_slices(1), total=len(scene), position=1): 243 | if self.missing_only: 244 | filename = self.save_path(scene_name, sample) 245 | if filename.exists(): 246 | continue 247 | # sample has annotations 248 | if sample["LIDAR_TOP.sample_data.token"].is_in(scene_with_annos["LIDAR_TOP.sample_data.token"]).item(): 249 | sample_with_annos = self.ds.join(sample.select("LIDAR_TOP.sample_data.token"), scene_with_annos) 250 | sample_with_ids = extract_instance_ids_chunked( 251 | sample, sample_with_annos, ego_volume, self.cartesian_shape, device=self.device 252 | ) 253 | else: 254 | # everything is background 255 | sample_with_ids = sample.select("LIDAR_TOP.sample_data.token").with_columns( 256 | torch_to_series( 257 | "LIDAR_TOP.scene_flow.scene_instance_index", 258 | torch.zeros(1, *self.cartesian_shape, dtype=torch.long), 259 | ) 260 | ) 261 | 262 | sample = self.ds.join( 263 | sample.select( 264 | "LIDAR_TOP.sample_data.token", 265 | "LIDAR_TOP.scene_flow.instance_from_global", 266 | "LIDAR_TOP.scene_flow.instance_from_global_mask", 267 | ), 268 | sample_with_ids.select( 269 | "LIDAR_TOP.sample_data.token", 270 | "LIDAR_TOP.scene_flow.scene_instance_index", 271 | ), 272 | ) 273 | sample = sample.with_columns(volume_lower, volume_upper) 274 | self.save_scene_flow_polars(scene_name, sample) 275 | 276 | return scene 277 | 278 | def process_data(self): 279 | """Process dataset.""" 280 | self.ds.load_sensor_sample_annotation("LIDAR_TOP") 281 | for scene_name in tqdm.tqdm(self.ds.scene["scene.name"], position=0, desc="Calculate scene flow from boxes"): 282 | self.process_scene(scene_name) 283 | 284 | 285 | def deform_grid(scene: pl.DataFrame, volume: Volume): 286 | """Deform grid considering dynmic objects.""" 287 | instance_ids = series_to_torch(scene["LIDAR_TOP.scene_flow.scene_instance_index"]) 288 | instance_from_global = series_to_torch(scene["LIDAR_TOP.scene_flow.instance_from_global"]) 289 | global_from_ego = series_to_torch(scene["LIDAR_TOP.transform.global_from_ego"]) 290 | instance_from_ego = torch.einsum("bnig,bge->bnie", instance_from_global, global_from_ego) 291 | instance_from_ego_src = instance_from_ego[1:] 292 | instance_from_ego_tgt = instance_from_ego[:-1] 293 | 294 | instance_from_mask = series_to_torch(scene["LIDAR_TOP.scene_flow.instance_from_global_mask"]) 295 | instance_from_src_mask = instance_from_mask[1:] 296 | instance_from_tgt_mask = instance_from_mask[:-1] 297 | instance_ids_src = instance_ids[1:] 298 | tgt_from_src = lookup_instance_transform( 299 | instance_ids_src, 300 | instance_from_ego_tgt, 301 | instance_from_ego_src, 302 | instance_from_tgt_mask, 303 | instance_from_src_mask, 304 | ) 305 | new_grid_points = volume.new_coord_grid(instance_ids.shape[-3:]) 306 | warped_grid_points = einsum_transform("bxyzon,bxyzn->bxyzo", tgt_from_src, points=new_grid_points) 307 | return warped_grid_points 308 | -------------------------------------------------------------------------------- /scene_reconstruction/occupancy/temporal_transmission_and_reflection.py: -------------------------------------------------------------------------------- 1 | """Temporal accumulation of transmission and reflection.""" 2 | 3 | from dataclasses import dataclass, field 4 | from pathlib import Path 5 | from typing import Any, Optional, Union 6 | 7 | import polars as pl 8 | import torch 9 | import tqdm 10 | from torch import Tensor 11 | 12 | from scene_reconstruction.core import einsum_transform 13 | from scene_reconstruction.core.volume import Volume 14 | from scene_reconstruction.data.nuscenes.dataset import NuscenesDataset 15 | from scene_reconstruction.data.nuscenes.polars_helpers import series_to_torch, torch_to_series 16 | from scene_reconstruction.data.nuscenes.scene_utils import scene_to_tensor_dict 17 | from scene_reconstruction.math.spherical_coordinate_system import ( 18 | cartesian_to_spherical, 19 | spherical_volume_element_center_and_voxel_size, 20 | ) 21 | 22 | 23 | def _load_flow_and_rt(ds: NuscenesDataset, scene: pl.DataFrame, extra_data_root: Union[Path, str]): 24 | scene = ds.load_reflection_and_transmission_spherical(scene, extra_data_root) 25 | scene = ds.load_scene_flow_polars(scene, extra_data_root) 26 | return scene 27 | 28 | 29 | class AccumulateFrames: 30 | """Accumlate reflection and transmission for reference frame.""" 31 | 32 | def __init__( 33 | self, 34 | ds: NuscenesDataset, 35 | ref_frame: pl.DataFrame, 36 | extra_data_root: Union[Path, str], 37 | icp_alignment: bool = False, 38 | batch_size: int = 2, 39 | max_num_frames: int = 50, 40 | max_ego_pose_difference: float = 20.0, 41 | device: str = "cuda", 42 | num_threads: int = 8, 43 | ) -> None: 44 | """Initialize refernce frame.""" 45 | self.extra_data_root = extra_data_root 46 | self.kwargs = {"device": device, "non_blocking": True} 47 | self.ds = ds 48 | self.use_icp_alignment = icp_alignment 49 | self.frame_dict = {k: v.to(**self.kwargs) for k, v in self._load_fn(ref_frame).items()} 50 | self.instance_ids = self.frame_dict["frame_instance_ids"] 51 | self.instance_from_global = self.frame_dict["frame_instance_from_global"] 52 | self.global_from_ego = series_to_torch(ref_frame["LIDAR_TOP.transform.global_from_ego"]).to(**self.kwargs) 53 | self.instance_from_ego = torch.einsum("bnig,bge->bnie", self.instance_from_global, self.global_from_ego) 54 | self.sph_volume = Volume( 55 | lower=series_to_torch(ref_frame["LIDAR_TOP.reflection_and_transmission_spherical.volume.lower"]).to( 56 | **self.kwargs 57 | ), 58 | upper=series_to_torch(ref_frame["LIDAR_TOP.reflection_and_transmission_spherical.volume.upper"]).to( 59 | **self.kwargs 60 | ), 61 | ) 62 | self.ego_volume = Volume( 63 | lower=series_to_torch(ref_frame["LIDAR_TOP.scene_flow.volume.lower"]).to(**self.kwargs), 64 | upper=series_to_torch(ref_frame["LIDAR_TOP.scene_flow.volume.upper"]).to(**self.kwargs), 65 | ) 66 | 67 | self.max_ego_pose_difference = max_ego_pose_difference 68 | self.max_num_frames = max_num_frames 69 | self.batch_size = batch_size 70 | self.num_threads = num_threads 71 | 72 | self._init_ref_frame() 73 | 74 | def _init_ref_frame(self): 75 | sample_points = self._generate_sample_points( 76 | self.frame_dict["frame_instance_ids"], 77 | self.frame_dict["frame_instance_from_global"], 78 | self.frame_dict["frame_ego_from_global"], 79 | self.frame_dict["frame_lidar_from_ego"], 80 | ) 81 | sample_points_spherical = cartesian_to_spherical(sample_points) 82 | spherical_rt = self.frame_dict["spherical_rt"] 83 | cartesian_rt = self.sph_volume.sample_volume(spherical_rt, sample_points_spherical) 84 | spherical_voxel_size = self.sph_volume.voxel_size(spherical_rt) 85 | cart_voxel_size = self.ego_volume.voxel_size(cartesian_rt) 86 | # there might be points containing "nan" values, since they have no valid transformation 87 | scale = ( 88 | cart_voxel_size.prod(-1)[:, None, None, None] 89 | / spherical_volume_element_center_and_voxel_size( 90 | sample_points_spherical, spherical_voxel_size[:, None, None, None, :] 91 | ) 92 | ).nan_to_num(0.0) 93 | cartesian_rt_scaled = scale.unsqueeze(1) * cartesian_rt 94 | assert cartesian_rt_scaled.shape[0] == 1, "Only one ref frame allowed" 95 | self.agg_rt = cartesian_rt_scaled.sum(0, keepdim=True) 96 | self.num_samples = cartesian_rt_scaled.shape[0] 97 | 98 | # def _icp_alignment( 99 | # self, 100 | # ref_points_global: Tensor, 101 | # ref_points_global_mask: Tensor, 102 | # new_points_global: Tensor, 103 | # new_points_global_mask: Tensor, 104 | # ): 105 | # # ref_points_global: [1, N, 3] 106 | # # new_points_global: [B, N, 3] 107 | # # new_points_global_mask: [B, N] 108 | # ref_points = ref_points_global[ref_points_global_mask] # [M, 3] 109 | # ref_from_global_list = [] 110 | # for points_, mask_ in zip(new_points_global, new_points_global_mask): 111 | # points_global = points_[mask_] # [N, 3] 112 | # ref_from_global = register_frame(ref_points, points_global) 113 | # ref_from_global_list.append(ref_from_global) 114 | # ref_from_global = torch.stack(ref_from_global_list) 115 | # return ref_from_global 116 | 117 | # def _icp_ref_from_global(self, scene: pl.DataFrame): 118 | # points_lidar = series_to_torch(scene["LIDAR_TOP.sample_data.points_lidar"]).to(**self.kwargs) 119 | # points_mask = series_to_torch(scene["LIDAR_TOP.sample_data.points_mask"]).to(**self.kwargs) 120 | # global_from_ego = series_to_torch(scene["LIDAR_TOP.transform.global_from_ego"]).to(**self.kwargs) 121 | # ego_from_lidar = series_to_torch(scene["LIDAR_TOP.transform.ego_from_sensor"]).to(**self.kwargs) 122 | # global_from_lidar = global_from_ego @ ego_from_lidar 123 | # points_global = einsum_transform("bgl,bnl->bng", global_from_lidar, points=points_lidar) 124 | # ref_from_global = self._icp_alignment(self.points_global, self.points_mask, points_global, points_mask) 125 | # return ref_from_global 126 | 127 | def _generate_sample_points( 128 | self, 129 | frame_instance_ids: Tensor, 130 | frame_instance_from_global: Tensor, 131 | frame_ego_from_global: Tensor, 132 | frame_lidar_from_ego: Tensor, 133 | ): 134 | frame_global_from_instance = frame_instance_from_global.inverse() 135 | if self.use_icp_alignment: 136 | raise NotImplementedError() 137 | 138 | # build transformations 139 | frame_lidar_from_global = frame_lidar_from_ego @ frame_ego_from_global 140 | frame_ego_from_instance = torch.einsum("beg,bngi->bnei", frame_ego_from_global, frame_global_from_instance) 141 | frame_lidar_from_instance = torch.einsum("blg,bngi->bnli", frame_lidar_from_global, frame_global_from_instance) 142 | 143 | frame_ego_from_self_ego = frame_ego_from_instance @ self.instance_from_ego 144 | frame_lidar_from_self_ego = frame_lidar_from_instance @ self.instance_from_ego 145 | self_instance_ids_flat = self.instance_ids.flatten(1) 146 | 147 | frame_ego_from_self_ego_gathered_flat = frame_ego_from_self_ego.gather( 148 | 1, self_instance_ids_flat[..., None, None].expand(frame_instance_ids.shape[0], -1, 4, 4) 149 | ) 150 | frame_ego_from_self_ego_dense = frame_ego_from_self_ego_gathered_flat.view( 151 | -1, *self.instance_ids.shape[1:], 4, 4 152 | ) 153 | frame_ego_points = einsum_transform( 154 | "bxyzfs,bxyzs->bxyzf", frame_ego_from_self_ego_dense, points=self.ego_volume.coord_grid(frame_instance_ids) 155 | ) 156 | # sample ids to check for valid transforms 157 | self_instance_ids_sampled = self.ego_volume.sample_volume( 158 | frame_instance_ids.unsqueeze(1).float(), frame_ego_points, mode="nearest", fill_invalid=float("nan") 159 | ) 160 | valid_transform = self_instance_ids_sampled.squeeze(1) == self.instance_ids.float() # [B, X, Y, Z] 161 | 162 | frame_lidar_from_self_ego_gathered_flat = frame_lidar_from_self_ego.gather( 163 | 1, self_instance_ids_flat[..., None, None].expand(frame_instance_ids.shape[0], -1, 4, 4) 164 | ) 165 | frame_lidar_from_self_ego_dense = frame_lidar_from_self_ego_gathered_flat.view( 166 | -1, *self.instance_ids.shape[1:], 4, 4 167 | ) 168 | frame_lidar_points = einsum_transform( 169 | "bxyzfs,bxyzs->bxyzf", 170 | frame_lidar_from_self_ego_dense, 171 | points=self.ego_volume.coord_grid(frame_instance_ids), 172 | ) 173 | 174 | frame_lidar_points_valid = torch.where(valid_transform[..., None], frame_lidar_points, float("nan")) 175 | 176 | return frame_lidar_points_valid 177 | 178 | def _load_fn(self, frame: pl.DataFrame): 179 | frame = _load_flow_and_rt(self.ds, frame, self.extra_data_root) 180 | frame_dict = scene_to_tensor_dict( 181 | frame, 182 | { 183 | "frame_instance_ids": "LIDAR_TOP.scene_flow.scene_instance_index", 184 | "frame_instance_from_global": "LIDAR_TOP.scene_flow.instance_from_global", 185 | "frame_ego_from_global": "LIDAR_TOP.transform.ego_from_global", 186 | "frame_lidar_from_ego": "LIDAR_TOP.transform.sensor_from_ego", 187 | "spherical_rt": "LIDAR_TOP.reflection_and_transmission_spherical", 188 | }, 189 | ) 190 | return frame_dict 191 | 192 | def _accumulate_frames( 193 | self, 194 | frame_instance_ids: Tensor, 195 | frame_instance_from_global: Tensor, 196 | frame_ego_from_global: Tensor, 197 | frame_lidar_from_ego: Tensor, 198 | spherical_rt: Tensor, 199 | ): 200 | sample_points = self._generate_sample_points( 201 | frame_instance_ids, 202 | frame_instance_from_global, 203 | frame_ego_from_global, 204 | frame_lidar_from_ego, 205 | ) 206 | sample_points_spherical = cartesian_to_spherical(sample_points) 207 | cartesian_rt = self.sph_volume.sample_volume(spherical_rt, sample_points_spherical) 208 | spherical_voxel_size = self.sph_volume.voxel_size(spherical_rt) 209 | cart_voxel_size = self.ego_volume.voxel_size(cartesian_rt) 210 | # there might be points containing "nan" values, since they have no valid transformation 211 | scale = ( 212 | cart_voxel_size.prod(-1)[:, None, None, None] 213 | / spherical_volume_element_center_and_voxel_size( 214 | sample_points_spherical, spherical_voxel_size[:, None, None, None, :] 215 | ) 216 | ).nan_to_num(0.0) 217 | cartesian_rt_scaled = scale.unsqueeze(1) * cartesian_rt 218 | 219 | self.agg_rt += cartesian_rt_scaled.sum(0, keepdim=True) 220 | self.num_samples += cartesian_rt_scaled.shape[0] 221 | 222 | def process_frames_in_radius(self, scene: pl.DataFrame): 223 | """Process frames in radius.""" 224 | scene = scene.sort("LIDAR_TOP.sample_data.timestamp") 225 | frame_ego_from_global = series_to_torch(scene["LIDAR_TOP.transform.ego_from_global"]) 226 | frame_ego_from_self_ego = frame_ego_from_global @ self.global_from_ego.cpu() 227 | pos_diff = frame_ego_from_self_ego[:, :3, 3].norm(dim=-1) 228 | frams_in_radius = pos_diff < self.max_ego_pose_difference 229 | scene = scene.with_columns(torch_to_series("frame_in_radius", frams_in_radius)) 230 | scene = scene.filter(pl.col("frame_in_radius")) 231 | if len(scene) > self.max_num_frames: 232 | indices = torch.linspace(-0.5, len(scene) + 0.5, self.max_num_frames).round().long() 233 | scene = scene.with_row_count("index").filter( 234 | pl.col("index").is_in(torch_to_series("select_indices", indices)) 235 | ) 236 | 237 | items = scene.select( 238 | "LIDAR_TOP.transform.ego_from_global", 239 | "LIDAR_TOP.transform.sensor_from_ego", 240 | "scene.name", 241 | "LIDAR_TOP.sample_data.token", 242 | ).iter_slices(self.batch_size) 243 | for sample in items: 244 | frame_dict = self._load_fn(sample) 245 | frame_dict = {k: v.to(**self.kwargs) for k, v in frame_dict.items()} 246 | self._accumulate_frames(**frame_dict) 247 | 248 | 249 | @dataclass 250 | class TemporalTransmissionAndReflection: 251 | """Temporal accumulation of reflection and transmission using scene flow info.""" 252 | 253 | ds: NuscenesDataset 254 | extra_data_root: Path 255 | reference_keyframes_only: bool = True 256 | frame_accumulation_kwargs: dict[str, Any] = field(default_factory=dict) 257 | missing_only: bool = False 258 | scene_offset: int = 0 259 | num_scenes: Optional[int] = None 260 | 261 | def save_accumulated_sample(self, scene_name: str, sample: pl.DataFrame): 262 | """Saves a single sample.""" 263 | assert len(sample) == 1 264 | filename = self.save_path(scene_name, sample) 265 | filename.parent.mkdir(exist_ok=True, parents=True) 266 | # remove batch dim 267 | sample.write_ipc(filename, compression="zstd") 268 | 269 | def save_path(self, scene_name: str, sample: pl.DataFrame): 270 | """Save path.""" 271 | assert len(sample) == 1 272 | filename = ( 273 | Path(self.extra_data_root) 274 | / "reflection_and_transmission_multi_frame" 275 | / scene_name 276 | / "LIDAR_TOP" 277 | / f"{sample.item(0, 'LIDAR_TOP.sample_data.token')}.arrow" 278 | ) 279 | return filename 280 | 281 | def process_scene(self, scene: pl.DataFrame): 282 | """Processes single scene.""" 283 | scene = self.ds.join(scene, self.ds.sample) 284 | scene = self.ds.load_sample_data(scene, "LIDAR_TOP", with_data=False) 285 | scene = self.ds.sort_by_time(scene) 286 | 287 | if self.reference_keyframes_only: 288 | reference_frames = scene.filter(pl.col("LIDAR_TOP.sample_data.is_key_frame")) 289 | else: 290 | reference_frames = scene 291 | for reference_frame in tqdm.tqdm(reference_frames.iter_slices(1), total=len(reference_frames), position=1): 292 | if self.missing_only: 293 | filename = self.save_path(reference_frame["scene.name"].item(), reference_frame) 294 | if filename.exists(): 295 | continue 296 | 297 | reference_frame = self.ds.load_reflection_and_transmission_spherical(reference_frame, self.extra_data_root) 298 | reference_frame = self.ds.load_scene_flow_polars(reference_frame, self.extra_data_root) 299 | 300 | agg_frames = AccumulateFrames( 301 | self.ds, 302 | ref_frame=reference_frame, 303 | extra_data_root=self.extra_data_root, 304 | **self.frame_accumulation_kwargs, 305 | ) 306 | 307 | agg_frames.process_frames_in_radius(scene) 308 | assert agg_frames.agg_rt is not None, "No frames accumulated" 309 | agg_frames_mean = agg_frames.agg_rt / agg_frames.num_samples 310 | data = ( 311 | reference_frame.select( 312 | "LIDAR_TOP.sample_data.token", 313 | "LIDAR_TOP.scene_flow.volume.lower", 314 | "LIDAR_TOP.scene_flow.volume.upper", 315 | ) 316 | .rename( 317 | { 318 | "LIDAR_TOP.scene_flow.volume.lower": "LIDAR_TOP.reflection_and_transmission_multi_frame.volume.lower", 319 | "LIDAR_TOP.scene_flow.volume.upper": "LIDAR_TOP.reflection_and_transmission_multi_frame.volume.upper", 320 | } 321 | ) 322 | .with_columns( 323 | torch_to_series("LIDAR_TOP.reflection_and_transmission_multi_frame", agg_frames_mean.cpu()) 324 | ) 325 | ) 326 | self.save_accumulated_sample(reference_frame["scene.name"].item(), data) 327 | 328 | def process_data(self) -> None: 329 | """Process dataset.""" 330 | # load all sample data with lidar token 331 | scene_to_process = self.ds.scene.slice(self.scene_offset, self.num_scenes) 332 | for scene in (tbar := tqdm.tqdm(scene_to_process.iter_slices(1), total=len(scene_to_process), position=0)): 333 | scene_name = scene["scene.name"].item() 334 | tbar.set_description_str(f"Processing {scene_name}") 335 | self.process_scene(scene) 336 | -------------------------------------------------------------------------------- /scene_reconstruction/eval/lidar_depth.py: -------------------------------------------------------------------------------- 1 | """Lidar rendering.""" 2 | import math 3 | from collections.abc import Sequence 4 | from datetime import datetime 5 | from pathlib import Path 6 | from typing import Any, Optional 7 | 8 | import einops 9 | import kaolin 10 | import polars as pl 11 | import torch 12 | import torch.nn as nn 13 | import tqdm 14 | from kaolin.ops.spc.points import morton_to_points, points_to_morton, unbatched_points_to_octree 15 | from kaolin.rep.spc import Spc 16 | from torch import Tensor 17 | from torchmetrics.classification import BinaryAccuracy 18 | from torchmetrics.collections import MetricCollection 19 | from torchmetrics.regression import MeanAbsoluteError, MeanSquaredError 20 | from typing_extensions import Literal 21 | 22 | from scene_reconstruction.core import einsum_transform 23 | from scene_reconstruction.core.volume import Volume 24 | from scene_reconstruction.data.nuscenes.dataset import NuscenesDataset 25 | from scene_reconstruction.data.nuscenes.scene_utils import scene_to_tensor_dict 26 | from scene_reconstruction.math.dempster_shafer import belief_from_reflection_and_transmission_stacked 27 | 28 | 29 | def unbatched_spc_from_volume(features: Tensor, volume: Volume, mask_value=0): 30 | """Converts volume to structured point cloud.""" 31 | assert features.shape[0] == 1 32 | assert features.is_cuda 33 | volume_grid = volume.coord_grid(features) # [1, X, Y, Z, 3] 34 | # max next power of 2 for octree level 35 | level = max([int(math.ceil(math.log2(x))) for x in volume.volume_shape(features)]) 36 | next_power_of_2 = 2**level 37 | volume_shape = torch.tensor([volume.volume_shape(features)], device=features.device) 38 | normed_limits = volume_shape / next_power_of_2 39 | 40 | # spc voxel index coords, range from [0, 0, 0] to [X, Y, Z] 41 | lower_voxel_spc = torch.zeros_like(volume.lower) 42 | upper_voxel_spc = volume_shape 43 | spc_voxel_volume = Volume(lower=lower_voxel_spc, upper=upper_voxel_spc) 44 | spc_voxel_volume_from_ego = spc_voxel_volume.self_from_other(volume) 45 | grid_points_spc_voxel_volume = einsum_transform("bse,bxyze->bxyzs", spc_voxel_volume_from_ego, points=volume_grid) 46 | grid_points_quantized = torch.floor(torch.clamp(grid_points_spc_voxel_volume, 0, next_power_of_2 - 1)).short() 47 | 48 | # spc normalized coords, range from -1.0 to 1.0 49 | lower_norm_spc = -torch.ones_like(volume.lower) 50 | upper_norm_spc = -torch.ones_like(volume.upper) + 2 * normed_limits 51 | spc_norm_volume = Volume(lower=lower_norm_spc, upper=upper_norm_spc) 52 | spc_norm_volume_from_ego = spc_norm_volume.self_from_other(volume) 53 | assert ( 54 | spc_norm_volume_from_ego[0, 0, 0] == spc_norm_volume_from_ego[0, 1, 1] == spc_norm_volume_from_ego[0, 2, 2] 55 | ), "voxel size must be uniform across dims" 56 | scale = 1 / spc_norm_volume_from_ego[0, 0, 0].item() 57 | 58 | non_zero_mask = (features != mask_value).any(1) # [B, X, Y, Z] 59 | 60 | features_flat = einops.rearrange(features, "b c x y z -> b x y z c")[non_zero_mask] # [N, C] 61 | points_quantized_flat = grid_points_quantized[non_zero_mask] # [N, 3] 62 | 63 | # sort points and features 64 | morton, keys = torch.sort(points_to_morton(points_quantized_flat.contiguous()).contiguous()) 65 | points_quantized_flat = morton_to_points(morton.contiguous()) 66 | octree = unbatched_points_to_octree(points_quantized_flat, level, sorted=True) 67 | features_flat = features_flat[keys] 68 | 69 | # A full SPC requires octree hierarchy + auxilary data structures 70 | lengths = torch.tensor([len(octree)], dtype=torch.int32) # Single entry batch 71 | return ( 72 | Spc(octrees=octree, lengths=lengths, features=features_flat, max_level=level), 73 | spc_norm_volume_from_ego[0], 74 | scale, 75 | ) 76 | 77 | 78 | def unbatched_raytrace_lidar_transmissions( 79 | occupancies: Tensor, # [B, C, X, Y, Z] 80 | volume: Volume, 81 | start_points: Tensor, # [N, 3] 82 | end_points: Tensor, # [N, 3] 83 | mask_value=0, 84 | ): 85 | """Render lidar in occupancy volume.""" 86 | spc, spc_from_ego, scale = unbatched_spc_from_volume(occupancies, volume=volume, mask_value=mask_value) 87 | start_points_spc = einsum_transform("se,ne->ns", spc_from_ego, points=start_points) 88 | end_points_spc = einsum_transform("se,ne->ns", spc_from_ego, points=end_points) 89 | ray_direction = end_points_spc - start_points_spc 90 | ray_length: Tensor = ray_direction.norm(dim=-1, keepdim=True) 91 | ray_direction_normalized = ray_direction / ray_length 92 | num_rays = start_points_spc.shape[0] 93 | ridx: Tensor 94 | pidx: Tensor 95 | depth: Tensor 96 | ridx, pidx, depth = kaolin.render.spc.unbatched_raytrace( # type: ignore 97 | spc.octrees, 98 | spc.point_hierarchies, 99 | spc.pyramids[0], 100 | spc.exsum, 101 | start_points_spc + 2**-spc.max_level, # this offset seems to fix points at [0.0, 0.0, 0.0] bug? 102 | ray_direction_normalized, 103 | spc.max_level, 104 | return_depth=True, 105 | with_exit=False, 106 | ) 107 | 108 | first_hits_mask: Tensor = kaolin.render.spc.mark_pack_boundaries(ridx) 109 | first_hit_ridx = ridx[first_hits_mask] 110 | first_hit_depth = torch.full((num_rays, 1), fill_value=float("inf"), device=depth.device) 111 | first_hit_depth.index_copy_(0, first_hit_ridx.long(), depth[first_hits_mask]) 112 | 113 | before_lidar_reflection = (depth < ray_length.index_select(0, ridx)).float() 114 | num_transmissions_per_ray = torch.full((num_rays, 1), fill_value=0.0, device=depth.device) 115 | num_transmissions_per_ray.index_add_(0, ridx, before_lidar_reflection) 116 | 117 | num_transmissions_per_voxel = torch.full_like(spc.features[:, :1], fill_value=0.0, dtype=torch.float) 118 | num_transmissions_per_voxel.index_add_( 119 | 0, 120 | pidx - spc.pyramids[0, 1, spc.max_level], 121 | before_lidar_reflection, 122 | ) 123 | 124 | spc_num_transmissions = Spc( 125 | octrees=spc.octrees, 126 | max_level=spc.max_level, 127 | lengths=spc.lengths, 128 | features=num_transmissions_per_voxel, 129 | ) 130 | 131 | first_hit_depth_unscaled = first_hit_depth * scale 132 | 133 | return first_hit_depth_unscaled, num_transmissions_per_ray, spc_num_transmissions 134 | 135 | 136 | class RenderLidarDepth(nn.Module): 137 | """Lidar rendering.""" 138 | 139 | def __init__( 140 | self, 141 | volume: Volume, 142 | eval_volume_ego: Volume, 143 | free_index: int = 17, 144 | min_distance: float = 2.5, 145 | volume_frame: Literal["ego", "lidar"] = "ego", 146 | ) -> None: 147 | """Lidar rendering in ego volume.""" 148 | super().__init__() 149 | self.volume = volume 150 | self.free_index = free_index 151 | self.min_distance = min_distance 152 | self.max_distance = (eval_volume_ego.upper - eval_volume_ego.lower).norm(dim=-1) 153 | self.volume_frame = volume_frame 154 | self.eval_volume_ego = eval_volume_ego 155 | 156 | def forward( 157 | self, 158 | occupancy: Tensor, 159 | points_lidar: Tensor, 160 | points_mask: Tensor, 161 | ego_from_lidar: Tensor, 162 | ): 163 | """Rendering.""" 164 | if self.volume_frame == "ego": 165 | return self._render_in_ego_volume(occupancy, points_lidar, points_mask, ego_from_lidar) 166 | elif self.volume_frame == "lidar": 167 | return self._render_in_lidar_volume(occupancy, points_lidar, points_mask, ego_from_lidar) 168 | else: 169 | raise NotImplementedError() 170 | 171 | def _render_in_ego_volume( 172 | self, occupancy: Tensor, points_lidar: Tensor, points_mask: Tensor, ego_from_lidar: Tensor 173 | ): 174 | """Rendering in ego volume.""" 175 | occupancy_with_channel = occupancy.unsqueeze(1) 176 | start_points_ego = einsum_transform( 177 | "bel,bnl->bne", ego_from_lidar, points=torch.zeros_like(points_lidar) 178 | ) # lidar origin 179 | end_points_ego = einsum_transform("bel,bnl->bne", ego_from_lidar, points=points_lidar) 180 | 181 | rendered_depths = torch.full_like(start_points_ego[..., 0], fill_value=float("nan")) 182 | num_transmissions = torch.full_like(start_points_ego[..., 0], fill_value=float("nan")) 183 | lidar_depth = points_lidar.norm(dim=-1) 184 | for ub_num_t, ub_rendered_depth, ub_occ, ub_start_points_ego, ub_end_points_ego, ub_masks in zip( 185 | num_transmissions.split(1), 186 | rendered_depths.split(1), 187 | occupancy_with_channel.split(1), 188 | start_points_ego.split(1), 189 | end_points_ego.split(1), 190 | points_mask.split(1), 191 | ): 192 | ( 193 | rendered_depth, 194 | num_transmissions_per_ray, 195 | spc_num_transmissions, 196 | ) = unbatched_raytrace_lidar_transmissions( 197 | occupancies=ub_occ, 198 | volume=self.volume, 199 | start_points=ub_start_points_ego[ub_masks], 200 | end_points=ub_end_points_ego[ub_masks], 201 | mask_value=self.free_index, 202 | ) 203 | ub_rendered_depth[ub_masks] = rendered_depth.squeeze(1) 204 | ub_num_t[ub_masks] = num_transmissions_per_ray.squeeze(1) 205 | # calculate rendered endpoints in ego frame to clamp them to the eval bounds 206 | direction = end_points_ego - start_points_ego 207 | direction /= direction.norm(dim=-1, keepdim=True) 208 | finite_depth = rendered_depths.isfinite() 209 | rendered_depths = rendered_depths.clamp_max(self.max_distance.unsqueeze(-1)) 210 | rendered_end_points_ego = start_points_ego + direction * rendered_depths.unsqueeze(-1) 211 | rendered_end_points_ego_clamped = self.eval_volume_ego.clamp_points_along_line( 212 | rendered_end_points_ego, start_points_ego 213 | ) 214 | # recalculate depth 215 | rendered_depths = (rendered_end_points_ego_clamped - start_points_ego).norm(dim=-1) 216 | points_in_volume = ( 217 | # ego volume check 218 | (end_points_ego >= self.eval_volume_ego.lower.unsqueeze(1)).all(-1) 219 | & (end_points_ego <= self.eval_volume_ego.upper.unsqueeze(1)).all(-1) 220 | & (start_points_ego >= self.eval_volume_ego.lower.unsqueeze(1)).all(-1) 221 | & (start_points_ego <= self.eval_volume_ego.upper.unsqueeze(1)).all(-1) 222 | # lidar points check 223 | & points_mask 224 | & (lidar_depth > self.min_distance) 225 | ) # [B, X, Y, Z] 226 | 227 | finite_depth = finite_depth & points_mask & (lidar_depth > self.min_distance) 228 | 229 | out_dict = { 230 | "rendered_depth": rendered_depths, 231 | "num_transmission": num_transmissions, 232 | "lidar_depth": lidar_depth, 233 | "points_in_ego_volume": points_in_volume, 234 | "finite_depth": finite_depth, 235 | } 236 | 237 | return out_dict 238 | 239 | def _render_in_lidar_volume( 240 | self, occupancy: Tensor, points_lidar: Tensor, points_mask: Tensor, ego_from_lidar: Tensor 241 | ): 242 | """Rendering in lidar volume.""" 243 | occupancy_with_channel = occupancy.unsqueeze(1) 244 | start_points_lidar = torch.zeros_like(points_lidar) # lidar origin 245 | end_points_lidar = points_lidar 246 | 247 | rendered_depths = torch.full_like(start_points_lidar[..., 0], fill_value=float("nan")) 248 | num_transmissions = torch.full_like(start_points_lidar[..., 0], fill_value=float("nan")) 249 | lidar_depth = end_points_lidar.norm(dim=-1) 250 | for ub_num_t, ub_rendered_depth, ub_occ, ub_start_points_ego, ub_end_points_ego, ub_masks in zip( 251 | num_transmissions.split(1), 252 | rendered_depths.split(1), 253 | occupancy_with_channel.split(1), 254 | start_points_lidar.split(1), 255 | end_points_lidar.split(1), 256 | points_mask.split(1), 257 | ): 258 | ( 259 | rendered_depth, 260 | num_transmissions_per_ray, 261 | spc_num_transmissions, 262 | ) = unbatched_raytrace_lidar_transmissions( 263 | occupancies=ub_occ, 264 | volume=self.volume, 265 | start_points=ub_start_points_ego[ub_masks], 266 | end_points=ub_end_points_ego[ub_masks], 267 | mask_value=self.free_index, 268 | ) 269 | ub_rendered_depth[ub_masks] = rendered_depth.squeeze(1) 270 | ub_num_t[ub_masks] = num_transmissions_per_ray.squeeze(1) 271 | start_points_ego = einsum_transform( 272 | "bel,bnl->bne", ego_from_lidar, points=torch.zeros_like(points_lidar) 273 | ) # lidar origin 274 | end_points_ego = einsum_transform("bel,bnl->bne", ego_from_lidar, points=points_lidar) 275 | # calculate rendered endpoints in ego frame to clamp them to the eval bounds 276 | finite_depth = rendered_depths.isfinite() 277 | rendered_depths = rendered_depths.clamp_max(self.max_distance.unsqueeze(-1)) 278 | direction = end_points_ego - start_points_ego 279 | direction /= direction.norm(dim=-1, keepdim=True) 280 | rendered_end_points_ego = start_points_ego + direction * rendered_depths.unsqueeze(-1) 281 | rendered_end_points_ego_clamped = self.eval_volume_ego.clamp_points_along_line( 282 | rendered_end_points_ego, start_points_ego 283 | ) 284 | # recalculate depth 285 | rendered_depths = (rendered_end_points_ego_clamped - start_points_ego).norm(dim=-1) 286 | points_in_volume = ( 287 | # ego volume check 288 | (end_points_ego >= self.eval_volume_ego.lower.unsqueeze(1)).all(-1) 289 | & (end_points_ego <= self.eval_volume_ego.upper.unsqueeze(1)).all(-1) 290 | & (start_points_ego >= self.eval_volume_ego.lower.unsqueeze(1)).all(-1) 291 | & (start_points_ego <= self.eval_volume_ego.upper.unsqueeze(1)).all(-1) 292 | # lidar volume check 293 | & (start_points_lidar >= self.volume.lower.unsqueeze(1)).all(-1) 294 | & (start_points_lidar <= self.volume.upper.unsqueeze(1)).all(-1) 295 | & (end_points_lidar >= self.volume.lower.unsqueeze(1)).all(-1) 296 | & (end_points_lidar <= self.volume.upper.unsqueeze(1)).all(-1) 297 | # lidar points checks 298 | & points_mask 299 | & (lidar_depth > self.min_distance) 300 | ) # [B, X, Y, Z] 301 | 302 | finite_depth = finite_depth & points_mask & (lidar_depth > self.min_distance) 303 | 304 | out_dict = { 305 | "rendered_depth": rendered_depths, 306 | "num_transmission": num_transmissions, 307 | "lidar_depth": lidar_depth, 308 | "points_in_ego_volume": points_in_volume, 309 | "finite_depth": finite_depth, 310 | } 311 | 312 | return out_dict 313 | 314 | 315 | METHODS = Literal["cvpr2023", "preds", "bba", "bba04", "open_occupancy"] 316 | 317 | 318 | class DeltaAccuracy(BinaryAccuracy): 319 | """Delta accuracy.""" 320 | 321 | def __init__( 322 | self, 323 | delta: float = 1.25, 324 | threshold: float = 0.5, 325 | multidim_average: Literal["global", "samplewise"] = "global", 326 | ignore_index: Optional[int] = None, 327 | validate_args: bool = True, 328 | **kwargs: Any, 329 | ) -> None: 330 | """Delta accuracy.""" 331 | super().__init__(threshold, multidim_average, ignore_index, validate_args, **kwargs) 332 | self.delta = delta 333 | 334 | def update(self, pred_depth: Tensor, gt_depth: Tensor): 335 | """Update.""" 336 | thresh = torch.maximum((gt_depth / pred_depth), (pred_depth / gt_depth)) 337 | binary = thresh < self.delta 338 | super().update(binary, torch.ones_like(binary)) 339 | 340 | 341 | class Rmse(MeanSquaredError): 342 | """Root Mean Square Error.""" 343 | 344 | def __init__(self, squared: bool = False, num_outputs: int = 1, **kwargs: Any) -> None: 345 | """Update.""" 346 | super().__init__(squared, num_outputs, **kwargs) 347 | 348 | 349 | class LogRmse(Rmse): 350 | """Log Root Mean Square Error.""" 351 | 352 | def update(self, preds: Tensor, target: Tensor) -> None: 353 | """Update.""" 354 | return super().update(preds.log(), target.log()) 355 | 356 | 357 | class RelAbs(MeanAbsoluteError): 358 | """Relative Absolute Error.""" 359 | 360 | def update(self, preds: Tensor, target: Tensor) -> None: 361 | """Update.""" 362 | return super().update(preds / target, target / target) 363 | 364 | 365 | class LidarDistanceEval: 366 | """Lidar distance evaluation.""" 367 | 368 | def __init__( 369 | self, 370 | ds: NuscenesDataset, 371 | method: METHODS, 372 | split: Literal["train", "val"] = "val", 373 | lower: Sequence[float] = (-40.0, -40.0, -1.0), 374 | upper: Sequence[float] = (40.0, 40.0, 5.4), 375 | volume_frame: Literal["ego", "lidar"] = "ego", 376 | min_distance: float = 2.5, 377 | batch_size: int = 1, 378 | p_fn: float = 0.8, 379 | p_fp: float = 0.05, 380 | eval_ego_lower: Sequence[float] = (-40.0, -40.0, -1.0), 381 | eval_ego_upper: Sequence[float] = (40.0, 40.0, 5.4), 382 | save_path: Optional[str] = None, 383 | pred_method: Optional[str] = None, 384 | occ_threshold: Optional[int] = None, 385 | ) -> None: 386 | """Lidar distance evaluation.""" 387 | self.ds = ds 388 | self.split = split 389 | self.volume = Volume.new_volume(lower=lower, upper=upper).cuda() 390 | self.eval_volume_ego = Volume.new_volume(lower=eval_ego_lower, upper=eval_ego_upper).cuda() 391 | self.volume_frame = volume_frame 392 | self.render_lidar_depth = RenderLidarDepth( 393 | self.volume, 394 | min_distance=min_distance, 395 | volume_frame=self.volume_frame, 396 | eval_volume_ego=self.eval_volume_ego, 397 | ).cuda() 398 | self.method: METHODS = method 399 | self.save_path = save_path 400 | self.pred_method = pred_method 401 | self.occ_threshold = occ_threshold 402 | 403 | metrics = MetricCollection( 404 | { 405 | "rmse": Rmse(squared=False), 406 | "log_rmse": LogRmse(), 407 | "mae": MeanAbsoluteError(), 408 | "delta_1_25": DeltaAccuracy(1.25), 409 | "delta_1_25_2": DeltaAccuracy(1.25**2), 410 | "delta_1_25_3": DeltaAccuracy(1.25**3), 411 | "rel_abs": RelAbs(), 412 | } 413 | ) 414 | self.metrics = metrics.to(self.volume.device) 415 | self.batch_size = batch_size 416 | 417 | self.p_fn = p_fn 418 | self.p_fp = p_fp 419 | 420 | def save_results(self, results: dict): 421 | """Save results to json.""" 422 | if self.save_path is not None: 423 | time_str = datetime.now().strftime("%Y_%m_%d-%H_%M_%S") 424 | filename = Path(self.save_path) / f"results-{time_str}.json" 425 | filename.parent.mkdir(exist_ok=True, parents=True) 426 | results_df = pl.DataFrame(results) 427 | results_df.write_json(filename) 428 | 429 | def eval(self): 430 | """Evaluation of dataset.""" 431 | scenes = self.ds.scenes_by_split(self.split) 432 | for scene in tqdm.tqdm(scenes.iter_slices(1), total=len(scenes)): 433 | scene = self.ds.join(scene, self.ds.sample) 434 | self.process_scene(scene) 435 | results = { 436 | "method": self.method, 437 | "version": self.ds.version, 438 | "split": self.split, 439 | "occ_threshold": self.occ_threshold, 440 | } 441 | if self.method in ["bba", "bba04"]: 442 | results["p_fn"] = self.p_fn 443 | results["p_fp"] = self.p_fp 444 | metrics = self.metrics.compute() 445 | metrics = {k: v.item() for k, v in metrics.items()} 446 | results.update(metrics) 447 | self.save_results(results) 448 | return results 449 | 450 | def process_scene(self, scene: pl.DataFrame): 451 | """Process single scene.""" 452 | for sample in scene.iter_slices(self.batch_size): 453 | sample = self.ds.load_sample_data(sample, "LIDAR_TOP") 454 | if self.method == "cvpr2023": 455 | sample = self.ds.load_cvpr2023_occupancy(sample) 456 | elif self.method == "bba" or self.method == "bba04": 457 | sample = self.ds.load_reflection_and_transmission_multi_frame(sample) 458 | elif self.method == "bba" or self.method == "bba04": 459 | sample = self.ds.load_reflection_and_transmission_multi_frame(sample) 460 | elif self.method == "open_occupancy": 461 | sample = self.ds.load_open_occpancy(sample) 462 | elif self.method == "surround_occ": 463 | sample = self.ds.load_surround_occ_occupancy(sample) 464 | elif self.method == "scene_as_occupancy": 465 | sample = self.ds.load_scene_as_occupancy_gt(sample) 466 | elif self.method == "preds": 467 | assert self.pred_method is not None 468 | sample = self.ds.load_occupancy_preds(sample, self.pred_method) 469 | self.process_sample(sample) 470 | 471 | def process_sample(self, sample: pl.DataFrame): 472 | """Process sample.""" 473 | scene_dict = self.load_sample(sample) 474 | rendered_depth_dict = self.render_lidar_depth( 475 | occupancy=scene_dict["occupancy"], 476 | points_lidar=scene_dict["points_lidar"], 477 | points_mask=scene_dict["points_mask"], 478 | ego_from_lidar=scene_dict["ego_from_lidar"], 479 | ) 480 | mask = rendered_depth_dict["points_in_ego_volume"] 481 | rendered_depth: Tensor = rendered_depth_dict["rendered_depth"][mask] 482 | lidar_depth: Tensor = rendered_depth_dict["lidar_depth"][mask] 483 | assert rendered_depth.isfinite().all(), "rendered depth not finite" 484 | assert lidar_depth.isfinite().all(), "lidar depth not finite" 485 | 486 | self.metrics.update(rendered_depth, lidar_depth) 487 | 488 | def load_sample(self, sample: pl.DataFrame): 489 | """Load sample.""" 490 | mapping = { 491 | "points_lidar": "LIDAR_TOP.sample_data.points_lidar", 492 | "points_mask": "LIDAR_TOP.sample_data.points_mask", 493 | "ego_from_lidar": "LIDAR_TOP.transform.ego_from_sensor", 494 | } 495 | if self.method == "cvpr2023": 496 | mapping["occupancy"] = "sample.occ_gt.semantics" 497 | elif self.method == "preds": 498 | mapping["occupancy"] = f"sample.{self.pred_method}.arr_0" 499 | elif self.method == "open_occupancy" or self.method == "open_occupancy04": 500 | mapping["occupancy"] = "sample.open_occupancy.semantics" 501 | elif self.method == "surround_occ": 502 | mapping["occupancy"] = "sample.surround_occ.semantics" 503 | elif self.method == "scene_as_occupancy": 504 | mapping["occupancy"] = "sample.scene_as_occupancy.semantics" 505 | elif self.method == "bba" or self.method == "bba04": 506 | mapping["reflection_and_transmission"] = "LIDAR_TOP.reflection_and_transmission_multi_frame" 507 | scene_dict = scene_to_tensor_dict(sample, mapping=mapping, device=self.volume.device) 508 | if self.occ_threshold is not None: 509 | scene_dict["occupancy"] = torch.where( 510 | scene_dict["occupancy"] >= self.occ_threshold, 0, 17 511 | ) # 17: free, otherwise occ 512 | if self.method == "bba": 513 | bba = belief_from_reflection_and_transmission_stacked( 514 | scene_dict["reflection_and_transmission"], p_fn=self.p_fn, p_fp=self.p_fp 515 | ) 516 | scene_dict["occupancy"] = torch.where(bba[:, 0] > bba[:, 1], 0, 17) 517 | if self.method == "bba04": 518 | rt = einops.reduce( 519 | scene_dict["reflection_and_transmission"], "b c (x 2) (y 2) (z 2) -> b c x y z", reduction="sum" 520 | ) 521 | bba04 = belief_from_reflection_and_transmission_stacked(rt, p_fn=self.p_fn, p_fp=self.p_fp) 522 | scene_dict["occupancy"] = torch.where(bba04[:, 0] > bba04[:, 1], 0, 17) 523 | return scene_dict 524 | -------------------------------------------------------------------------------- /scene_reconstruction/data/nuscenes/dataset.py: -------------------------------------------------------------------------------- 1 | """NuScenes Dataset.""" 2 | 3 | import math 4 | from collections.abc import Sequence 5 | from pathlib import Path 6 | from typing import Literal, Optional, Union, get_args 7 | 8 | import numpy as np 9 | import polars as pl 10 | import tqdm 11 | from PIL import Image 12 | from polars.type_aliases import JoinStrategy, JoinValidation 13 | 14 | from .parse import ANNOTATION_FILES, SELECTION 15 | from .polars_helpers import ( 16 | join_on_token, 17 | numpy_to_series, 18 | pad_intrinsics_from_colums, 19 | series_to_numpy, 20 | transform_from_columns, 21 | ) 22 | from .typing import CAMERA_CHANNELS, LIDAR_CHANNELS, RADAR_CHANNELS, SENSOR_CHANNELS 23 | 24 | # pylint: disable=C0103,R0902 25 | 26 | CAMERAS = get_args(CAMERA_CHANNELS) 27 | LIDARS = get_args(LIDAR_CHANNELS) 28 | ALL_CHANNELS = get_args(SENSOR_CHANNELS) 29 | 30 | VERSION = Literal["v1.0-mini", "v1.0-trainval"] 31 | 32 | 33 | class NuscenesDataset: 34 | """Nuscenes Dataset.""" 35 | 36 | def __init__( 37 | self, 38 | data_root: Union[Path, str] = "./data/nuscenes/", 39 | extra_data_root: Optional[Union[Path, str]] = None, 40 | version: VERSION = "v1.0-mini", 41 | verbose: bool = False, 42 | key_frames_only: bool = True, 43 | with_sample_annotation: bool = False, 44 | with_lidarseg: bool = False, 45 | with_panoptic: bool = False, 46 | ) -> None: 47 | """Initializes NuScenes dataset from data root.""" 48 | self.data_root = Path(data_root) 49 | self.extra_data_root = Path(extra_data_root) if extra_data_root is not None else None 50 | self.version: VERSION = version 51 | self.key_frames_only = key_frames_only 52 | 53 | # load annotation files 54 | # annotations schema: https://www.nuscenes.org/nuscenes#data-annotation 55 | self.log = self.read_json("log") 56 | self.scene = self.read_json("scene") 57 | self.instance = self.read_json("instance") 58 | self.category = self.read_json("category") 59 | self.map = self.read_json("map") 60 | self.sample = self.read_json("sample") 61 | self.calibrated_sensor = self.read_json("calibrated_sensor") 62 | self.sample_data = self.read_json("sample_data") 63 | self.sample_annotation = self.read_json("sample_annotation") 64 | self.attribute = self.read_json("attribute") 65 | self.sensor = self.read_json("sensor") 66 | self.ego_pose = self.read_json("ego_pose") 67 | self.visibility = self.read_json("visibility") 68 | if with_lidarseg: 69 | self.lidarseg = self.read_json("lidarseg") 70 | if with_panoptic: 71 | self.panoptic = self.read_json("panoptic") 72 | self.sample_data_dict: dict[SENSOR_CHANNELS, pl.DataFrame] 73 | self.sample_annotation_dict: dict[SENSOR_CHANNELS, pl.DataFrame] = {} 74 | 75 | self.lidar_pad_length = 35_000 76 | self.verbose = verbose 77 | self.with_sample_annotation = with_sample_annotation 78 | 79 | self._prepare(key_frames_only) 80 | 81 | def read_json(self, name: ANNOTATION_FILES): 82 | """Reads NuScenes json files.""" 83 | return pl.read_json(self.data_root / self.version / f"{name}.json").select(SELECTION[name]) 84 | 85 | def _assign_sample_keyframe_index(self, sample: pl.DataFrame): 86 | sample = sample.sort("scene.token", "sample.timestamp").with_columns( 87 | pl.arange(0, pl.count()).over("scene.token").alias("sample.key_frame_index") 88 | ) 89 | return sample 90 | 91 | @staticmethod 92 | def _box_to_dict(sample_data_token_name: str, sample_data_token_name_value: str, box): 93 | return { 94 | sample_data_token_name: sample_data_token_name_value, 95 | "sample_annotation.token": box.token, 96 | "sample_annotation.translation": list(box.center), 97 | "sample_annotation.size": list(box.wlh), 98 | "sample_annotation.rotation": list(box.orientation.q), 99 | } 100 | 101 | def load_sensor_sample_annotation(self, sensor_channel: SENSOR_CHANNELS): 102 | """Loads sensor synced sample annotation from cache or builds it from scratch.""" 103 | if self.extra_data_root is not None: 104 | filename = self.extra_data_root / "sample_annotation_cache" / f"{sensor_channel}.arrow" 105 | if filename.exists(): 106 | self.sample_annotation_dict[sensor_channel] = pl.read_ipc(filename, memory_map=False) 107 | return 108 | self.build_sensor_sample_annotation(sensor_channel) 109 | if self.extra_data_root is not None: 110 | filename.parent.mkdir(exist_ok=True, parents=True) 111 | self.sample_annotation_dict[sensor_channel].write_ipc(filename, compression="zstd") 112 | 113 | def build_sensor_sample_annotation(self, sensor_channel: SENSOR_CHANNELS): 114 | """Interpolates sample annotation to given sample data token.""" 115 | import nuscenes 116 | 117 | nuscenes = nuscenes.NuScenes(verbose=False, version=self.version, dataroot=self.data_root) 118 | 119 | sample_data_token_name = f"{sensor_channel}.sample_data.token" 120 | sample_data_tokens = self.sample_data_dict[sensor_channel][sample_data_token_name] 121 | 122 | boxes = [] 123 | for sample_data_token in tqdm.tqdm( 124 | sample_data_tokens, 125 | desc=f"Build sample annotation for sensor {sensor_channel}", 126 | ): 127 | boxes.extend( 128 | [ 129 | self._box_to_dict(sample_data_token_name, sample_data_token, box) 130 | for box in nuscenes.get_boxes(sample_data_token) 131 | ] 132 | ) 133 | interpolated_annotations = pl.DataFrame(boxes).with_columns( 134 | pl.col(sample_data_token_name), 135 | pl.col("sample_annotation.token"), 136 | pl.col("sample_annotation.translation").cast(pl.Array(width=3, inner=pl.Float32)), 137 | pl.col("sample_annotation.size").cast(pl.Array(width=3, inner=pl.Float32)), 138 | pl.col("sample_annotation.rotation").cast(pl.Array(width=4, inner=pl.Float32)), 139 | ) 140 | # join additional information 141 | interpolated_annotations = self.join( 142 | interpolated_annotations, 143 | self.sample_annotation.select( 144 | pl.exclude( 145 | "sample_annotation.translation", 146 | "sample_annotation.size", 147 | "sample_annotation.rotation", 148 | "transform.global_from_obj", 149 | "transform.obj_from_global", 150 | "sample_annotation.size_aligned", 151 | ) 152 | ), 153 | ) 154 | interpolated_annotations = self.prepare_sample_annotation(interpolated_annotations) 155 | 156 | self.sample_annotation_dict[sensor_channel] = interpolated_annotations 157 | 158 | def load_cvpr2023_occupancy(self, scene: pl.DataFrame, root_path: Optional[Union[Path, str]] = None): 159 | """Loads occupancy ground truth from disk.""" 160 | 161 | root_path = Path(root_path) if root_path is not None else self.extra_data_root 162 | 163 | def make_filename(x): 164 | scene_name, sample_token = x 165 | return root_path / "nuscenes_occ3d" / "gts" / scene_name / sample_token / "labels.npz" 166 | 167 | filenames = list(map(make_filename, scene.select("scene.name", "sample.token").iter_rows())) 168 | data_dict: dict[str, list[np.ndarray]] = {} 169 | for filename in filenames: 170 | for k, v in np.load(filename).items(): 171 | if k not in data_dict: 172 | data_dict[k] = [] 173 | data_dict[k].append(v) 174 | list_of_series = [numpy_to_series(f"sample.occ_gt.{k}", np.stack(v)) for k, v in data_dict.items()] 175 | scene = scene.with_columns(*list_of_series) 176 | scene = scene.sort("sample.timestamp") 177 | return scene 178 | 179 | def load_open_occpancy(self, scene: pl.DataFrame, root_path: Optional[Union[Path, str]] = None): 180 | """Loads occupancy ground truth from disk.""" 181 | 182 | root_path = Path(root_path) if root_path is not None else self.extra_data_root 183 | 184 | def make_filename(x): 185 | scene_token, lidar_token = x 186 | return root_path / "nuScenes-Occupancy-v0.1" / f"scene_{scene_token}" / "occupancy" / f"{lidar_token}.npy" 187 | 188 | filenames = list( 189 | map( 190 | make_filename, 191 | scene.select("scene.token", "LIDAR_TOP.sample_data.token").iter_rows(), 192 | ) 193 | ) 194 | volume_shape = [512, 512, 40] 195 | num_samples = len(filenames) 196 | lower = np.broadcast_to(np.array([-51.2, -51.2, -5.0], dtype=np.float32), [len(filenames), 3]) 197 | upper = np.broadcast_to(np.array([51.2, 51.2, 3.0], dtype=np.float32), [len(filenames), 3]) 198 | dense_occupancy = np.full([num_samples] + volume_shape, fill_value=17, dtype=np.int32) 199 | for i, filename in enumerate(filenames): 200 | # https://github.com/JeffWang987/OpenOccupancy/blob/main/docs/prepare_data.md 201 | # https://github.com/JeffWang987/OpenOccupancy/blob/main/projects/occ_plugin/datasets/pipelines/loading.py 202 | # [z y x cls] 203 | pcd = np.load(filename) 204 | pcd_label = pcd[..., -1] 205 | # map free and noise to 17 206 | pcd_label[pcd_label == 0] = 17 207 | pcd_label[pcd_label == 255] = 17 208 | z, y, x = pcd[:, 0], pcd[:, 1], pcd[:, 2] 209 | batch_idx = np.broadcast_to(np.array(i), x.shape) 210 | dense_occupancy[batch_idx, x, y, z] = pcd_label 211 | scene = scene.with_columns( 212 | numpy_to_series("sample.open_occupancy.semantics", dense_occupancy), 213 | numpy_to_series("sample.open_occupancy.volume.lower", lower), 214 | numpy_to_series("sample.open_occupancy.volume.upper", upper), 215 | ) 216 | scene = scene.sort("sample.timestamp") 217 | return scene 218 | 219 | def load_surround_occ_occupancy( 220 | self, 221 | scene: pl.DataFrame, 222 | root_path: Optional[Union[Path, str]] = None, 223 | subdir: str = "surround_occ_occupancy", 224 | ): 225 | """Loads occupancy ground truth from disk.""" 226 | 227 | root_path = Path(root_path) if root_path is not None else self.extra_data_root 228 | 229 | def make_filename(x): 230 | token, pcd_filename = x 231 | filename = Path(pcd_filename).name 232 | return root_path / subdir / "samples" / f"{filename}.npy" 233 | 234 | filenames = list( 235 | map( 236 | make_filename, 237 | scene.select("scene.token", "LIDAR_TOP.sample_data.filename").iter_rows(), 238 | ) 239 | ) 240 | volume_shape = [200, 200, 16] 241 | num_samples = len(filenames) 242 | lower = np.broadcast_to(np.array([-50.0, -50.0, -5.0], dtype=np.float32), [len(filenames), 3]) 243 | upper = np.broadcast_to(np.array([50.0, 50.0, 3.0], dtype=np.float32), [len(filenames), 3]) 244 | dense_occupancy = np.full([num_samples] + volume_shape, fill_value=17, dtype=np.int32) 245 | for i, filename in enumerate(filenames): 246 | # https://github.com/weiyithu/SurroundOcc/blob/main/docs/data.md 247 | # https://github.com/weiyithu/SurroundOcc/blob/main/projects/mmdet3d_plugin/datasets/pipelines/loading.py 248 | # The ground truth is a (N, 4) tensor, N is the occupied voxel number, 249 | # The first three channels represent xyz voxel coordinate and last channel is semantic class. 250 | # [z y x cls] 251 | pcd = np.load(filename) 252 | pcd_label = pcd[..., -1] 253 | # map free and noise to 17 254 | pcd_label[pcd_label == 0] = 17 255 | pcd_label[pcd_label == 255] = 17 256 | x, y, z = pcd[:, 0], pcd[:, 1], pcd[:, 2] 257 | batch_idx = np.broadcast_to(np.array(i), x.shape) 258 | dense_occupancy[batch_idx, x, y, z] = pcd_label 259 | scene = scene.with_columns( 260 | numpy_to_series("sample.surround_occ.semantics", dense_occupancy), 261 | numpy_to_series("sample.surround_occ.volume.lower", lower), 262 | numpy_to_series("sample.surround_occ.volume.upper", upper), 263 | ) 264 | scene = scene.sort("sample.timestamp") 265 | return scene 266 | 267 | def load_scene_as_occupancy_gt( 268 | self, 269 | scene: pl.DataFrame, 270 | root_path: Optional[Union[Path, str]] = None, 271 | subdir: str = "occ_gt_release_v1_0", 272 | ): 273 | """Loads occupancy ground truth from disk.""" 274 | 275 | root_path = Path(root_path) if root_path is not None else self.extra_data_root 276 | 277 | def make_filename(x): 278 | scene_token, scene_name, key_frame_index = x 279 | return root_path / subdir / "trainval" / scene_name / f"{key_frame_index:03d}_occ.npy" 280 | 281 | filenames = list( 282 | map( 283 | make_filename, 284 | scene.select("scene.token", "scene.name", "sample.key_frame_index").iter_rows(), 285 | ) 286 | ) 287 | volume_shape = [200, 200, 16] 288 | num_voxels = math.prod(volume_shape) 289 | num_samples = len(filenames) 290 | lower = np.broadcast_to(np.array([-50.0, -50.0, -5.0], dtype=np.float32), [len(filenames), 3]) 291 | upper = np.broadcast_to(np.array([50.0, 50.0, 3.0], dtype=np.float32), [len(filenames), 3]) 292 | dense_occupancy = np.full([num_samples] + volume_shape, fill_value=17, dtype=np.int32) 293 | for i, filename in enumerate(filenames): 294 | # https://github.com/OpenDriveLab/OccNet/blob/main/projects/mmdet3d_plugin/datasets/nuscenes_dataset.py 295 | # https://github.com/OpenDriveLab/OccNet/blob/main/docs/prepare_dataset.md 296 | # The ground truth is a (N, 4) tensor, N is the occupied voxel number, 297 | # The first three channels represent xyz voxel coordinate and last channel is semantic class. 298 | # [z y x cls]num_voxels = math.prod(volume_shape) 299 | data = np.load(filename) 300 | flat_zyx_index = data[:, 0] 301 | label = data[:, 1] 302 | dense_voxel_grid = np.full(num_voxels, fill_value=17, dtype=np.int32) # [Z, Y, X] 303 | dense_voxel_grid[flat_zyx_index] = label 304 | # transpose to convert ZYX to XYZ 305 | dense_voxel_grid = dense_voxel_grid.reshape(volume_shape[::-1]).transpose(2, 1, 0) 306 | dense_occupancy[i] = dense_voxel_grid 307 | scene = scene.with_columns( 308 | numpy_to_series("sample.scene_as_occupancy.semantics", dense_occupancy), 309 | numpy_to_series("sample.scene_as_occupancy.volume.lower", lower), 310 | numpy_to_series("sample.scene_as_occupancy.volume.upper", upper), 311 | ) 312 | scene = scene.sort("sample.timestamp") 313 | return scene 314 | 315 | def load_reflection_and_transmission_spherical( 316 | self, scene: pl.DataFrame, root_path: Optional[Union[Path, str]] = None 317 | ): 318 | """Loads number of transmissions and reflections in sperical coords.""" 319 | root_path = Path(root_path) if root_path is not None else self.extra_data_root 320 | 321 | def make_filename(x): 322 | scene_name, lidar_token = x 323 | return ( 324 | root_path / "reflection_and_transmission_spherical" / scene_name / "LIDAR_TOP" / f"{lidar_token}.arrow" 325 | ) 326 | 327 | filenames = list( 328 | map( 329 | make_filename, 330 | scene.select("scene.name", "LIDAR_TOP.sample_data.token").iter_rows(), 331 | ) 332 | ) 333 | df = pl.scan_ipc(filenames, memory_map=False) 334 | scene = scene.lazy().join(df, on="LIDAR_TOP.sample_data.token", validate="1:1", how="left").collect() 335 | return scene 336 | 337 | def load_reflection_and_transmission_multi_frame( 338 | self, scene: pl.DataFrame, root_path: Optional[Union[Path, str]] = None 339 | ): 340 | """Loads number of transmissions and reflections in ego frame accumulated over multiple frames coords.""" 341 | root_path = Path(root_path) if root_path is not None else self.extra_data_root 342 | 343 | def make_filename(x): 344 | scene_name, lidar_token = x 345 | return ( 346 | root_path 347 | / "reflection_and_transmission_multi_frame" 348 | / scene_name 349 | / "LIDAR_TOP" 350 | / f"{lidar_token}.arrow" 351 | ) 352 | 353 | filenames = list( 354 | map( 355 | make_filename, 356 | scene.select("scene.name", "LIDAR_TOP.sample_data.token").iter_rows(), 357 | ) 358 | ) 359 | df = pl.scan_ipc(filenames, memory_map=False) 360 | scene = scene.lazy().join(df, on="LIDAR_TOP.sample_data.token", validate="1:1", how="left").collect() 361 | return scene 362 | 363 | def load_scene_flow_polars(self, scene: pl.DataFrame, root_path: Optional[Union[Path, str]] = None): 364 | """Loads scene flow information.""" 365 | root_path = Path(root_path) if root_path is not None else self.extra_data_root 366 | 367 | def make_filename(x): 368 | scene_name, lidar_token = x 369 | return root_path / "scene_flow" / scene_name / "LIDAR_TOP" / f"{lidar_token}.arrow" 370 | 371 | filenames = list( 372 | map( 373 | make_filename, 374 | scene.select("scene.name", "LIDAR_TOP.sample_data.token").iter_rows(), 375 | ) 376 | ) 377 | df = pl.scan_ipc(filenames, memory_map=False) 378 | scene = scene.lazy().join(df, on="LIDAR_TOP.sample_data.token", validate="1:1", how="left").collect() 379 | return scene 380 | 381 | def load_occupancy_preds( 382 | self, 383 | scene: pl.DataFrame, 384 | name: str, 385 | root_path: Optional[Union[Path, str]] = None, 386 | ): 387 | """Loads occupancy preds truth from disk.""" 388 | 389 | root_path = Path(root_path) if root_path is not None else self.extra_data_root 390 | 391 | def make_filename(x): 392 | scene_name, sample_token = x 393 | return root_path / name / f"{sample_token}.npz" 394 | 395 | filenames = list(map(make_filename, scene.select("scene.name", "sample.token").iter_rows())) 396 | data_dict: dict[str, list[np.ndarray]] = {} 397 | for filename in filenames: 398 | for k, v in np.load(filename).items(): 399 | if k not in data_dict: 400 | data_dict[k] = [] 401 | data_dict[k].append(v) 402 | list_of_series = [numpy_to_series(f"sample.{name}.{k}", np.stack(v)) for k, v in data_dict.items()] 403 | scene = scene.with_columns(*list_of_series) 404 | scene = scene.sort("sample.timestamp") 405 | return scene 406 | 407 | def sort_by_time(self, scene: pl.DataFrame): 408 | """Sort scene by time.""" 409 | return scene.sort("LIDAR_TOP.sample_data.timestamp") 410 | 411 | def save_scenes(self): 412 | """Save scenes to disk.""" 413 | for scene_name in tqdm.tqdm(self.scene["scene.name"]): 414 | scene = self.load_scene(scene_name) 415 | scene_dir = self.data_root / "scenes" 416 | scene_dir.mkdir(exist_ok=True) 417 | scene_name = scene.item(0, "scene.name") 418 | scene.write_ipc( 419 | scene_dir / f"{scene_name}.arrow", 420 | compression="zstd", 421 | ) 422 | 423 | def load_scene_from_file(self, scene_name: str): 424 | """Load scene from file.""" 425 | return pl.read_ipc( 426 | self.data_root / "scenes" / f"{scene_name}.arrow", 427 | memory_map=False, 428 | ) 429 | 430 | def build_transformations(self): 431 | """Build homogenous transformations.""" 432 | # sensor / ego 433 | self.calibrated_sensor = transform_from_columns( 434 | self.calibrated_sensor, 435 | "calibrated_sensor.rotation", 436 | "calibrated_sensor.translation", 437 | transform="transform.ego_from_sensor", 438 | transform_inv="transform.sensor_from_ego", 439 | ) 440 | # sensor intrinsics 441 | self.calibrated_sensor = pad_intrinsics_from_colums( 442 | self.calibrated_sensor, 443 | "calibrated_sensor.camera_intrinsic", 444 | "transform.image_from_sensor", 445 | ) 446 | # ego / global 447 | self.ego_pose = transform_from_columns( 448 | self.ego_pose, 449 | "ego_pose.rotation", 450 | "ego_pose.translation", 451 | transform="transform.global_from_ego", 452 | transform_inv="transform.ego_from_global", 453 | ) 454 | 455 | def build_sample_annotation_transforms(self, sample_annotation: pl.DataFrame): 456 | """Build sample annotation transforms.""" 457 | sample_annotation = transform_from_columns( 458 | sample_annotation, 459 | "sample_annotation.rotation", 460 | "sample_annotation.translation", 461 | transform="transform.global_from_obj", 462 | transform_inv="transform.obj_from_global", 463 | ) 464 | return sample_annotation 465 | 466 | def filter_key_frames_only(self): 467 | """Filter by key frame.""" 468 | self.sample_data = self.sample_data.filter(pl.col("sample_data.is_key_frame")) 469 | 470 | def build_sample_data_dict(self): 471 | """Build sample data dict per sensor.""" 472 | self.sample_data_dict: dict[SENSOR_CHANNELS, pl.DataFrame] = { 473 | channel: df.select(pl.all().name.prefix(f"{channel}.")) 474 | for (channel,), df in self.sample_data.partition_by("sensor.channel", as_dict=True).items() 475 | } 476 | 477 | def load_sample_data( 478 | self, 479 | sample: pl.DataFrame, 480 | sensor_channel: SENSOR_CHANNELS, 481 | with_data: bool = True, 482 | ): 483 | """Load sample data.""" 484 | sample = self.join( 485 | sample, 486 | self.sample_data_dict[sensor_channel].rename({f"{sensor_channel}.sample.token": "sample.token"}), 487 | # validate="1:1" if self.key_frames_only else "1:m", 488 | ) 489 | 490 | if with_data: 491 | sample = self._load_sensor_data(sample, sensor_channel) 492 | sample = self.sort_by_time(sample) 493 | return sample 494 | 495 | def load_panoptic( 496 | self, 497 | df: pl.DataFrame, 498 | with_data: bool = False, 499 | ): 500 | """Load panoptic info.""" 501 | df = self.join(df, self.lidarseg, validate="1:1") 502 | if with_data: 503 | df = self._load_panoptic_data(df) 504 | return df 505 | 506 | def _load_panoptic_data( 507 | self, 508 | df: pl.DataFrame, 509 | ): 510 | """Load panoptic data.""" 511 | data = np.stack(list(map(self._load_lidarseg_from_file, df["LIDAR_TOP.panoptic.filename"]))) 512 | df = df.with_columns(numpy_to_series("LIDAR_TOP.panoptic.data", data)) 513 | return df 514 | 515 | def load_lidarseg( 516 | self, 517 | df: pl.DataFrame, 518 | with_data: bool = False, 519 | ): 520 | """Load lidar segmentation data.""" 521 | df = self.join(df, self.lidarseg, validate="1:1") 522 | if with_data: 523 | df = self._load_lidarseg_data(df) 524 | return df 525 | 526 | def _load_lidarseg_data( 527 | self, 528 | df: pl.DataFrame, 529 | ): 530 | data = np.stack( 531 | df["LIDAR_TOP.lidarseg.filename"] 532 | .map_elements(self._load_lidarseg_from_file, return_dtype=pl.Object()) 533 | .to_list() 534 | ) 535 | df = df.with_columns(numpy_to_series("LIDAR_TOP.points_category_idx", data)) 536 | return df 537 | 538 | def _load_sensor_data( 539 | self, 540 | df: pl.DataFrame, 541 | sensor_channel: SENSOR_CHANNELS, 542 | ): 543 | """Loads sensor data.""" 544 | if sensor_channel in get_args(CAMERA_CHANNELS): 545 | load_fn = self._load_image_from_file 546 | prefix = f"{sensor_channel}.sample_data" 547 | filename = f"{sensor_channel}.sample_data.filename" 548 | elif sensor_channel in get_args(LIDAR_CHANNELS): 549 | load_fn = self._load_lidar_from_file 550 | prefix = f"{sensor_channel}.sample_data" 551 | filename = f"{sensor_channel}.sample_data.filename" 552 | elif sensor_channel in get_args(RADAR_CHANNELS): 553 | raise NotImplementedError("Loading radar data is not yet supported") 554 | else: 555 | raise ValueError(f"Unknown channel: {sensor_channel}") 556 | list_of_dicts = [load_fn(f) for f in df[filename]] 557 | dict_of_data = {f"{prefix}.{k}": np.stack([d[k] for d in list_of_dicts]) for k in list_of_dicts[0].keys()} 558 | return df.with_columns(*[numpy_to_series(k, v) for k, v in dict_of_data.items()]) 559 | 560 | def join( 561 | self, 562 | df: pl.DataFrame, 563 | other: pl.DataFrame, 564 | how: JoinStrategy = "inner", 565 | *, 566 | allow_duplicates: bool = False, 567 | suffix: str = "_right", 568 | validate: JoinValidation = "m:m", 569 | ): 570 | """Join frames on common tokens.""" 571 | return join_on_token( 572 | df, 573 | other, 574 | how=how, 575 | suffix=suffix, 576 | validate=validate, 577 | verbose=self.verbose, 578 | allow_duplicates=allow_duplicates, 579 | ) 580 | 581 | def _load_image_from_file(self, filename: str): 582 | """Load single image from file.""" 583 | return {"image": np.asarray(Image.open(self.data_root / filename))} 584 | 585 | def _load_lidar_from_file(self, filename: str): 586 | """Loads LIDAR data from binary numpy format. Data is stored as (x, y, z, intensity, ring index).""" 587 | dims_to_load = [0, 1, 2] # [x, y, z] 588 | scan = np.fromfile(self.data_root / filename, dtype=np.float32) 589 | points = scan.reshape((-1, 5))[:, dims_to_load] 590 | points_padded = np.empty((self.lidar_pad_length, 3), dtype=points.dtype) 591 | num_points = len(points) 592 | points_padded[:num_points] = points 593 | points_padded[num_points:] = float("nan") 594 | mask = np.empty((self.lidar_pad_length,), dtype=bool) 595 | mask[:num_points] = 1 596 | mask[num_points:] = 0 597 | return {"points_lidar": points_padded, "points_mask": mask} 598 | 599 | def _load_lidarseg_from_file(self, filename: str): 600 | """Loads lidar seg data from file.""" 601 | data = np.fromfile(self.data_root / filename, dtype=np.uint8) 602 | data_padded = np.empty(self.lidar_pad_length, dtype=data.dtype) 603 | num_points = len(data) 604 | data_padded[:num_points] = data 605 | data_padded[num_points:] = 255 606 | return {"points_semantics": data_padded} 607 | 608 | def prepare_sample_data(self): 609 | """Prepares sample data.""" 610 | self.sample_data = self.join(self.sample_data, self.ego_pose, validate="1:1") 611 | self.sample_data = self.join(self.sample_data, self.calibrated_sensor, validate="m:1") 612 | self.sample_data = self.join(self.sample_data, self.sensor, validate="m:1") 613 | 614 | def prepare_sample_annotation(self, sample_annotation: pl.DataFrame): 615 | """Perpares sample annotation.""" 616 | size = series_to_numpy(sample_annotation["sample_annotation.size"]) 617 | size_aligned = size[..., [1, 0, 2]] # front, left, up 618 | 619 | sample_annotation = sample_annotation.with_columns( 620 | numpy_to_series("sample_annotation.size_aligned", size_aligned) 621 | ) 622 | 623 | # obj / global 624 | sample_annotation = self.build_sample_annotation_transforms(sample_annotation) 625 | return sample_annotation 626 | 627 | def assign_instance_index(self): 628 | """Assign unique index to each instance.""" 629 | self.instance = ( 630 | self.instance.sort("instance.token") 631 | .with_row_count("instance.index", offset=1) # 0 will be used as "no instance" 632 | .with_columns(pl.col("instance.index").cast(pl.Int64)) 633 | ) 634 | 635 | def _prepare(self, key_frames_only: bool = True): 636 | """Prepare dataset for usage.""" 637 | self.sample = self._assign_sample_keyframe_index(sample=self.sample) 638 | self.assign_instance_index() 639 | if key_frames_only: 640 | self.filter_key_frames_only() 641 | self.build_transformations() 642 | self.prepare_sample_data() 643 | if self.with_sample_annotation: 644 | self.sample_annotation = self.prepare_sample_annotation(self.sample_annotation) 645 | self.build_sample_data_dict() 646 | 647 | def sample_scene(self): 648 | """Samples and returns a random scene.""" 649 | return self.get_scene(self.scene.sample(1)) 650 | 651 | def get_scene(self, scene: pl.DataFrame): 652 | """Get samples for a given scene.""" 653 | return self.join(scene, self.sample) 654 | 655 | def load_cameras_and_lidar( 656 | self, 657 | scene: pl.DataFrame, 658 | cameras: Sequence[CAMERA_CHANNELS] = CAMERAS, 659 | lidar: LIDAR_CHANNELS = "LIDAR_TOP", 660 | ): 661 | """Loads camera and lidar data for a given scene.""" 662 | scene = self.load_sample_data(scene, lidar, with_data=True) 663 | for cam in cameras: 664 | scene = self.load_sample_data(scene, cam, with_data=True) 665 | return scene 666 | 667 | def scene_by_name(self, scene_name: str): 668 | """Select scene by name.""" 669 | scene = self.get_scene(self.scene.filter(pl.col("scene.name") == scene_name)) 670 | return scene 671 | 672 | def load_scene(self, scene_name: str): 673 | """Loads a scene by name.""" 674 | scene = self.scene_by_name(scene_name) 675 | scene = self.load_cameras_and_lidar(scene) 676 | scene = self.load_lidarseg(scene, with_data=True) 677 | scene = scene.sort("sample.timestamp") 678 | return scene 679 | 680 | def train_scenes(self): 681 | """Training scenes.""" 682 | import nuscenes.utils.splits as splits 683 | 684 | if self.version == "v1.0-mini": 685 | scene_names = pl.Series("scene.name", splits.mini_train) 686 | elif self.version == "v1.0-trainval": 687 | scene_names = pl.Series("scene.name", splits.train) 688 | else: 689 | raise ValueError("Invalid version") 690 | return self.scene.filter(pl.col("scene.name").is_in(scene_names)) 691 | 692 | def val_scenes(self): 693 | """Validation scenes.""" 694 | import nuscenes.utils.splits as splits 695 | 696 | if self.version == "v1.0-mini": 697 | scene_names = pl.Series("scene.name", splits.mini_val) 698 | elif self.version == "v1.0-trainval": 699 | scene_names = pl.Series("scene.name", splits.val) 700 | else: 701 | raise ValueError("Invalid version") 702 | return self.scene.filter(pl.col("scene.name").is_in(scene_names)) 703 | 704 | def scenes_by_split(self, split: Literal["train", "val", "trainval"]): 705 | """Scenes in selected splits.""" 706 | if split == "train": 707 | return self.train_scenes() 708 | elif split == "val": 709 | return self.val_scenes() 710 | elif split == "trainval": 711 | return pl.concat([self.train_scenes(), self.val_scenes()], how="vertical") 712 | else: 713 | raise ValueError("Invalid split") 714 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | GNU AFFERO GENERAL PUBLIC LICENSE 2 | Version 3, 19 November 2007 3 | 4 | Copyright (C) 2007 Free Software Foundation, Inc. 5 | Everyone is permitted to copy and distribute verbatim copies 6 | of this license document, but changing it is not allowed. 7 | 8 | Preamble 9 | 10 | The GNU Affero General Public License is a free, copyleft license for 11 | software and other kinds of works, specifically designed to ensure 12 | cooperation with the community in the case of network server software. 13 | 14 | The licenses for most software and other practical works are designed 15 | to take away your freedom to share and change the works. By contrast, 16 | our General Public Licenses are intended to guarantee your freedom to 17 | share and change all versions of a program--to make sure it remains free 18 | software for all its users. 19 | 20 | When we speak of free software, we are referring to freedom, not 21 | price. Our General Public Licenses are designed to make sure that you 22 | have the freedom to distribute copies of free software (and charge for 23 | them if you wish), that you receive source code or can get it if you 24 | want it, that you can change the software or use pieces of it in new 25 | free programs, and that you know you can do these things. 26 | 27 | Developers that use our General Public Licenses protect your rights 28 | with two steps: (1) assert copyright on the software, and (2) offer 29 | you this License which gives you legal permission to copy, distribute 30 | and/or modify the software. 31 | 32 | A secondary benefit of defending all users' freedom is that 33 | improvements made in alternate versions of the program, if they 34 | receive widespread use, become available for other developers to 35 | incorporate. Many developers of free software are heartened and 36 | encouraged by the resulting cooperation. However, in the case of 37 | software used on network servers, this result may fail to come about. 38 | The GNU General Public License permits making a modified version and 39 | letting the public access it on a server without ever releasing its 40 | source code to the public. 41 | 42 | The GNU Affero General Public License is designed specifically to 43 | ensure that, in such cases, the modified source code becomes available 44 | to the community. It requires the operator of a network server to 45 | provide the source code of the modified version running there to the 46 | users of that server. Therefore, public use of a modified version, on 47 | a publicly accessible server, gives the public access to the source 48 | code of the modified version. 49 | 50 | An older license, called the Affero General Public License and 51 | published by Affero, was designed to accomplish similar goals. This is 52 | a different license, not a version of the Affero GPL, but Affero has 53 | released a new version of the Affero GPL which permits relicensing under 54 | this license. 55 | 56 | The precise terms and conditions for copying, distribution and 57 | modification follow. 58 | 59 | TERMS AND CONDITIONS 60 | 61 | 0. Definitions. 62 | 63 | "This License" refers to version 3 of the GNU Affero General Public License. 64 | 65 | "Copyright" also means copyright-like laws that apply to other kinds of 66 | works, such as semiconductor masks. 67 | 68 | "The Program" refers to any copyrightable work licensed under this 69 | License. Each licensee is addressed as "you". "Licensees" and 70 | "recipients" may be individuals or organizations. 71 | 72 | To "modify" a work means to copy from or adapt all or part of the work 73 | in a fashion requiring copyright permission, other than the making of an 74 | exact copy. The resulting work is called a "modified version" of the 75 | earlier work or a work "based on" the earlier work. 76 | 77 | A "covered work" means either the unmodified Program or a work based 78 | on the Program. 79 | 80 | To "propagate" a work means to do anything with it that, without 81 | permission, would make you directly or secondarily liable for 82 | infringement under applicable copyright law, except executing it on a 83 | computer or modifying a private copy. Propagation includes copying, 84 | distribution (with or without modification), making available to the 85 | public, and in some countries other activities as well. 86 | 87 | To "convey" a work means any kind of propagation that enables other 88 | parties to make or receive copies. Mere interaction with a user through 89 | a computer network, with no transfer of a copy, is not conveying. 90 | 91 | An interactive user interface displays "Appropriate Legal Notices" 92 | to the extent that it includes a convenient and prominently visible 93 | feature that (1) displays an appropriate copyright notice, and (2) 94 | tells the user that there is no warranty for the work (except to the 95 | extent that warranties are provided), that licensees may convey the 96 | work under this License, and how to view a copy of this License. If 97 | the interface presents a list of user commands or options, such as a 98 | menu, a prominent item in the list meets this criterion. 99 | 100 | 1. Source Code. 101 | 102 | The "source code" for a work means the preferred form of the work 103 | for making modifications to it. "Object code" means any non-source 104 | form of a work. 105 | 106 | A "Standard Interface" means an interface that either is an official 107 | standard defined by a recognized standards body, or, in the case of 108 | interfaces specified for a particular programming language, one that 109 | is widely used among developers working in that language. 110 | 111 | The "System Libraries" of an executable work include anything, other 112 | than the work as a whole, that (a) is included in the normal form of 113 | packaging a Major Component, but which is not part of that Major 114 | Component, and (b) serves only to enable use of the work with that 115 | Major Component, or to implement a Standard Interface for which an 116 | implementation is available to the public in source code form. A 117 | "Major Component", in this context, means a major essential component 118 | (kernel, window system, and so on) of the specific operating system 119 | (if any) on which the executable work runs, or a compiler used to 120 | produce the work, or an object code interpreter used to run it. 121 | 122 | The "Corresponding Source" for a work in object code form means all 123 | the source code needed to generate, install, and (for an executable 124 | work) run the object code and to modify the work, including scripts to 125 | control those activities. However, it does not include the work's 126 | System Libraries, or general-purpose tools or generally available free 127 | programs which are used unmodified in performing those activities but 128 | which are not part of the work. For example, Corresponding Source 129 | includes interface definition files associated with source files for 130 | the work, and the source code for shared libraries and dynamically 131 | linked subprograms that the work is specifically designed to require, 132 | such as by intimate data communication or control flow between those 133 | subprograms and other parts of the work. 134 | 135 | The Corresponding Source need not include anything that users 136 | can regenerate automatically from other parts of the Corresponding 137 | Source. 138 | 139 | The Corresponding Source for a work in source code form is that 140 | same work. 141 | 142 | 2. Basic Permissions. 143 | 144 | All rights granted under this License are granted for the term of 145 | copyright on the Program, and are irrevocable provided the stated 146 | conditions are met. This License explicitly affirms your unlimited 147 | permission to run the unmodified Program. The output from running a 148 | covered work is covered by this License only if the output, given its 149 | content, constitutes a covered work. This License acknowledges your 150 | rights of fair use or other equivalent, as provided by copyright law. 151 | 152 | You may make, run and propagate covered works that you do not 153 | convey, without conditions so long as your license otherwise remains 154 | in force. You may convey covered works to others for the sole purpose 155 | of having them make modifications exclusively for you, or provide you 156 | with facilities for running those works, provided that you comply with 157 | the terms of this License in conveying all material for which you do 158 | not control copyright. Those thus making or running the covered works 159 | for you must do so exclusively on your behalf, under your direction 160 | and control, on terms that prohibit them from making any copies of 161 | your copyrighted material outside their relationship with you. 162 | 163 | Conveying under any other circumstances is permitted solely under 164 | the conditions stated below. Sublicensing is not allowed; section 10 165 | makes it unnecessary. 166 | 167 | 3. Protecting Users' Legal Rights From Anti-Circumvention Law. 168 | 169 | No covered work shall be deemed part of an effective technological 170 | measure under any applicable law fulfilling obligations under article 171 | 11 of the WIPO copyright treaty adopted on 20 December 1996, or 172 | similar laws prohibiting or restricting circumvention of such 173 | measures. 174 | 175 | When you convey a covered work, you waive any legal power to forbid 176 | circumvention of technological measures to the extent such circumvention 177 | is effected by exercising rights under this License with respect to 178 | the covered work, and you disclaim any intention to limit operation or 179 | modification of the work as a means of enforcing, against the work's 180 | users, your or third parties' legal rights to forbid circumvention of 181 | technological measures. 182 | 183 | 4. Conveying Verbatim Copies. 184 | 185 | You may convey verbatim copies of the Program's source code as you 186 | receive it, in any medium, provided that you conspicuously and 187 | appropriately publish on each copy an appropriate copyright notice; 188 | keep intact all notices stating that this License and any 189 | non-permissive terms added in accord with section 7 apply to the code; 190 | keep intact all notices of the absence of any warranty; and give all 191 | recipients a copy of this License along with the Program. 192 | 193 | You may charge any price or no price for each copy that you convey, 194 | and you may offer support or warranty protection for a fee. 195 | 196 | 5. Conveying Modified Source Versions. 197 | 198 | You may convey a work based on the Program, or the modifications to 199 | produce it from the Program, in the form of source code under the 200 | terms of section 4, provided that you also meet all of these conditions: 201 | 202 | a) The work must carry prominent notices stating that you modified 203 | it, and giving a relevant date. 204 | 205 | b) The work must carry prominent notices stating that it is 206 | released under this License and any conditions added under section 207 | 7. This requirement modifies the requirement in section 4 to 208 | "keep intact all notices". 209 | 210 | c) You must license the entire work, as a whole, under this 211 | License to anyone who comes into possession of a copy. This 212 | License will therefore apply, along with any applicable section 7 213 | additional terms, to the whole of the work, and all its parts, 214 | regardless of how they are packaged. This License gives no 215 | permission to license the work in any other way, but it does not 216 | invalidate such permission if you have separately received it. 217 | 218 | d) If the work has interactive user interfaces, each must display 219 | Appropriate Legal Notices; however, if the Program has interactive 220 | interfaces that do not display Appropriate Legal Notices, your 221 | work need not make them do so. 222 | 223 | A compilation of a covered work with other separate and independent 224 | works, which are not by their nature extensions of the covered work, 225 | and which are not combined with it such as to form a larger program, 226 | in or on a volume of a storage or distribution medium, is called an 227 | "aggregate" if the compilation and its resulting copyright are not 228 | used to limit the access or legal rights of the compilation's users 229 | beyond what the individual works permit. Inclusion of a covered work 230 | in an aggregate does not cause this License to apply to the other 231 | parts of the aggregate. 232 | 233 | 6. Conveying Non-Source Forms. 234 | 235 | You may convey a covered work in object code form under the terms 236 | of sections 4 and 5, provided that you also convey the 237 | machine-readable Corresponding Source under the terms of this License, 238 | in one of these ways: 239 | 240 | a) Convey the object code in, or embodied in, a physical product 241 | (including a physical distribution medium), accompanied by the 242 | Corresponding Source fixed on a durable physical medium 243 | customarily used for software interchange. 244 | 245 | b) Convey the object code in, or embodied in, a physical product 246 | (including a physical distribution medium), accompanied by a 247 | written offer, valid for at least three years and valid for as 248 | long as you offer spare parts or customer support for that product 249 | model, to give anyone who possesses the object code either (1) a 250 | copy of the Corresponding Source for all the software in the 251 | product that is covered by this License, on a durable physical 252 | medium customarily used for software interchange, for a price no 253 | more than your reasonable cost of physically performing this 254 | conveying of source, or (2) access to copy the 255 | Corresponding Source from a network server at no charge. 256 | 257 | c) Convey individual copies of the object code with a copy of the 258 | written offer to provide the Corresponding Source. This 259 | alternative is allowed only occasionally and noncommercially, and 260 | only if you received the object code with such an offer, in accord 261 | with subsection 6b. 262 | 263 | d) Convey the object code by offering access from a designated 264 | place (gratis or for a charge), and offer equivalent access to the 265 | Corresponding Source in the same way through the same place at no 266 | further charge. You need not require recipients to copy the 267 | Corresponding Source along with the object code. If the place to 268 | copy the object code is a network server, the Corresponding Source 269 | may be on a different server (operated by you or a third party) 270 | that supports equivalent copying facilities, provided you maintain 271 | clear directions next to the object code saying where to find the 272 | Corresponding Source. Regardless of what server hosts the 273 | Corresponding Source, you remain obligated to ensure that it is 274 | available for as long as needed to satisfy these requirements. 275 | 276 | e) Convey the object code using peer-to-peer transmission, provided 277 | you inform other peers where the object code and Corresponding 278 | Source of the work are being offered to the general public at no 279 | charge under subsection 6d. 280 | 281 | A separable portion of the object code, whose source code is excluded 282 | from the Corresponding Source as a System Library, need not be 283 | included in conveying the object code work. 284 | 285 | A "User Product" is either (1) a "consumer product", which means any 286 | tangible personal property which is normally used for personal, family, 287 | or household purposes, or (2) anything designed or sold for incorporation 288 | into a dwelling. In determining whether a product is a consumer product, 289 | doubtful cases shall be resolved in favor of coverage. For a particular 290 | product received by a particular user, "normally used" refers to a 291 | typical or common use of that class of product, regardless of the status 292 | of the particular user or of the way in which the particular user 293 | actually uses, or expects or is expected to use, the product. A product 294 | is a consumer product regardless of whether the product has substantial 295 | commercial, industrial or non-consumer uses, unless such uses represent 296 | the only significant mode of use of the product. 297 | 298 | "Installation Information" for a User Product means any methods, 299 | procedures, authorization keys, or other information required to install 300 | and execute modified versions of a covered work in that User Product from 301 | a modified version of its Corresponding Source. The information must 302 | suffice to ensure that the continued functioning of the modified object 303 | code is in no case prevented or interfered with solely because 304 | modification has been made. 305 | 306 | If you convey an object code work under this section in, or with, or 307 | specifically for use in, a User Product, and the conveying occurs as 308 | part of a transaction in which the right of possession and use of the 309 | User Product is transferred to the recipient in perpetuity or for a 310 | fixed term (regardless of how the transaction is characterized), the 311 | Corresponding Source conveyed under this section must be accompanied 312 | by the Installation Information. But this requirement does not apply 313 | if neither you nor any third party retains the ability to install 314 | modified object code on the User Product (for example, the work has 315 | been installed in ROM). 316 | 317 | The requirement to provide Installation Information does not include a 318 | requirement to continue to provide support service, warranty, or updates 319 | for a work that has been modified or installed by the recipient, or for 320 | the User Product in which it has been modified or installed. Access to a 321 | network may be denied when the modification itself materially and 322 | adversely affects the operation of the network or violates the rules and 323 | protocols for communication across the network. 324 | 325 | Corresponding Source conveyed, and Installation Information provided, 326 | in accord with this section must be in a format that is publicly 327 | documented (and with an implementation available to the public in 328 | source code form), and must require no special password or key for 329 | unpacking, reading or copying. 330 | 331 | 7. Additional Terms. 332 | 333 | "Additional permissions" are terms that supplement the terms of this 334 | License by making exceptions from one or more of its conditions. 335 | Additional permissions that are applicable to the entire Program shall 336 | be treated as though they were included in this License, to the extent 337 | that they are valid under applicable law. If additional permissions 338 | apply only to part of the Program, that part may be used separately 339 | under those permissions, but the entire Program remains governed by 340 | this License without regard to the additional permissions. 341 | 342 | When you convey a copy of a covered work, you may at your option 343 | remove any additional permissions from that copy, or from any part of 344 | it. (Additional permissions may be written to require their own 345 | removal in certain cases when you modify the work.) You may place 346 | additional permissions on material, added by you to a covered work, 347 | for which you have or can give appropriate copyright permission. 348 | 349 | Notwithstanding any other provision of this License, for material you 350 | add to a covered work, you may (if authorized by the copyright holders of 351 | that material) supplement the terms of this License with terms: 352 | 353 | a) Disclaiming warranty or limiting liability differently from the 354 | terms of sections 15 and 16 of this License; or 355 | 356 | b) Requiring preservation of specified reasonable legal notices or 357 | author attributions in that material or in the Appropriate Legal 358 | Notices displayed by works containing it; or 359 | 360 | c) Prohibiting misrepresentation of the origin of that material, or 361 | requiring that modified versions of such material be marked in 362 | reasonable ways as different from the original version; or 363 | 364 | d) Limiting the use for publicity purposes of names of licensors or 365 | authors of the material; or 366 | 367 | e) Declining to grant rights under trademark law for use of some 368 | trade names, trademarks, or service marks; or 369 | 370 | f) Requiring indemnification of licensors and authors of that 371 | material by anyone who conveys the material (or modified versions of 372 | it) with contractual assumptions of liability to the recipient, for 373 | any liability that these contractual assumptions directly impose on 374 | those licensors and authors. 375 | 376 | All other non-permissive additional terms are considered "further 377 | restrictions" within the meaning of section 10. If the Program as you 378 | received it, or any part of it, contains a notice stating that it is 379 | governed by this License along with a term that is a further 380 | restriction, you may remove that term. If a license document contains 381 | a further restriction but permits relicensing or conveying under this 382 | License, you may add to a covered work material governed by the terms 383 | of that license document, provided that the further restriction does 384 | not survive such relicensing or conveying. 385 | 386 | If you add terms to a covered work in accord with this section, you 387 | must place, in the relevant source files, a statement of the 388 | additional terms that apply to those files, or a notice indicating 389 | where to find the applicable terms. 390 | 391 | Additional terms, permissive or non-permissive, may be stated in the 392 | form of a separately written license, or stated as exceptions; 393 | the above requirements apply either way. 394 | 395 | 8. Termination. 396 | 397 | You may not propagate or modify a covered work except as expressly 398 | provided under this License. Any attempt otherwise to propagate or 399 | modify it is void, and will automatically terminate your rights under 400 | this License (including any patent licenses granted under the third 401 | paragraph of section 11). 402 | 403 | However, if you cease all violation of this License, then your 404 | license from a particular copyright holder is reinstated (a) 405 | provisionally, unless and until the copyright holder explicitly and 406 | finally terminates your license, and (b) permanently, if the copyright 407 | holder fails to notify you of the violation by some reasonable means 408 | prior to 60 days after the cessation. 409 | 410 | Moreover, your license from a particular copyright holder is 411 | reinstated permanently if the copyright holder notifies you of the 412 | violation by some reasonable means, this is the first time you have 413 | received notice of violation of this License (for any work) from that 414 | copyright holder, and you cure the violation prior to 30 days after 415 | your receipt of the notice. 416 | 417 | Termination of your rights under this section does not terminate the 418 | licenses of parties who have received copies or rights from you under 419 | this License. If your rights have been terminated and not permanently 420 | reinstated, you do not qualify to receive new licenses for the same 421 | material under section 10. 422 | 423 | 9. Acceptance Not Required for Having Copies. 424 | 425 | You are not required to accept this License in order to receive or 426 | run a copy of the Program. Ancillary propagation of a covered work 427 | occurring solely as a consequence of using peer-to-peer transmission 428 | to receive a copy likewise does not require acceptance. However, 429 | nothing other than this License grants you permission to propagate or 430 | modify any covered work. These actions infringe copyright if you do 431 | not accept this License. Therefore, by modifying or propagating a 432 | covered work, you indicate your acceptance of this License to do so. 433 | 434 | 10. Automatic Licensing of Downstream Recipients. 435 | 436 | Each time you convey a covered work, the recipient automatically 437 | receives a license from the original licensors, to run, modify and 438 | propagate that work, subject to this License. You are not responsible 439 | for enforcing compliance by third parties with this License. 440 | 441 | An "entity transaction" is a transaction transferring control of an 442 | organization, or substantially all assets of one, or subdividing an 443 | organization, or merging organizations. If propagation of a covered 444 | work results from an entity transaction, each party to that 445 | transaction who receives a copy of the work also receives whatever 446 | licenses to the work the party's predecessor in interest had or could 447 | give under the previous paragraph, plus a right to possession of the 448 | Corresponding Source of the work from the predecessor in interest, if 449 | the predecessor has it or can get it with reasonable efforts. 450 | 451 | You may not impose any further restrictions on the exercise of the 452 | rights granted or affirmed under this License. For example, you may 453 | not impose a license fee, royalty, or other charge for exercise of 454 | rights granted under this License, and you may not initiate litigation 455 | (including a cross-claim or counterclaim in a lawsuit) alleging that 456 | any patent claim is infringed by making, using, selling, offering for 457 | sale, or importing the Program or any portion of it. 458 | 459 | 11. Patents. 460 | 461 | A "contributor" is a copyright holder who authorizes use under this 462 | License of the Program or a work on which the Program is based. The 463 | work thus licensed is called the contributor's "contributor version". 464 | 465 | A contributor's "essential patent claims" are all patent claims 466 | owned or controlled by the contributor, whether already acquired or 467 | hereafter acquired, that would be infringed by some manner, permitted 468 | by this License, of making, using, or selling its contributor version, 469 | but do not include claims that would be infringed only as a 470 | consequence of further modification of the contributor version. For 471 | purposes of this definition, "control" includes the right to grant 472 | patent sublicenses in a manner consistent with the requirements of 473 | this License. 474 | 475 | Each contributor grants you a non-exclusive, worldwide, royalty-free 476 | patent license under the contributor's essential patent claims, to 477 | make, use, sell, offer for sale, import and otherwise run, modify and 478 | propagate the contents of its contributor version. 479 | 480 | In the following three paragraphs, a "patent license" is any express 481 | agreement or commitment, however denominated, not to enforce a patent 482 | (such as an express permission to practice a patent or covenant not to 483 | sue for patent infringement). To "grant" such a patent license to a 484 | party means to make such an agreement or commitment not to enforce a 485 | patent against the party. 486 | 487 | If you convey a covered work, knowingly relying on a patent license, 488 | and the Corresponding Source of the work is not available for anyone 489 | to copy, free of charge and under the terms of this License, through a 490 | publicly available network server or other readily accessible means, 491 | then you must either (1) cause the Corresponding Source to be so 492 | available, or (2) arrange to deprive yourself of the benefit of the 493 | patent license for this particular work, or (3) arrange, in a manner 494 | consistent with the requirements of this License, to extend the patent 495 | license to downstream recipients. "Knowingly relying" means you have 496 | actual knowledge that, but for the patent license, your conveying the 497 | covered work in a country, or your recipient's use of the covered work 498 | in a country, would infringe one or more identifiable patents in that 499 | country that you have reason to believe are valid. 500 | 501 | If, pursuant to or in connection with a single transaction or 502 | arrangement, you convey, or propagate by procuring conveyance of, a 503 | covered work, and grant a patent license to some of the parties 504 | receiving the covered work authorizing them to use, propagate, modify 505 | or convey a specific copy of the covered work, then the patent license 506 | you grant is automatically extended to all recipients of the covered 507 | work and works based on it. 508 | 509 | A patent license is "discriminatory" if it does not include within 510 | the scope of its coverage, prohibits the exercise of, or is 511 | conditioned on the non-exercise of one or more of the rights that are 512 | specifically granted under this License. You may not convey a covered 513 | work if you are a party to an arrangement with a third party that is 514 | in the business of distributing software, under which you make payment 515 | to the third party based on the extent of your activity of conveying 516 | the work, and under which the third party grants, to any of the 517 | parties who would receive the covered work from you, a discriminatory 518 | patent license (a) in connection with copies of the covered work 519 | conveyed by you (or copies made from those copies), or (b) primarily 520 | for and in connection with specific products or compilations that 521 | contain the covered work, unless you entered into that arrangement, 522 | or that patent license was granted, prior to 28 March 2007. 523 | 524 | Nothing in this License shall be construed as excluding or limiting 525 | any implied license or other defenses to infringement that may 526 | otherwise be available to you under applicable patent law. 527 | 528 | 12. No Surrender of Others' Freedom. 529 | 530 | If conditions are imposed on you (whether by court order, agreement or 531 | otherwise) that contradict the conditions of this License, they do not 532 | excuse you from the conditions of this License. If you cannot convey a 533 | covered work so as to satisfy simultaneously your obligations under this 534 | License and any other pertinent obligations, then as a consequence you may 535 | not convey it at all. For example, if you agree to terms that obligate you 536 | to collect a royalty for further conveying from those to whom you convey 537 | the Program, the only way you could satisfy both those terms and this 538 | License would be to refrain entirely from conveying the Program. 539 | 540 | 13. Remote Network Interaction; Use with the GNU General Public License. 541 | 542 | Notwithstanding any other provision of this License, if you modify the 543 | Program, your modified version must prominently offer all users 544 | interacting with it remotely through a computer network (if your version 545 | supports such interaction) an opportunity to receive the Corresponding 546 | Source of your version by providing access to the Corresponding Source 547 | from a network server at no charge, through some standard or customary 548 | means of facilitating copying of software. This Corresponding Source 549 | shall include the Corresponding Source for any work covered by version 3 550 | of the GNU General Public License that is incorporated pursuant to the 551 | following paragraph. 552 | 553 | Notwithstanding any other provision of this License, you have 554 | permission to link or combine any covered work with a work licensed 555 | under version 3 of the GNU General Public License into a single 556 | combined work, and to convey the resulting work. The terms of this 557 | License will continue to apply to the part which is the covered work, 558 | but the work with which it is combined will remain governed by version 559 | 3 of the GNU General Public License. 560 | 561 | 14. Revised Versions of this License. 562 | 563 | The Free Software Foundation may publish revised and/or new versions of 564 | the GNU Affero General Public License from time to time. Such new versions 565 | will be similar in spirit to the present version, but may differ in detail to 566 | address new problems or concerns. 567 | 568 | Each version is given a distinguishing version number. If the 569 | Program specifies that a certain numbered version of the GNU Affero General 570 | Public License "or any later version" applies to it, you have the 571 | option of following the terms and conditions either of that numbered 572 | version or of any later version published by the Free Software 573 | Foundation. If the Program does not specify a version number of the 574 | GNU Affero General Public License, you may choose any version ever published 575 | by the Free Software Foundation. 576 | 577 | If the Program specifies that a proxy can decide which future 578 | versions of the GNU Affero General Public License can be used, that proxy's 579 | public statement of acceptance of a version permanently authorizes you 580 | to choose that version for the Program. 581 | 582 | Later license versions may give you additional or different 583 | permissions. However, no additional obligations are imposed on any 584 | author or copyright holder as a result of your choosing to follow a 585 | later version. 586 | 587 | 15. Disclaimer of Warranty. 588 | 589 | THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY 590 | APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT 591 | HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY 592 | OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, 593 | THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 594 | PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM 595 | IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF 596 | ALL NECESSARY SERVICING, REPAIR OR CORRECTION. 597 | 598 | 16. Limitation of Liability. 599 | 600 | IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING 601 | WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS 602 | THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY 603 | GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE 604 | USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF 605 | DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD 606 | PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), 607 | EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF 608 | SUCH DAMAGES. 609 | 610 | 17. Interpretation of Sections 15 and 16. 611 | 612 | If the disclaimer of warranty and limitation of liability provided 613 | above cannot be given local legal effect according to their terms, 614 | reviewing courts shall apply local law that most closely approximates 615 | an absolute waiver of all civil liability in connection with the 616 | Program, unless a warranty or assumption of liability accompanies a 617 | copy of the Program in return for a fee. 618 | 619 | END OF TERMS AND CONDITIONS 620 | 621 | How to Apply These Terms to Your New Programs 622 | 623 | If you develop a new program, and you want it to be of the greatest 624 | possible use to the public, the best way to achieve this is to make it 625 | free software which everyone can redistribute and change under these terms. 626 | 627 | To do so, attach the following notices to the program. It is safest 628 | to attach them to the start of each source file to most effectively 629 | state the exclusion of warranty; and each file should have at least 630 | the "copyright" line and a pointer to where the full notice is found. 631 | 632 | 633 | Copyright (C) 634 | 635 | This program is free software: you can redistribute it and/or modify 636 | it under the terms of the GNU Affero General Public License as published 637 | by the Free Software Foundation, either version 3 of the License, or 638 | (at your option) any later version. 639 | 640 | This program is distributed in the hope that it will be useful, 641 | but WITHOUT ANY WARRANTY; without even the implied warranty of 642 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 643 | GNU Affero General Public License for more details. 644 | 645 | You should have received a copy of the GNU Affero General Public License 646 | along with this program. If not, see . 647 | 648 | Also add information on how to contact you by electronic and paper mail. 649 | 650 | If your software can interact with users remotely through a computer 651 | network, you should also make sure that it provides a way for users to 652 | get its source. For example, if your program is a web application, its 653 | interface could display a "Source" link that leads users to an archive 654 | of the code. There are many ways you could offer source, and different 655 | solutions will be better for different programs; see section 13 for the 656 | specific requirements. 657 | 658 | You should also get your employer (if you work as a programmer) or school, 659 | if any, to sign a "copyright disclaimer" for the program, if necessary. 660 | For more information on this, and how to apply and follow the GNU AGPL, see 661 | . 662 | --------------------------------------------------------------------------------