diff --git a/docs/GET_STARTED.md b/docs/GET_STARTED.md
index 31dd91e..7e41955 100644
--- a/docs/GET_STARTED.md
+++ b/docs/GET_STARTED.md
@@ -1,2 +1,58 @@
# Getting Started
+## Reconstruction
+
+Reconstruction metrics follow the same WorldLens YAML and evaluator flow as other
+video-generation metrics. They read 4DGS assets that already exist under
+`generated_results/reconstruction` and compute metrics in the normal WorldLens
+environment.
+
+See `worldbench/videogen/reconstruction/README.md` for the full Reconstruction
+usage guide.
+
+For new videos, first run the vendored DriveStudio Reconstruction backend to
+prepare assets. The backend runs on top of the WorldLens environment with the
+Reconstruction extras installed. It trains 4DGS, renders train/depth/novel
+outputs, and writes the asset layout consumed by WorldLens.
+
+Install the Reconstruction extras after setting up the base WorldLens
+environment:
+
+```bash
+pip install -r worldbench/videogen/reconstruction/requirements-reconstruction-extra.txt
+```
+
+Example asset preparation:
+
+```bash
+cd /path/to/WorldLens
+python worldbench/videogen/reconstruction/prepare_assets.py \
+ --method-name magicdrive \
+ --clips 081 \
+ --reconstruction-root generated_results/reconstruction \
+ --work-root work_dirs \
+ --data-root data/nuscenes_trainval/raw \
+ --target-dir data/nuscenes_trainval/processed
+```
+
+Example WorldLens metric:
+
+```yaml
+reconstruction:
+ photometric_error:
+ - name: photometric_error
+ method_name: ${method_name}
+ reconstruction_root: generated_results/reconstruction
+ clips: ["081"]
+```
+
+DriveStudio training expects generated-data preprocessing dependencies to be
+available under the configured data roots, including fine dynamic masks and
+`humanpose/smpl.pkl` for each clip. The backend validates those inputs before
+launching training so missing preprocessing is reported before a long 4DGS run.
+Generated videos should use the DriveStudio `nuscenes_gen` layout for the
+selected `method_name`.
+
+`geometric_discrepancy` also needs reference/GT depth assets and GT
+Grounded-SAM2 evaluation masks. RGB videos alone are not sufficient for that
+metric.
diff --git a/tools/configs/videogen.yaml b/tools/configs/videogen.yaml
index c25000c..dab4beb 100644
--- a/tools/configs/videogen.yaml
+++ b/tools/configs/videogen.yaml
@@ -79,32 +79,29 @@ dimensions:
repeat_times: 1
# reconstruction:
-
- # novel_view_fidelity:
- # - name: novel_view_fidelity
+
+ # photometric_error:
+ # - name: photometric_error
# method_name: ${method_name}
- # need_preprocessing: true
- # method: videogpt
- # generated_data_path: ${generated_data_path}
- # pretrained_model_path: pretrained_models/fvd/i3d_pretrained_400.pt
- # data_info_path: data/nuscenes_mmdet3d-12Hz/nuscenes_interp_12Hz_infos_track2_eval.pkl
+ # reconstruction_root: generated_results/reconstruction
+ # clips: ["081"]
- # novel_view_fidelity:
- # - name: novel_view_fidelity
+ # geometric_discrepancy:
+ # - name: geometric_discrepancy
# method_name: ${method_name}
- # need_preprocessing: true
- # generated_data_path: ${generated_data_path}
+ # reconstruction_root: generated_results/reconstruction
+ # clips: ["081"]
- # novel_view_consistency:
- # - name: novel_view_consistency
+ # novel_view_fidelity:
+ # - name: novel_view_fidelity
# method_name: ${method_name}
- # need_preprocessing: true
- # repeat_times: 1
- # local_save_path: pretrained_models/dino/dino_vitbase16_pretrain.pth
+ # reconstruction_root: generated_results/reconstruction
+ # clips: ["081"]
# novel_view_consistency:
# - name: novel_view_consistency
# method_name: ${method_name}
- # need_preprocessing: true
+ # reconstruction_root: generated_results/reconstruction
+ # clips: ["081"]
# repeat_times: 1
- # local_save_path: pretrained_models/clip/ViT-B-32.pt
\ No newline at end of file
+ # local_save_path: pretrained_models/clip/ViT-B-32.pt
diff --git a/worldbench/third_party/drivestudio/LICENSE b/worldbench/third_party/drivestudio/LICENSE
new file mode 100644
index 0000000..ada6ec8
--- /dev/null
+++ b/worldbench/third_party/drivestudio/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Ziyu Chen
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/README.md b/worldbench/third_party/drivestudio/README.md
new file mode 100644
index 0000000..78b65fd
--- /dev/null
+++ b/worldbench/third_party/drivestudio/README.md
@@ -0,0 +1,227 @@
+
+
+
+
+
+
+
+A 3DGS framework for omni urban scene reconstruction and simulation!
+
+
+
+
+
+
+
+
+
+
+
+## About
+DriveStudio is a 3DGS codebase for urban scene reconstruction/simulation. It offers a system with multiple Gaussian representations to jointly reconstruct backgrounds, vehicles, and non-rigid categories (pedestrians, cyclists, etc.) from driving logs. DriveStudio also provides a unified data system supporting various popular driving datasets, including [Waymo](https://waymo.com/open/), [PandaSet](https://pandaset.org/), [Argoverse2](https://www.argoverse.org/av2.html), [KITTI](http://www.cvlibs.net/datasets/kitti/), [NuScenes](https://www.nuscenes.org/), and [NuPlan](https://www.nuscenes.org/nuplan).
+
+This codebase also contains the **official implementation** of:
+ > **OmniRe: Omni Urban Scene Reconstruction**
+Click to expand functionality details
+
+We have implemented interesting and useful functionalities:
+
+1. **Flexible multi-camera training:** Choose any combination of cameras for training - single, multiple, or all. You can set these up by **SIMPLY** configuring your selection in the config file.
+
+2. **Powered by gsplat** Integrated [gsplat](https://github.com/nerfstudio-project/gsplat) rasterization kernel with its advanced functions, e.g. absolute gradients, anti-aliasing, etc.
+
+3. **Camera Pose Refinement:** Recognizing that camera poses may not always be sufficiently accurate, we provide a method to refine and optimize these poses.
+
+4. **Objects' GT Bounding Box Refinement:** To address noise in ground truth boxes, we've added this feature to further improve accuracy and robustness.
+
+5. **Affine Transformation:** This feature handles camera exposure and other related issues, enhancing the quality of scene reconstruction.
+
+6. ...
+
+These functionalities are designed to enhance the overall performance and flexibility of our system, allowing for more accurate and adaptable scene reconstruction across various datasets and conditions.
+
+
+## 📢 Updates
+
+**[Jan 2025]** OmniRe is accepted by ICLR 2025 as a spotlight paper!
+
+**[Aug 2024]** Release code of DriveStudio!
+
+## 🔨 Installation
+
+Run the following commands to set up the environment:
+
+```shell
+# Clone the repository with submodules
+git clone --recursive https://github.com/ziyc/drivestudio.git
+cd drivestudio
+
+# Create the environment
+conda create -n drivestudio python=3.9 -y
+conda activate drivestudio
+pip install -r requirements.txt
+pip install git+https://github.com/nerfstudio-project/gsplat.git@v1.3.0
+pip install git+https://github.com/facebookresearch/pytorch3d.git
+pip install git+https://github.com/NVlabs/nvdiffrast
+
+# Set up for SMPL Gaussians
+cd third_party/smplx/
+pip install -e .
+cd ../..
+```
+
+## 📊 Prepare Data
+We support most popular public driving datasets. Detailed instructions for downloading and processing each dataset are available in the following documents:
+
+- Waymo: [Data Process Instruction](docs/Waymo.md)
+- NuScenes: [Data Process Instruction](docs/NuScenes.md)
+- NuPlan: [Data Process Instruction](docs/Nuplan.md)
+- ArgoVerse: [Data Process Instruction](docs/ArgoVerse.md)
+- PandaSet: [Data Process Instruction](docs/Pandaset.md)
+- KITTI: [Data Process Instruction](docs/KITTI.md)
+
+## 🚀 Running
+### Training
+```shell
+export PYTHONPATH=$(pwd)
+start_timestep=0 # start frame index for training
+end_timestep=-1 # end frame index, -1 for the last frame
+
+python tools/train.py \
+ --config_file configs/omnire.yaml \
+ --output_root $output_root \
+ --project $project \
+ --run_name $expname \
+ dataset=waymo/3cams \
+ data.scene_idx=$scene_idx \
+ data.start_timestep=$start_timestep \
+ data.end_timestep=$end_timestep
+```
+
+- To run other methods, change `--config_file`. See `configs/` for more options.
+- Specify dataset and number of cameras by setting `dataset`. Examples: `waymo/1cams`, `waymo/5cams`, `pandaset/6cams`, `argoverse/7cams`, etc.
+ You can set up arbitrary camera combinations for each dataset. See `configs/datasets/` for custom configuration details.
+- For over 3 cameras or 450+ images, we recommend using `omnire_extended_cam.yaml`. It works better in practice.
+### Evaluation
+```shell
+python tools/eval.py --resume_from $ckpt_path
+```
+
+
+Compare with OmniRe
+To compare with OmniRe, run the following commands to reproduce the paper’s checkpoints.
+
+```
+# Scene Reconstruction
+for scene_idx in {0..7}; do
+ IFS=',' read scene_idx seg_name start_timestep end_timestep <<< $(awk -v idx="$(($scene_idx + 2))" \
+ 'NR==idx {print $1, $2, $3, $4}' data/waymo_example_scenes.txt)
+
+ python tools/train.py \
+ --config_file configs/paper_legacy/omnire.yaml \
+ --output_root logs/omnire_waymo/ \
+ --project recon \
+ --run_name ${scene_idx} \
+ data.scene_idx=$scene_idx \
+ data.start_timestep=$start_timestep \
+ data.end_timestep=$end_timestep \
+ data.pixel_source.test_image_stride=0
+done
+```
+
+```
+# Novel View Synthesis
+for scene_idx in {0..7}; do
+ IFS=',' read scene_idx seg_name start_timestep end_timestep <<< $(awk -v idx="$(($scene_idx + 2))" \
+ 'NR==idx {print $1, $2, $3, $4}' data/waymo_example_scenes.txt)
+
+ python tools/train.py \
+ --config_file configs/paper_legacy/omnire.yaml \
+ --output_root logs/omnire_waymo/ \
+ --project nvs \
+ --run_name ${scene_idx} \
+ data.scene_idx=$scene_idx \
+ data.start_timestep=$start_timestep \
+ data.end_timestep=$end_timestep \
+ data.pixel_source.test_image_stride=10
+done
+```
+
+and run `utils/gather_results.py` to get the results.
+
+
+## 👏 Contributions
+We're improving our project to develop a robust driving recom/sim system. Some areas we're focusing on:
+
+- A real-time viewer for background and foreground visualization
+- Scene editing and simulation tools
+- Other Gaussian representations (e.g., 2DGS, surfels)
+
+We welcome pull requests and collaborations. If you'd like to contribute or have questions, feel free to open an issue or contact [Ziyu Chen](https://github.com/ziyc) (ziyu.sjtu@gmail.com).
+
+## 🙏 Acknowledgments
+We utilize the rasterization kernel from [gsplat](https://github.com/nerfstudio-project/gsplat). Parts of our implementation are based on work from [EmerNeRF](https://github.com/NVlabs/EmerNeRF), [NerfStudio](https://github.com/nerfstudio-project/nerfstudio), [GART](https://github.com/JiahuiLei/GART), and [Neuralsim](https://github.com/PJLab-ADG/neuralsim). We've also implemented unofficial versions of [Deformable-GS](https://github.com/ingra14m/Deformable-3D-Gaussians), [PVG](https://github.com/fudan-zvg/PVG), and [Street Gaussians](https://github.com/zju3dv/street_gaussians), with reference to their original codebases.
+
+We extend our deepest gratitude to the authors for their contributions to the community, which have greatly supported our research.
+
+## Citation
+If you find this codebase helpful, please kindly cite:
+```
+@inproceedings{
+ chen2025omnire,
+ title={OmniRe: Omni Urban Scene Reconstruction},
+ author={Ziyu Chen and Jiawei Yang and Jiahui Huang and Riccardo de Lutio and Janick Martinez Esturo and Boris Ivanovic and Or Litany and Zan Gojcic and Sanja Fidler and Marco Pavone and Li Song and Yue Wang},
+ booktitle={The Thirteenth International Conference on Learning Representations},
+ year={2025}
+}
+```
diff --git a/worldbench/third_party/drivestudio/configs/datasets/argoverse/1cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/argoverse/1cams.yaml
new file mode 100644
index 0000000..16ef7ef
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/argoverse/1cams.yaml
@@ -0,0 +1,42 @@
+# ArgoVerse2 dataset configuration
+# Processed cameras:
+# idx camera original size
+# 0 ring_front_center (1550, 2048)
+# 1 ring_front_left (2048, 1550)
+# 2 ring_front_right (2048, 1550)
+# 3 ring_side_left (2048, 1550)
+# 4 ring_side_right (2048, 1550)
+# 5 ring_rear_left (2048, 1550)
+# 6 ring_rear_right (2048, 1550)
+
+data:
+ data_root: data/argoverse/processed/training # data root for the dataset
+ dataset: argoverse # dataset type
+ scene_idx: 0 # which scene to use
+ start_timestep: 0 # which timestep to start from
+ end_timestep: -1 # which timestep to end at, -1 means the last timestep
+ preload_device: cuda # choose from ["cpu", "cuda"], cache the data on this device
+ pixel_source: # image source and object annotations
+ type: datasets.argoverse.argoverse_sourceloader.ArgoVersePixelSource
+ cameras: [0] # which cameras to use
+ downscale_when_loading: [2] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: False # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.argoverse.argoverse_sourceloader.ArgoVerseLiDARSource
+ load_lidar: True # whether to load lidar
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/datasets/argoverse/3cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/argoverse/3cams.yaml
new file mode 100644
index 0000000..919b1f8
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/argoverse/3cams.yaml
@@ -0,0 +1,42 @@
+# ArgoVerse2 dataset configuration
+# Processed cameras:
+# idx camera original size
+# 0 ring_front_center (1550, 2048)
+# 1 ring_front_left (2048, 1550)
+# 2 ring_front_right (2048, 1550)
+# 3 ring_side_left (2048, 1550)
+# 4 ring_side_right (2048, 1550)
+# 5 ring_rear_left (2048, 1550)
+# 6 ring_rear_right (2048, 1550)
+
+data:
+ data_root: data/argoverse/processed/training # data root for the dataset
+ dataset: argoverse # dataset type
+ scene_idx: 0 # which scene to use
+ start_timestep: 0 # which timestep to start from
+ end_timestep: -1 # which timestep to end at, -1 means the last timestep
+ preload_device: cuda # choose from ["cpu", "cuda"], cache the data on this device
+ pixel_source: # image source and object annotations
+ type: datasets.argoverse.argoverse_sourceloader.ArgoVersePixelSource
+ cameras: [0, 1, 2] # which cameras to use
+ downscale_when_loading: [2, 2, 2] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: False # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.argoverse.argoverse_sourceloader.ArgoVerseLiDARSource
+ load_lidar: True # whether to load lidar
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/datasets/argoverse/5cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/argoverse/5cams.yaml
new file mode 100644
index 0000000..2c193d9
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/argoverse/5cams.yaml
@@ -0,0 +1,42 @@
+# ArgoVerse2 dataset configuration
+# Processed cameras:
+# idx camera original size
+# 0 ring_front_center (1550, 2048)
+# 1 ring_front_left (2048, 1550)
+# 2 ring_front_right (2048, 1550)
+# 3 ring_side_left (2048, 1550)
+# 4 ring_side_right (2048, 1550)
+# 5 ring_rear_left (2048, 1550)
+# 6 ring_rear_right (2048, 1550)
+
+data:
+ data_root: data/argoverse/processed/training # data root for the dataset
+ dataset: argoverse # dataset type
+ scene_idx: 0 # which scene to use
+ start_timestep: 0 # which timestep to start from
+ end_timestep: -1 # which timestep to end at, -1 means the last timestep
+ preload_device: cpu # choose from ["cpu", "cuda"], cache the data on this device
+ pixel_source: # image source and object annotations
+ type: datasets.argoverse.argoverse_sourceloader.ArgoVersePixelSource
+ cameras: [0, 1, 2, 3, 4] # which cameras to use
+ downscale_when_loading: [3, 3, 3, 3, 3] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: False # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.argoverse.argoverse_sourceloader.ArgoVerseLiDARSource
+ load_lidar: True # whether to load lidar
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/datasets/argoverse/7cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/argoverse/7cams.yaml
new file mode 100644
index 0000000..ce9cda0
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/argoverse/7cams.yaml
@@ -0,0 +1,42 @@
+# ArgoVerse2 dataset configuration
+# Processed cameras:
+# idx camera original size
+# 0 ring_front_center (1550, 2048)
+# 1 ring_front_left (2048, 1550)
+# 2 ring_front_right (2048, 1550)
+# 3 ring_side_left (2048, 1550)
+# 4 ring_side_right (2048, 1550)
+# 5 ring_rear_left (2048, 1550)
+# 6 ring_rear_right (2048, 1550)
+
+data:
+ data_root: data/argoverse/processed/training # data root for the dataset
+ dataset: argoverse # dataset type
+ scene_idx: 0 # which scene to use
+ start_timestep: 0 # which timestep to start from
+ end_timestep: -1 # which timestep to end at, -1 means the last timestep
+ preload_device: cpu # choose from ["cpu", "cuda"], cache the data on this device
+ pixel_source: # image source and object annotations
+ type: datasets.argoverse.argoverse_sourceloader.ArgoVersePixelSource
+ cameras: [0, 1, 2, 3, 4, 5, 6] # which cameras to use
+ downscale_when_loading: [3, 3, 3, 3, 3, 3, 3] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: False # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.argoverse.argoverse_sourceloader.ArgoVerseLiDARSource
+ load_lidar: True # whether to load lidar
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/datasets/kitti/1cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/kitti/1cams.yaml
new file mode 100644
index 0000000..8ba7251
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/kitti/1cams.yaml
@@ -0,0 +1,43 @@
+# KITTI dataset configuration
+# Processed cameras:
+# idx camera original size
+# 0 CAM_LEFT (375, 1242)
+# 1 CAM_RIGHT (375, 1242)
+
+data:
+ data_root: data/kitti/processed # data root for the dataset
+ dataset: kitti # dataset type
+ scene_idx: 2011_09_26_drive_0091_sync # which scene to use, specified by date and drive number
+ start_timestep: 0 # which timestep to start from
+ end_timestep: 200 # which timestep to end at, -1 means the last timestep
+ preload_device: cuda # choose from ["cpu", "cuda"], cache the data on this device
+ pixel_source: # image source and object annotations
+ type: datasets.kitti.kitti_sourceloader.KITTIPixelSource
+ cameras: [0] # which cameras to use
+ downscale_when_loading: [1] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: False # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians (not available for KITTI)
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.kitti.kitti_sourceloader.KITTILiDARSource
+ load_lidar: True # whether to load lidar
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
+
+# dataset specific configurations
+model:
+ RigidNodes:
+ ctrl:
+ stop_split_at: 30000
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/datasets/kitti/2cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/kitti/2cams.yaml
new file mode 100644
index 0000000..c1659b0
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/kitti/2cams.yaml
@@ -0,0 +1,43 @@
+# KITTI dataset configuration
+# Processed cameras:
+# idx camera original size
+# 0 CAM_LEFT (375, 1242)
+# 1 CAM_RIGHT (375, 1242)
+
+data:
+ data_root: data/kitti/processed # data root for the dataset
+ dataset: kitti # dataset type
+ scene_idx: 2011_09_26_drive_0091_sync # which scene to use, specified by date and drive number
+ start_timestep: 0 # which timestep to start from
+ end_timestep: 200 # which timestep to end at, -1 means the last timestep
+ preload_device: cuda # choose from ["cpu", "cuda"], cache the data on this device
+ pixel_source: # image source and object annotations
+ type: datasets.kitti.kitti_sourceloader.KITTIPixelSource
+ cameras: [0, 1] # which cameras to use
+ downscale_when_loading: [2, 2] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: False # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians (not available for KITTI)
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.kitti.kitti_sourceloader.KITTILiDARSource
+ load_lidar: True # whether to load lidar
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
+
+# dataset specific configurations
+model:
+ RigidNodes:
+ ctrl:
+ stop_split_at: 30000
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/datasets/nuplan/1cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/nuplan/1cams.yaml
new file mode 100644
index 0000000..bd97b54
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/nuplan/1cams.yaml
@@ -0,0 +1,43 @@
+# NuPlan dataset configuration
+# Processed cameras:
+# idx camera original size egocar_visible
+# 0 CAM_F0 (1080, 1920) False
+# 1 CAM_L0 (1080, 1920) True
+# 2 CAM_R0 (1080, 1920) True
+# 3 CAM_L1 (1080, 1920) False
+# 4 CAM_R1 (1080, 1920) False
+# 5 CAM_L2 (1080, 1920) True
+# 6 CAM_R2 (1080, 1920) True
+# 7 CAM_B0 (1080, 1920) False
+
+data:
+ data_root: data/nuplan/processed/mini # data root for the dataset
+ dataset: nuplan # dataset type
+ scene_idx: 2021.05.12.22.00.38_veh-35_01008_01518 # which scene to use, specific to NuPlan's naming convention
+ start_timestep: 0 # which timestep to start from
+ end_timestep: -1 # which timestep to end at, -1 means the last timestep
+ preload_device: cuda # choose from ["cpu", "cuda"], cache the data on this device
+ pixel_source: # image source and object annotations
+ type: datasets.nuplan.nuplan_sourceloader.NuPlanPixelSource
+ cameras: [0] # which cameras to use
+ downscale_when_loading: [2] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: False # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.nuplan.nuplan_sourceloader.NuPlanLiDARSource
+ load_lidar: True # whether to load lidar
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/datasets/nuplan/3cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/nuplan/3cams.yaml
new file mode 100644
index 0000000..ef3a348
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/nuplan/3cams.yaml
@@ -0,0 +1,43 @@
+# NuPlan dataset configuration
+# Processed cameras:
+# idx camera original size egocar_visible
+# 0 CAM_F0 (1080, 1920) False
+# 1 CAM_L0 (1080, 1920) True
+# 2 CAM_R0 (1080, 1920) True
+# 3 CAM_L1 (1080, 1920) False
+# 4 CAM_R1 (1080, 1920) False
+# 5 CAM_L2 (1080, 1920) True
+# 6 CAM_R2 (1080, 1920) True
+# 7 CAM_B0 (1080, 1920) False
+
+data:
+ data_root: data/nuplan/processed/mini # data root for the dataset
+ dataset: nuplan # dataset type
+ scene_idx: 2021.05.12.22.00.38_veh-35_01008_01518 # which scene to use, specific to NuPlan's naming convention
+ start_timestep: 0 # which timestep to start from
+ end_timestep: -1 # which timestep to end at, -1 means the last timestep
+ preload_device: cuda # choose from ["cpu", "cuda"], cache the data on this device
+ pixel_source: # image source and object annotations
+ type: datasets.nuplan.nuplan_sourceloader.NuPlanPixelSource
+ cameras: [0, 1, 2] # which cameras to use
+ downscale_when_loading: [2, 2, 2] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: False # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.nuplan.nuplan_sourceloader.NuPlanLiDARSource
+ load_lidar: True # whether to load lidar
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/datasets/nuplan/5cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/nuplan/5cams.yaml
new file mode 100644
index 0000000..3be9185
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/nuplan/5cams.yaml
@@ -0,0 +1,43 @@
+# NuPlan dataset configuration
+# Processed cameras:
+# idx camera original size egocar_visible
+# 0 CAM_F0 (1080, 1920) False
+# 1 CAM_L0 (1080, 1920) True
+# 2 CAM_R0 (1080, 1920) True
+# 3 CAM_L1 (1080, 1920) False
+# 4 CAM_R1 (1080, 1920) False
+# 5 CAM_L2 (1080, 1920) True
+# 6 CAM_R2 (1080, 1920) True
+# 7 CAM_B0 (1080, 1920) False
+
+data:
+ data_root: data/nuplan/processed/mini # data root for the dataset
+ dataset: nuplan # dataset type
+ scene_idx: 2021.05.12.22.00.38_veh-35_01008_01518 # which scene to use, specific to NuPlan's naming convention
+ start_timestep: 0 # which timestep to start from
+ end_timestep: -1 # which timestep to end at, -1 means the last timestep
+ preload_device: cpu # choose from ["cpu", "cuda"], cache the data on this device
+ pixel_source: # image source and object annotations
+ type: datasets.nuplan.nuplan_sourceloader.NuPlanPixelSource
+ cameras: [0, 1, 2, 3, 4] # which cameras to use
+ downscale_when_loading: [3, 3, 3, 3, 3] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: False # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.nuplan.nuplan_sourceloader.NuPlanLiDARSource
+ load_lidar: True # whether to load lidar
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/datasets/nuplan/8cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/nuplan/8cams.yaml
new file mode 100644
index 0000000..06dd536
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/nuplan/8cams.yaml
@@ -0,0 +1,43 @@
+# NuPlan dataset configuration
+# Processed cameras:
+# idx camera original size egocar_visible
+# 0 CAM_F0 (1080, 1920) False
+# 1 CAM_L0 (1080, 1920) True
+# 2 CAM_R0 (1080, 1920) True
+# 3 CAM_L1 (1080, 1920) False
+# 4 CAM_R1 (1080, 1920) False
+# 5 CAM_L2 (1080, 1920) True
+# 6 CAM_R2 (1080, 1920) True
+# 7 CAM_B0 (1080, 1920) False
+
+data:
+ data_root: data/nuplan/processed/mini # data root for the dataset
+ dataset: nuplan # dataset type
+ scene_idx: 2021.05.12.22.00.38_veh-35_01008_01518 # which scene to use, specific to NuPlan's naming convention
+ start_timestep: 0 # which timestep to start from
+ end_timestep: -1 # which timestep to end at, -1 means the last timestep
+ preload_device: cpu # choose from ["cpu", "cuda"], cache the data on this device
+ pixel_source: # image source and object annotations
+ type: datasets.nuplan.nuplan_sourceloader.NuPlanPixelSource
+ cameras: [0, 1, 2, 3, 4, 5, 6, 7] # which cameras to use
+ downscale_when_loading: [3, 3, 3, 3, 3, 3, 3, 3] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: False # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.nuplan.nuplan_sourceloader.NuPlanLiDARSource
+ load_lidar: True # whether to load lidar
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/datasets/nuscenes/1cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/nuscenes/1cams.yaml
new file mode 100644
index 0000000..986dbcb
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/nuscenes/1cams.yaml
@@ -0,0 +1,41 @@
+# NuScenes dataset configuration
+# Processed cameras:
+# idx camera original size
+# 0 CAM_FRONT (900, 1600)
+# 1 CAM_FRONT_LEFT (900, 1600)
+# 2 CAM_FRONT_RIGHT (900, 1600)
+# 3 CAM_BACK_LEFT (900, 1600)
+# 4 CAM_BACK_RIGHT (900, 1600)
+# 5 CAM_BACK (900, 1600)
+
+data:
+ data_root: data/nuscenes/processed_10Hz/mini # data root for the dataset
+ dataset: nuscenes # dataset type
+ scene_idx: 0 # which scene to use, [0, 849] for nuscenes's train/val sets, inclusive
+ start_timestep: 0 # which timestep to start from
+ end_timestep: -1 # which timestep to end at, -1 means the last timestep
+ preload_device: cuda # choose from ["cpu", "cuda"], cache the data on this device
+ pixel_source: # image source and object annotations
+ type: datasets.nuscenes.nuscenes_sourceloader.NuScenesPixelSource
+ cameras: [0] # which cameras to use
+ downscale_when_loading: [2] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: False # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.nuscenes.nuscenes_sourceloader.NuScenesLiDARSource
+ load_lidar: True # whether to load lidar
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/datasets/nuscenes/3cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/nuscenes/3cams.yaml
new file mode 100644
index 0000000..1f9fea8
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/nuscenes/3cams.yaml
@@ -0,0 +1,41 @@
+# NuScenes dataset configuration
+# Processed cameras:
+# idx camera original size
+# 0 CAM_FRONT (900, 1600)
+# 1 CAM_FRONT_LEFT (900, 1600)
+# 2 CAM_FRONT_RIGHT (900, 1600)
+# 3 CAM_BACK_LEFT (900, 1600)
+# 4 CAM_BACK_RIGHT (900, 1600)
+# 5 CAM_BACK (900, 1600)
+
+data:
+ data_root: data/nuscenes/processed_10Hz/mini # data root for the dataset
+ dataset: nuscenes # dataset type
+ scene_idx: 0 # which scene to use, [0, 849] for nuscenes's train/val sets, inclusive
+ start_timestep: 0 # which timestep to start from
+ end_timestep: -1 # which timestep to end at, -1 means the last timestep
+ preload_device: cuda # choose from ["cpu", "cuda"], cache the data on this device
+ pixel_source: # image source and object annotations
+ type: datasets.nuscenes.nuscenes_sourceloader.NuScenesPixelSource
+ cameras: [0, 1, 2] # which cameras to use
+ downscale_when_loading: [2, 2, 2] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: False # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.nuscenes.nuscenes_sourceloader.NuScenesLiDARSource
+ load_lidar: True # whether to load lidar
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/datasets/nuscenes/6cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/nuscenes/6cams.yaml
new file mode 100644
index 0000000..b165ede
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/nuscenes/6cams.yaml
@@ -0,0 +1,41 @@
+# NuScenes dataset configuration
+# Processed cameras:
+# idx camera original size
+# 0 CAM_FRONT (900, 1600)
+# 1 CAM_FRONT_LEFT (900, 1600)
+# 2 CAM_FRONT_RIGHT (900, 1600)
+# 3 CAM_BACK_LEFT (900, 1600)
+# 4 CAM_BACK_RIGHT (900, 1600)
+# 5 CAM_BACK (900, 1600)
+
+data:
+ data_root: data/nuscenes/processed_10Hz/mini # data root for the dataset
+ dataset: nuscenes # dataset type
+ scene_idx: 0 # which scene to use, [0, 849] for nuscenes's train/val sets, inclusive
+ start_timestep: 0 # which timestep to start from
+ end_timestep: -1 # which timestep to end at, -1 means the last timestep
+ preload_device: cpu # choose from ["cpu", "cuda"], cache the data on this device
+ pixel_source: # image source and object annotations
+ type: datasets.nuscenes.nuscenes_sourceloader.NuScenesPixelSource
+ cameras: [0, 1, 2, 3, 4, 5] # which cameras to use
+ downscale_when_loading: [3, 3, 3, 3, 3, 3] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: False # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.nuscenes.nuscenes_sourceloader.NuScenesLiDARSource
+ load_lidar: True # whether to load lidar
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/datasets/nuscenes_gen/6cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/nuscenes_gen/6cams.yaml
new file mode 100644
index 0000000..efb11d2
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/nuscenes_gen/6cams.yaml
@@ -0,0 +1,42 @@
+# NuScenes dataset configuration
+# Processed cameras:
+# idx camera original size
+# 0 CAM_FRONT (900, 1600)
+# 1 CAM_FRONT_LEFT (900, 1600)
+# 2 CAM_FRONT_RIGHT (900, 1600)
+# 3 CAM_BACK_LEFT (900, 1600)
+# 4 CAM_BACK_RIGHT (900, 1600)
+# 5 CAM_BACK (900, 1600)
+
+data:
+ # data_root: data/nuscenes_trainval/processed/advanced_12Hz_trainval # data root for the dataset
+ data_root: None
+ dataset: nuscenes # dataset type
+ scene_idx: 0 # which scene to use, [0, 849] for nuscenes's train/val sets, inclusive
+ start_timestep: 0 # which timestep to start from
+ end_timestep: -1 # which timestep to end at, -1 means the last timestep
+ preload_device: cpu # choose from ["cpu", "cuda"], cache the data on this device
+ pixel_source: # image source and object annotations
+ type: datasets.nuscenes.nuscenes_sourceloader.NuScenesPixelSource
+ cameras: [0, 1, 2, 3, 4, 5] # which cameras to use
+ downscale_when_loading: [3, 3, 3, 3, 3, 3] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: False # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.nuscenes.nuscenes_sourceloader.NuScenesLiDARSource
+ load_lidar: True # whether to load lidar
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/datasets/pandaset/1cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/pandaset/1cams.yaml
new file mode 100644
index 0000000..79278ff
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/pandaset/1cams.yaml
@@ -0,0 +1,42 @@
+# Pandaset dataset configuration
+# Processed cameras:
+# idx camera original size
+# 0 front_camera (1920, 1080)
+# 1 front_left_camera (1920, 1080)
+# 2 front_right_camera (1920, 1080)
+# 3 left_camera (1920, 1080)
+# 4 right_camera (1920, 1080)
+# 5 back_camera (1920, 1080)
+
+data:
+ data_root: data/pandaset/processed # data root for the dataset
+ dataset: pandaset # dataset type
+ scene_idx: 0 # which scene to use
+ start_timestep: 0 # which timestep to start from
+ end_timestep: -1 # which timestep to end at, -1 means the last timestep
+ preload_device: cuda # choose from ["cpu", "cuda"], cache the data on this device
+ pixel_source: # image source and object annotations
+ type: datasets.pandaset.pandaset_sourceloader.PandaPixelSource
+ cameras: [0] # which cameras to use
+ downscale_when_loading: [1] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: False # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.pandaset.pandaset_sourceloader.PandaLiDARSource
+ load_lidar: True # whether to load lidar
+ only_use_360_lidar: False # whether to only use the 360-degree lidar
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/datasets/pandaset/3cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/pandaset/3cams.yaml
new file mode 100644
index 0000000..9d9c401
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/pandaset/3cams.yaml
@@ -0,0 +1,42 @@
+# Pandaset dataset configuration
+# Processed cameras:
+# idx camera original size
+# 0 front_camera (1920, 1080)
+# 1 front_left_camera (1920, 1080)
+# 2 front_right_camera (1920, 1080)
+# 3 left_camera (1920, 1080)
+# 4 right_camera (1920, 1080)
+# 5 back_camera (1920, 1080)
+
+data:
+ data_root: data/pandaset/processed # data root for the dataset
+ dataset: pandaset # dataset type
+ scene_idx: 0 # which scene to use
+ start_timestep: 0 # which timestep to start from
+ end_timestep: -1 # which timestep to end at, -1 means the last timestep
+ preload_device: cuda # choose from ["cpu", "cuda"], cache the data on this device
+ pixel_source: # image source and object annotations
+ type: datasets.pandaset.pandaset_sourceloader.PandaPixelSource
+ cameras: [0, 1, 2] # which cameras to use
+ downscale_when_loading: [2, 2, 2] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: False # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.pandaset.pandaset_sourceloader.PandaLiDARSource
+ load_lidar: True # whether to load lidar
+ only_use_360_lidar: False # whether to only use the 360-degree lidar
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/datasets/pandaset/6cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/pandaset/6cams.yaml
new file mode 100644
index 0000000..2cf1b6a
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/pandaset/6cams.yaml
@@ -0,0 +1,42 @@
+# Pandaset dataset configuration
+# Processed cameras:
+# idx camera original size
+# 0 front_camera (1920, 1080)
+# 1 front_left_camera (1920, 1080)
+# 2 front_right_camera (1920, 1080)
+# 3 left_camera (1920, 1080)
+# 4 right_camera (1920, 1080)
+# 5 back_camera (1920, 1080)
+
+data:
+ data_root: data/pandaset/processed # data root for the dataset
+ dataset: pandaset # dataset type
+ scene_idx: 0 # which scene to use
+ start_timestep: 0 # which timestep to start from
+ end_timestep: -1 # which timestep to end at, -1 means the last timestep
+ preload_device: cpu # choose from ["cpu", "cuda"], cache the data on this device
+ pixel_source: # image source and object annotations
+ type: datasets.pandaset.pandaset_sourceloader.PandaPixelSource
+ cameras: [0, 1, 2, 3, 4, 5] # which cameras to use
+ downscale_when_loading: [2, 2, 2, 2, 2, 2] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: False # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.pandaset.pandaset_sourceloader.PandaLiDARSource
+ load_lidar: True # whether to load lidar
+ only_use_360_lidar: False # whether to only use the 360-degree lidar
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/datasets/waymo/1cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/waymo/1cams.yaml
new file mode 100644
index 0000000..5412dbd
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/waymo/1cams.yaml
@@ -0,0 +1,43 @@
+# Waymo dataset configuration
+# Processed cameras:
+# idx camera original size
+# 0 front_camera (1920, 1280)
+# 1 front_left_camera (1920, 1280)
+# 2 front_right_camera (1920, 1280)
+# 3 left_camera (1920, 866)
+# 4 right_camera (1920, 866)
+
+data:
+ data_root: data/waymo/processed/training # data root for the dataset
+ dataset: waymo
+ scene_idx: 0 # which scene to use, [0, 798] for waymo's training set and [0, 849] for nuscenes's train/val sets, inclusive
+ start_timestep: 0 # which timestep to start from
+ end_timestep: -1 # which timestep to end at, -1 means the last timestep
+ preload_device: cuda # choose from ["cpu", "cuda"], cache the data on this device.
+ pixel_source: # image source and object annotations
+ type: datasets.waymo.waymo_sourceloader.WaymoPixelSource
+ cameras: [0] # which cameras to use
+ downscale_when_loading: [1] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: True # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.waymo.waymo_sourceloader.WaymoLiDARSource
+ load_lidar: True # whether to load lidar
+ only_use_top_lidar: False # whether to only use the top lidar
+ truncated_max_range: 80 # max range for truncated lidar in a ego-centric coordinate system
+ truncated_min_range: -2 # min range for truncated lidar in a ego-centric coordinate system.
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/datasets/waymo/3cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/waymo/3cams.yaml
new file mode 100644
index 0000000..26bc035
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/waymo/3cams.yaml
@@ -0,0 +1,43 @@
+# Waymo dataset configuration
+# Processed cameras:
+# idx camera original size
+# 0 front_camera (1920, 1280)
+# 1 front_left_camera (1920, 1280)
+# 2 front_right_camera (1920, 1280)
+# 3 left_camera (1920, 866)
+# 4 right_camera (1920, 866)
+
+data:
+ data_root: data/waymo/processed/training # data root for the dataset
+ dataset: waymo
+ scene_idx: 0 # which scene to use, [0, 798] for waymo's training set and [0, 849] for nuscenes's train/val sets, inclusive
+ start_timestep: 0 # which timestep to start from
+ end_timestep: -1 # which timestep to end at, -1 means the last timestep
+ preload_device: cuda # choose from ["cpu", "cuda"], cache the data on this device.
+ pixel_source: # image source and object annotations
+ type: datasets.waymo.waymo_sourceloader.WaymoPixelSource
+ cameras: [0, 1, 2] # which cameras to use
+ downscale_when_loading: [2, 2, 2] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: True # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.waymo.waymo_sourceloader.WaymoLiDARSource
+ load_lidar: True # whether to load lidar
+ only_use_top_lidar: False # whether to only use the top lidar
+ truncated_max_range: 80 # max range for truncated lidar in a ego-centric coordinate system
+ truncated_min_range: -2 # min range for truncated lidar in a ego-centric coordinate system.
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/datasets/waymo/5cams.yaml b/worldbench/third_party/drivestudio/configs/datasets/waymo/5cams.yaml
new file mode 100644
index 0000000..11aae38
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/datasets/waymo/5cams.yaml
@@ -0,0 +1,43 @@
+# Waymo dataset configuration
+# Processed cameras:
+# idx camera original size
+# 0 front_camera (1920, 1280)
+# 1 front_left_camera (1920, 1280)
+# 2 front_right_camera (1920, 1280)
+# 3 left_camera (1920, 866)
+# 4 right_camera (1920, 866)
+
+data:
+ data_root: data/waymo/processed/training # data root for the dataset
+ dataset: waymo
+ scene_idx: 0 # which scene to use, [0, 798] for waymo's training set and [0, 849] for nuscenes's train/val sets, inclusive
+ start_timestep: 0 # which timestep to start from
+ end_timestep: -1 # which timestep to end at, -1 means the last timestep
+ preload_device: cpu # choose from ["cpu", "cuda"], cache the data on this device.
+ pixel_source: # image source and object annotations
+ type: datasets.waymo.waymo_sourceloader.WaymoPixelSource
+ cameras: [0, 1, 2, 3, 4] # which cameras to use
+ downscale_when_loading: [2, 2, 2, 2, 2] # the size of the images to load
+ downscale: 1 # downscale factor wrt to the downscale_when_loading
+ undistort: True # whether to undistort the images
+ test_image_stride: 0 # use every Nth timestep for the test set. if 0, use all images for training and none for testing
+ load_sky_mask: True # whether to load sky mask
+ load_dynamic_mask: True # whether to load dynamic mask
+ load_objects: True # whether to load object bounding boxes
+ load_smpl: True # whether to load SMPL template for pedestrians
+ sampler: # error based image sampler
+ buffer_downscale: 8 # downscale factor for the buffer wrt load_size
+ buffer_ratio: 0.5 # the percentage of images sampled according to the error buffer
+ start_enhance_weight: 3 # give more chance to sample starting frames, which usually have more errors
+ lidar_source: # everything related to "lidar" --- from lidar points
+ type: datasets.waymo.waymo_sourceloader.WaymoLiDARSource
+ load_lidar: True # whether to load lidar
+ only_use_top_lidar: False # whether to only use the top lidar
+ truncated_max_range: 80 # max range for truncated lidar in a ego-centric coordinate system
+ truncated_min_range: -2 # min range for truncated lidar in a ego-centric coordinate system.
+ # ---- compute aabb from lidar ---- #
+ # if load_lidar is True, we compute aabb from lidar, otherwise we compute aabb from cameras
+ # 1) downsample lidar by random sampling to 1/lidar_downsample_factor number of points
+ # 2) compute aabb from the downsampled lidar points by using the percentile of lidar_percentiles
+ lidar_downsample_factor: 4 # downsample lidar by this factor to compute percentile
+ lidar_percentile: 0.02 # percentile to compute aabb from lidar
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/deformablegs.yaml b/worldbench/third_party/drivestudio/configs/deformablegs.yaml
new file mode 100644
index 0000000..1e7fff3
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/deformablegs.yaml
@@ -0,0 +1,142 @@
+seed: 0
+dataset: waymo/3cams
+
+# ------------- Trainer ------------ #
+trainer:
+ type: models.trainers.SingleTrainer
+ optim:
+ num_iters: 30000
+ use_grad_scaler: false
+ cache_buffer_freq: -1
+ render:
+ near_plane: 0.1 # near plane for rendering
+ far_plane: 10000000000.0 # far plane for rendering
+ antialiased: false # whether to use antialiasing for gaussian rendering, supported by gsplat kernel
+ packed: false # whether to use packed rendering, supported by gsplat kernel
+ absgrad: true # whether to use absolute gradient for rendering, supported by gsplat kernel
+ sparse_grad: false # whether to use sparse gradient for rendering, supported by gsplat kernel
+ batch_size: 1 # batch size for rendering, currently only support 1
+ losses:
+ rgb:
+ w: 0.8
+ ssim:
+ w: 0.2
+ mask:
+ w: 0.05
+ opacity_loss_type: bce # choose from [bce, safe_bce]
+ depth:
+ w: 0.01 # weight of depth loss
+ inverse_depth: False # whether to use inverse depth, NOTE that when set to True, must normalize=True
+ normalize: False # whether to normalize depth loss
+ loss_type: l1 # choose from ["l1", "l2"]
+ affine:
+ w: 0.00001 # weight of affine regularization
+ res_schedule:
+ double_steps: 250 # training starts at 1/d resolution, every n steps this is doubled
+ downscale_times: 2 # at the beginning, resolution is 1/2^d, where d is this number
+ gaussian_optim_general_cfg:
+ xyz:
+ lr: 1.6e-04
+ lr_final: 1.6e-06
+ scale_factor: scene_radius # str or float, if "scene_scale", scale the learning rate by the scene scale
+ sh_dc:
+ lr: 0.0025
+ sh_rest:
+ lr: 0.000125
+ opacity:
+ lr: 0.05
+ scaling:
+ lr: 0.005
+ rotation:
+ lr: 0.001
+ gaussian_ctrl_general_cfg:
+ warmup_steps: 500 # warmup steps for alpha
+ reset_alpha_interval: 3000 # reset alpha every n steps
+ refine_interval: 100 # refine gaussians every n steps
+ sh_degree_interval: 1000 # every n intervals turn on another sh degree
+ n_split_samples: 2 # number of samples to split gaussians into
+ # may differ in different models
+ reset_alpha_value: 0.01 # reset alpha to this value
+ densify_grad_thresh: 0.001 # above this grad, gaussians are densified
+ densify_size_thresh: 0.003 # below this size, gaussians are *duplicated*, otherwise split
+ cull_alpha_thresh: 0.005 # threshold of opacity for culling gaussians
+ cull_scale_thresh: 0.5 # threshold of scale for culling gaussians
+ cull_screen_size: 0.15 # if a gaussian is more than this percent of screen space, cull it
+ split_screen_size: 0.05 # if a gaussian is more than this percent of screen space, split it
+ stop_screen_size_at: 4000 # stop culling/splitting at this step WRT screen size of gaussians
+ stop_split_at: 15000 # stop splitting at this step
+ sh_degree: 3 # sh degree for gaussians
+
+# ------------- Model ------------ #
+model:
+ Background:
+ type: models.gaussians.DeformableGaussians
+ init:
+ from_lidar:
+ num_samples: 800_000
+ return_color: True
+ near_randoms: 100_000
+ far_randoms: 100_000
+ networks:
+ D: 8
+ W: 256
+ x_multires: 10 # default 10
+ t_multires: 10 # default 10
+ reg:
+ sharp_shape_reg:
+ w: 1.
+ step_interval: 10
+ max_gauss_ratio: 10. # threshold of ratio of gaussian max to min scale before applying regularization loss from the PhysGaussian paper
+ ctrl:
+ coarse_train_interval: 3000
+ delta_xyz_rescale: False
+ optim:
+ deform_network:
+ lr: 1.6e-04
+ lr_final: 1.6e-06
+ lr_pre_warmup: 0.01
+ scale_factor: 5 # str or float, if "scene_scale", scale the learning rate by the scene scale
+ max_steps: 40000
+ Sky:
+ type: models.modules.EnvLight
+ params:
+ resolution: 1024
+ optim:
+ all:
+ lr: 0.01
+ Affine:
+ type: models.modules.AffineTransform
+ params:
+ embedding_dim: 4
+ base_mlp_layer_width: 64
+ pixel_affine: False
+ optim:
+ all:
+ lr: 1.0e-5
+ weight_decay: 1.0e-6
+ CamPose:
+ type: models.modules.CameraOptModule
+ optim:
+ all:
+ lr: 1.0e-5
+ weight_decay: 1.0e-6
+
+# ------------- render ------------ #
+render:
+ fps: 10 # frames per second for the main rendered output
+ render_full: True # whether to render full resolution videos
+ render_test: True # whether to render test set
+ render_novel:
+ traj_types:
+ - front_center_interp # type of trajectory for novel view synthesis
+ fps: 24 # frames per second for novel view rendering
+ vis_lidar: False # whether to visualize lidar points on ground truth images
+ vis_sky: False # whether to include "rgb_sky" and "rgb_sky_blend" in rendered keys
+ vis_error: False # whether to include "rgb_error_map" in rendered keys
+
+# ------------- logging ------------ #
+logging:
+ vis_freq: 2000 # how often to visualize training stats
+ print_freq: 500 # how often to print training stats
+ saveckpt_freq: 15000 # how often to save checkpoints
+ save_seperate_video: True # whether to save seperate videos for each scene
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/omnire.yaml b/worldbench/third_party/drivestudio/configs/omnire.yaml
new file mode 100644
index 0000000..0b59335
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/omnire.yaml
@@ -0,0 +1,278 @@
+seed: 0
+dataset: waymo/3cams
+
+# ------------- Trainer ------------ #
+trainer:
+ type: models.trainers.MultiTrainer
+ optim:
+ num_iters: 30000
+ use_grad_scaler: false
+ cache_buffer_freq: -1 # if > 0, use error based image sampler for training
+ render:
+ near_plane: 0.1 # near plane for rendering
+ far_plane: 10000000000.0 # far plane for rendering
+ antialiased: false # whether to use antialiasing for gaussian rendering, supported by gsplat kernel
+ packed: false # whether to use packed rendering, supported by gsplat kernel
+ absgrad: true # whether to use absolute gradient for rendering, supported by gsplat kernel
+ sparse_grad: false # whether to use sparse gradient for rendering, supported by gsplat kernel
+ batch_size: 1 # batch size for rendering, currently only support 1
+ losses:
+ rgb:
+ w: 0.8
+ ssim:
+ w: 0.2
+ mask:
+ w: 0.05
+ opacity_loss_type: bce # choose from [bce, safe_bce]
+ depth:
+ w: 0.01 # weight of depth loss
+ inverse_depth: False # whether to use inverse depth, NOTE that when set to True, must normalize=True
+ normalize: False # whether to normalize depth loss
+ loss_type: l1 # choose from ["l1", "l2"]
+ affine:
+ w: 0.00001 # weight of affine regularization
+ res_schedule:
+ double_steps: 250 # training starts at 1/d resolution, every n steps this is doubled
+ downscale_times: 2 # at the beginning, resolution is 1/2^d, where d is this number
+ gaussian_optim_general_cfg:
+ xyz:
+ lr: 1.6e-04
+ lr_final: 1.6e-06
+ scale_factor: scene_radius # str or float, if "scene_scale", scale the learning rate by the scene scale
+ sh_dc:
+ lr: 0.0025
+ sh_rest:
+ lr: 0.000125
+ opacity:
+ lr: 0.05
+ scaling:
+ lr: 0.005
+ rotation:
+ lr: 0.001
+ gaussian_ctrl_general_cfg:
+ warmup_steps: 500 # warmup steps for alpha
+ reset_alpha_interval: 3000 # reset alpha every n steps
+ refine_interval: 100 # refine gaussians every n steps
+ sh_degree_interval: 1000 # every n intervals turn on another sh degree
+ n_split_samples: 2 # number of samples to split gaussians into
+ # may differ in different models
+ reset_alpha_value: 0.01 # reset alpha to this value
+ densify_grad_thresh: 0.0005 # above this grad, gaussians are densified
+ densify_size_thresh: 0.003 # below this size, gaussians are *duplicated*, otherwise split
+ cull_alpha_thresh: 0.005 # threshold of opacity for culling gaussians
+ cull_scale_thresh: 0.5 # threshold of scale for culling gaussians
+ cull_screen_size: 0.15 # if a gaussian is more than this percent of screen space, cull it
+ split_screen_size: 0.05 # if a gaussian is more than this percent of screen space, split it
+ stop_screen_size_at: 4000 # stop culling/splitting at this step WRT screen size of gaussians
+ stop_split_at: 15000 # stop splitting at this step
+ sh_degree: 3 # sh degree for gaussians
+
+# ------------- Model ------------ #
+model:
+ Background:
+ type: models.gaussians.VanillaGaussians
+ init:
+ from_lidar:
+ num_samples: 800_000
+ return_color: True
+ near_randoms: 100_000
+ far_randoms: 100_000
+ reg:
+ sharp_shape_reg:
+ w: 1.
+ step_interval: 10
+ max_gauss_ratio: 10. # threshold of ratio of gaussian max to min scale before applying regularization loss from the PhysGaussian paper
+ RigidNodes:
+ type: models.nodes.RigidNodes
+ init:
+ instance_max_pts: 5000 # max initial points for each instance
+ only_moving: true # only optimize moving instances
+ traj_length_thres: 1.0 # threshold of trajectory length for moving instances
+ ctrl:
+ cull_scale_thresh: 0.1
+ stop_screen_size_at: 30000
+ stop_split_at: 20000
+ cull_out_of_bound: true
+ reg:
+ sharp_shape_reg:
+ w: 1.
+ step_interval: 10
+ max_gauss_ratio: 10.
+ temporal_smooth_reg:
+ trans:
+ w: 0.01 # optimal value may vary
+ smooth_range: 5 # no ablation
+ optim:
+ ins_rotation:
+ lr: 0.00001
+ lr_final: 0.000005
+ ins_translation:
+ lr: 0.0005
+ lr_final: 0.0001
+ DeformableNodes:
+ type: models.nodes.DeformableNodes
+ init:
+ instance_max_pts: 5000
+ only_moving: true
+ traj_length_thres: 0.5
+ networks:
+ D: 8
+ W: 256
+ embed_dim: 16
+ x_multires: 10 # default 10
+ t_multires: 10 # default 10
+ deform_quat: True
+ deform_scale: False
+ ctrl:
+ reset_alpha_value: 0.1
+ cull_scale_thresh: 0.1
+ stop_screen_size_at: 30000
+ stop_split_at: 20000
+ sh_degree: 3
+ cull_out_of_bound: false
+ use_deformgs_for_nonrigid: True # default True
+ use_deformgs_after: 3000 # default 3000
+ stop_optimizing_canonical_xyz: True
+ reg:
+ sharp_shape_reg:
+ w: 1.
+ step_interval: 10
+ max_gauss_ratio: 10.
+ out_of_bound_loss:
+ w: 1.
+ stop_after: 30000
+ temporal_smooth_reg:
+ trans:
+ w: 0.01 # optimal value may vary
+ smooth_range: 5 # optimal value may vary
+ optim:
+ xyz:
+ lr: 1.6e-04
+ lr_final: 1.6e-06
+ scale_factor: 2.
+ embedding:
+ lr: 0.001
+ lr_final: 0.0001
+ deform_network:
+ lr: 1.6e-03
+ lr_final: 1.6e-04
+ scale_factor: 5.
+ ins_rotation:
+ lr: 0.00001
+ lr_final: 0.000005
+ ins_translation:
+ lr: 0.0005
+ lr_final: 0.0001
+ SMPLNodes:
+ type: models.nodes.SMPLNodes
+ init:
+ only_moving: true
+ traj_length_thres: 0.5
+ reg:
+ sharp_shape_reg:
+ w: 1.
+ step_interval: 10
+ max_gauss_ratio: 10.
+ temporal_smooth_reg:
+ trans:
+ w: 0.01 # optimal value may vary
+ smooth_range: 5 # range of frames to consider for smoothing
+ # joint_smooth:
+ # w: 0.01 # optimal value may vary
+ knn_reg:
+ # lambda_std_x: 0.0001
+ lambda_std_q: 0.001
+ lambda_std_s: 0.001
+ lambda_std_o: 0.001
+ lambda_std_shs_dc: 0.001
+ lambda_std_shs_rest: 0.001
+ max_s_square_reg:
+ w: 0.05
+ x_offset:
+ w: 0.2
+ voxel_deformer_reg:
+ lambda_std_w: 0.6
+ lambda_std_w_rest: 0.5
+ lambda_w_norm: 0.6
+ lambda_w_rest_norm: 0.3
+ optim:
+ xyz:
+ lr: 1.6e-04
+ lr_final: 1.6e-06
+ scale_factor: 3.0
+ rotation:
+ lr: 0.005
+ sh_dc:
+ lr: 0.005
+ sh_rest:
+ lr: 0.00025
+ ins_rotation:
+ lr: 0.00001
+ lr_final: 0.000005
+ ins_translation:
+ lr: 0.0005
+ lr_final: 0.0001
+ smpl_rotation:
+ lr: 0.00005
+ lr_final: 0.00001
+ w_dc_vox:
+ lr: 0.0001
+ lr_final: 0.00001
+ opt_after: 10000
+ ctrl:
+ sh_degree: 1
+ opacity_init_value: 0.99
+ ball_gaussians: false
+ constrain_xyz_offset: false
+ knn_update_interval: 100
+ knn_neighbors: 3
+ use_voxel_deformer: true
+ freeze_x: false
+ freeze_o: false
+ freeze_q: false
+ freeze_s: false
+ freeze_shs_dc: false
+ freeze_shs_rest: false
+ Sky:
+ type: models.modules.EnvLight
+ params:
+ resolution: 1024
+ optim:
+ all:
+ lr: 0.01
+ Affine:
+ type: models.modules.AffineTransform
+ params:
+ embedding_dim: 4
+ base_mlp_layer_width: 64
+ pixel_affine: False
+ optim:
+ all:
+ lr: 1.0e-5
+ weight_decay: 1.0e-6
+ CamPose:
+ type: models.modules.CameraOptModule
+ optim:
+ all:
+ lr: 1.0e-5
+ weight_decay: 1.0e-6
+
+# ------------- render ------------ #
+render:
+ fps: 10 # frames per second for the main rendered output
+ render_full: True # whether to render full resolution videos
+ render_test: True # whether to render test set
+ render_novel:
+ traj_types:
+ - front_center_interp # type of trajectory for novel view synthesis
+ fps: 24 # frames per second for novel view rendering
+ vis_lidar: False # whether to visualize lidar points on ground truth images
+ vis_sky: False # whether to include "rgb_sky" and "rgb_sky_blend" in rendered keys
+ vis_error: False # whether to include "rgb_error_map" in rendered keys
+
+# ------------- logging ------------ #
+logging:
+ vis_freq: 2000 # how often to visualize training stats
+ print_freq: 500 # how often to print training stats
+ saveckpt_freq: 15000 # how often to save checkpoints
+ save_seperate_video: True # whether to save seperate videos for each scene
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/omnire_extended_cam.yaml b/worldbench/third_party/drivestudio/configs/omnire_extended_cam.yaml
new file mode 100644
index 0000000..e9c31c2
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/omnire_extended_cam.yaml
@@ -0,0 +1,279 @@
+seed: 0
+dataset: waymo/3cams
+
+# ------------- Trainer ------------ #
+trainer:
+ type: models.trainers.MultiTrainer
+ optim:
+ num_iters: 40000
+ use_grad_scaler: false
+ cache_buffer_freq: 2000 # we use error based image sampler for training with more than 3 cameras
+ render:
+ near_plane: 0.1 # near plane for rendering
+ far_plane: 10000000000.0 # far plane for rendering
+ antialiased: false # whether to use antialiasing for gaussian rendering, supported by gsplat kernel
+ packed: false # whether to use packed rendering, supported by gsplat kernel
+ absgrad: true # whether to use absolute gradient for rendering, supported by gsplat kernel
+ sparse_grad: false # whether to use sparse gradient for rendering, supported by gsplat kernel
+ batch_size: 1 # batch size for rendering, currently only support 1
+ losses:
+ rgb:
+ w: 0.8
+ ssim:
+ w: 0.2
+ mask:
+ w: 0.05
+ opacity_loss_type: bce # choose from [bce, safe_bce]
+ depth:
+ w: 0.01 # weight of depth loss
+ inverse_depth: False # whether to use inverse depth, NOTE that when set to True, must normalize=True
+ normalize: False # whether to normalize depth loss
+ loss_type: l1 # choose from ["l1", "l2"]
+ affine:
+ w: 0.00001 # weight of affine regularization
+ res_schedule:
+ double_steps: 250 # training starts at 1/d resolution, every n steps this is doubled
+ downscale_times: 2 # at the beginning, resolution is 1/2^d, where d is this number
+ gaussian_optim_general_cfg:
+ xyz:
+ lr: 1.6e-04
+ lr_final: 1.6e-06
+ scale_factor: scene_radius # str or float, if "scene_scale", scale the learning rate by the scene scale
+ sh_dc:
+ lr: 0.0025
+ sh_rest:
+ lr: 0.000125
+ opacity:
+ lr: 0.05
+ scaling:
+ lr: 0.005
+ rotation:
+ lr: 0.001
+ gaussian_ctrl_general_cfg:
+ warmup_steps: 500 # warmup steps for alpha
+ reset_alpha_interval: 4000 # reset alpha every n steps
+ refine_interval: 100 # refine gaussians every n steps
+ sh_degree_interval: 1000 # every n intervals turn on another sh degree
+ n_split_samples: 2 # number of samples to split gaussians into
+ # may differ in different models
+ reset_alpha_value: 0.01 # reset alpha to this value
+ densify_grad_thresh: 0.0003 # above this grad, gaussians are densified
+ densify_size_thresh: 0.002 # below this size, gaussians are *duplicated*, otherwise split
+ cull_alpha_thresh: 0.005 # threshold of opacity for culling gaussians
+ cull_scale_thresh: 0.5 # threshold of scale for culling gaussians
+ cull_screen_size: 0.15 # if a gaussian is more than this percent of screen space, cull it
+ split_screen_size: 0.05 # if a gaussian is more than this percent of screen space, split it
+ stop_screen_size_at: 5000 # stop culling/splitting at this step WRT screen size of gaussians
+ stop_split_at: 20000 # stop splitting at this step
+ sh_degree: 3 # sh degree for gaussians
+
+# ------------- Model ------------ #
+model:
+ Background:
+ type: models.gaussians.VanillaGaussians
+ init:
+ from_lidar:
+ num_samples: 800_000
+ return_color: True
+ near_randoms: 100_000
+ far_randoms: 100_000
+ reg:
+ sharp_shape_reg:
+ w: 1.
+ step_interval: 10
+ max_gauss_ratio: 10. # threshold of ratio of gaussian max to min scale before applying regularization loss from the PhysGaussian paper
+ RigidNodes:
+ type: models.nodes.RigidNodes
+ init:
+ instance_max_pts: 5000 # max initial points for each instance
+ only_moving: true # only optimize moving instances
+ traj_length_thres: 1.0 # threshold of trajectory length for moving instances
+ ctrl:
+ cull_scale_thresh: 0.1
+ stop_screen_size_at: 40000
+ stop_split_at: 30000
+ cull_out_of_bound: true
+ densify_grad_thresh: 0.0002
+ reg:
+ sharp_shape_reg:
+ w: 1.
+ step_interval: 10
+ max_gauss_ratio: 10.
+ temporal_smooth_reg:
+ trans:
+ w: 0.01 # optimal value may vary
+ smooth_range: 5 # no ablation
+ optim:
+ ins_rotation:
+ lr: 0.00001
+ lr_final: 0.000005
+ ins_translation:
+ lr: 0.0005
+ lr_final: 0.0001
+ DeformableNodes:
+ type: models.nodes.DeformableNodes
+ init:
+ instance_max_pts: 5000
+ only_moving: true
+ traj_length_thres: 0.5
+ networks:
+ D: 8
+ W: 256
+ embed_dim: 16
+ x_multires: 10 # default 10
+ t_multires: 10 # default 10
+ deform_quat: True
+ deform_scale: False
+ ctrl:
+ reset_alpha_value: 0.1
+ cull_scale_thresh: 0.1
+ stop_screen_size_at: 40000
+ stop_split_at: 30000
+ sh_degree: 3
+ cull_out_of_bound: false
+ use_deformgs_for_nonrigid: True # default True
+ use_deformgs_after: 4000 # default 3000
+ stop_optimizing_canonical_xyz: True
+ reg:
+ sharp_shape_reg:
+ w: 1.
+ step_interval: 10
+ max_gauss_ratio: 10.
+ out_of_bound_loss:
+ w: 1.
+ stop_after: 40000
+ temporal_smooth_reg:
+ trans:
+ w: 0.01 # optimal value may vary
+ smooth_range: 5 # optimal value may vary
+ optim:
+ xyz:
+ lr: 1.6e-04
+ lr_final: 1.6e-06
+ scale_factor: 2.
+ embedding:
+ lr: 0.001
+ lr_final: 0.0001
+ deform_network:
+ lr: 1.6e-03
+ lr_final: 1.6e-04
+ scale_factor: 5.
+ ins_rotation:
+ lr: 0.00001
+ lr_final: 0.000005
+ ins_translation:
+ lr: 0.0005
+ lr_final: 0.0001
+ SMPLNodes:
+ type: models.nodes.SMPLNodes
+ init:
+ only_moving: true
+ traj_length_thres: 0.5
+ reg:
+ sharp_shape_reg:
+ w: 1.
+ step_interval: 10
+ max_gauss_ratio: 10.
+ temporal_smooth_reg:
+ trans:
+ w: 0.01 # optimal value may vary
+ smooth_range: 5 # range of frames to consider for smoothing
+ # joint_smooth:
+ # w: 0.01 # optimal value may vary
+ knn_reg:
+ # lambda_std_x: 0.0001
+ lambda_std_q: 0.001
+ lambda_std_s: 0.001
+ lambda_std_o: 0.001
+ lambda_std_shs_dc: 0.001
+ lambda_std_shs_rest: 0.001
+ max_s_square_reg:
+ w: 0.05
+ x_offset:
+ w: 0.2
+ voxel_deformer_reg:
+ lambda_std_w: 0.6
+ lambda_std_w_rest: 0.5
+ lambda_w_norm: 0.6
+ lambda_w_rest_norm: 0.3
+ optim:
+ xyz:
+ lr: 1.6e-04
+ lr_final: 1.6e-06
+ scale_factor: 3.0
+ rotation:
+ lr: 0.005
+ sh_dc:
+ lr: 0.005
+ sh_rest:
+ lr: 0.00025
+ ins_rotation:
+ lr: 0.00001
+ lr_final: 0.000005
+ ins_translation:
+ lr: 0.0005
+ lr_final: 0.0001
+ smpl_rotation:
+ lr: 0.00005
+ lr_final: 0.00001
+ w_dc_vox:
+ lr: 0.0001
+ lr_final: 0.00001
+ opt_after: 10000
+ ctrl:
+ sh_degree: 1
+ opacity_init_value: 0.99
+ ball_gaussians: false
+ constrain_xyz_offset: false
+ knn_update_interval: 100
+ knn_neighbors: 3
+ use_voxel_deformer: true
+ freeze_x: false
+ freeze_o: false
+ freeze_q: false
+ freeze_s: false
+ freeze_shs_dc: false
+ freeze_shs_rest: false
+ Sky:
+ type: models.modules.EnvLight
+ params:
+ resolution: 1024
+ optim:
+ all:
+ lr: 0.01
+ Affine:
+ type: models.modules.AffineTransform
+ params:
+ embedding_dim: 4
+ base_mlp_layer_width: 64
+ pixel_affine: False
+ optim:
+ all:
+ lr: 1.0e-5
+ weight_decay: 1.0e-6
+ CamPose:
+ type: models.modules.CameraOptModule
+ optim:
+ all:
+ lr: 1.0e-5
+ weight_decay: 1.0e-6
+
+# ------------- render ------------ #
+render:
+ fps: 10 # frames per second for the main rendered output
+ render_full: True # whether to render full resolution videos
+ render_test: True # whether to render test set
+ render_novel:
+ traj_types:
+ - front_center_interp # type of trajectory for novel view synthesis
+ fps: 24 # frames per second for novel view rendering
+ vis_lidar: False # whether to visualize lidar points on ground truth images
+ vis_sky: False # whether to include "rgb_sky" and "rgb_sky_blend" in rendered keys
+ vis_error: False # whether to include "rgb_error_map" in rendered keys
+
+# ------------- logging ------------ #
+logging:
+ vis_freq: 2000 # how often to visualize training stats
+ print_freq: 500 # how often to print training stats
+ saveckpt_freq: 20000 # how often to save checkpoints
+ save_seperate_video: True # whether to save seperate videos for each scene
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/paper_legacy/omnire.yaml b/worldbench/third_party/drivestudio/configs/paper_legacy/omnire.yaml
new file mode 100644
index 0000000..84ff298
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/paper_legacy/omnire.yaml
@@ -0,0 +1,268 @@
+seed: 0
+dataset: waymo/3cams
+
+# ------------- Trainer ------------ #
+trainer:
+ type: models.trainers.MultiTrainer
+ optim:
+ num_iters: 30000
+ use_grad_scaler: false
+ cache_buffer_freq: -1
+ render:
+ near_plane: 0.1
+ far_plane: 10000000000.0
+ antialiased: false
+ packed: false
+ absgrad: true
+ sparse_grad: false
+ batch_size: 1
+ losses:
+ rgb:
+ w: 0.8
+ ssim:
+ w: 0.2
+ mask:
+ w: 0.05
+ opacity_loss_type: bce # choose from [bce, safe_bce]
+ depth:
+ w: 0.1 # weight of depth loss
+ inverse_depth: True # whether to use inverse depth, NOTE that when set to True, must normalize=True
+ normalize: False # whether to normalize depth loss
+ loss_type: l1 # choose from ["l1", "l2"]
+ opacity_entropy:
+ w: 0.05
+ inverse_depth_smoothness:
+ w: 0.001
+ dynamic_region:
+ factor: 5
+ start_from: 20000
+ res_schedule:
+ double_steps: 250 # training starts at 1/d resolution, every n steps this is doubled
+ downscale_times: 2 # at the beginning, resolution is 1/2^d, where d is this number
+ gaussian_optim_general_cfg:
+ xyz:
+ lr: 1.6e-04
+ lr_final: 1.6e-06
+ scale_factor: scene_radius # str or float, if "scene_scale", scale the learning rate by the scene scale
+ sh_dc:
+ lr: 0.0025
+ sh_rest:
+ lr: 0.000125
+ opacity:
+ lr: 0.05
+ scaling:
+ lr: 0.005
+ rotation:
+ lr: 0.001
+ gaussian_ctrl_general_cfg:
+ warmup_steps: 500 # warmup steps for alpha
+ reset_alpha_interval: 3000 # reset alpha every n steps
+ refine_interval: 100 # refine gaussians every n steps
+ sh_degree_interval: 1000 # every n intervals turn on another sh degree
+ n_split_samples: 2 # number of samples to split gaussians into
+ # may differ in different models
+ reset_alpha_value: 0.01 # reset alpha to this value
+ densify_grad_thresh: 0.0003 # above this grad, gaussians are densified
+ densify_size_thresh: 0.003 # below this size, gaussians are *duplicated*, otherwise split
+ cull_alpha_thresh: 0.005 # threshold of opacity for culling gaussians
+ cull_scale_thresh: 0.5 # threshold of scale for culling gaussians
+ cull_screen_size: 0.15 # if a gaussian is more than this percent of screen space, cull it
+ split_screen_size: 0.05 # if a gaussian is more than this percent of screen space, split it
+ stop_screen_size_at: 4000 # stop culling/splitting at this step WRT screen size of gaussians
+ stop_split_at: 15000 # stop splitting at this step
+ sh_degree: 3 # sh degree for gaussians
+
+# ------------- Model ------------ #
+model:
+ Background:
+ type: models.gaussians.VanillaGaussians
+ init:
+ from_lidar:
+ num_samples: 600_000
+ return_color: True
+ near_randoms: 200_000
+ far_randoms: 200_000
+ reg:
+ sharp_shape_reg:
+ w: 1.
+ step_interval: 10
+ max_gauss_ratio: 10. # threshold of ratio of gaussian max to min scale before applying regularization loss from the PhysGaussian paper
+ RigidNodes:
+ type: models.nodes.RigidNodes
+ init:
+ instance_max_pts: 5000
+ only_moving: true
+ traj_length_thres: 1.0
+ ctrl:
+ cull_scale_thresh: 0.1
+ stop_screen_size_at: 30000
+ stop_split_at: 20000
+ cull_out_of_bound: true
+ reg:
+ sharp_shape_reg:
+ w: 1.
+ step_interval: 10
+ max_gauss_ratio: 10.
+ temporal_smooth_reg:
+ trans:
+ w: 0.01 # no ablation, maybe some other values will be better
+ smooth_range: 5 # no ablation
+ optim:
+ ins_rotation:
+ lr: 0.00001
+ lr_final: 0.000005
+ ins_translation:
+ lr: 0.0005
+ lr_final: 0.0001
+ DeformableNodes:
+ type: models.nodes.DeformableNodes
+ init:
+ instance_max_pts: 5000
+ only_moving: true
+ traj_length_thres: 0.5
+ networks:
+ D: 8
+ W: 256
+ embed_dim: 16
+ x_multires: 10 # default 10
+ t_multires: 10 # default 10
+ deform_quat: True
+ deform_scale: False
+ ctrl:
+ reset_alpha_value: 0.1
+ cull_scale_thresh: 0.1
+ stop_screen_size_at: 30000
+ stop_split_at: 20000
+ sh_degree: 3
+ cull_out_of_bound: false
+ use_deformgs_for_nonrigid: True # default True
+ use_deformgs_after: 3000 # default 3000
+ stop_optimizing_canonical_xyz: True
+ reg:
+ sharp_shape_reg:
+ w: 1.
+ step_interval: 10
+ max_gauss_ratio: 10.
+ out_of_bound_loss:
+ w: 1.
+ stop_after: 30000
+ temporal_smooth_reg:
+ trans:
+ w: 0.01 # no ablation, maybe some other values will be better
+ smooth_range: 5 # no ablation
+ optim:
+ xyz:
+ lr: 1.6e-04
+ lr_final: 1.6e-06
+ scale_factor: 2.
+ embedding:
+ lr: 0.001
+ lr_final: 0.0001
+ deform_network:
+ lr: 1.6e-03
+ lr_final: 1.6e-04
+ scale_factor: 5.
+ ins_rotation:
+ lr: 0.00001
+ lr_final: 0.000005
+ ins_translation:
+ lr: 0.0005
+ lr_final: 0.0001
+ SMPLNodes:
+ type: models.nodes.SMPLNodes
+ init:
+ only_moving: true
+ traj_length_thres: 0.5
+ reg:
+ sharp_shape_reg:
+ w: 1.
+ step_interval: 10
+ max_gauss_ratio: 10.
+ scaling_reg:
+ w: 1.
+ precentile: 0.99
+ stop_after: 30000
+ start_after: 6000
+ temporal_smooth_reg:
+ trans:
+ w: 0.01 # no ablation, maybe some other values will be better
+ smooth_range: 5 # no ablation
+ # joint_smooth:
+ # w: 0.01 # no ablation, maybe some other values will be better
+ knn_reg:
+ # lambda_std_x: 0.0001
+ lambda_std_q: 0.001
+ lambda_std_s: 0.001
+ lambda_std_o: 0.001
+ lambda_std_shs_dc: 0.001
+ lambda_std_shs_rest: 0.001
+ max_s_square_reg:
+ w: 0.05
+ x_offset:
+ w: 0.2
+ voxel_deformer_reg:
+ lambda_std_w: 0.6 # default 0.3
+ lambda_std_w_rest: 0.5
+ lambda_w_norm: 0.6 # default 0.3
+ lambda_w_rest_norm: 0.3
+ optim:
+ xyz:
+ lr: 1.6e-04
+ lr_final: 1.6e-06
+ scale_factor: 3.0
+ rotation:
+ lr: 0.005
+ sh_dc:
+ lr: 0.005
+ sh_rest:
+ lr: 0.00025
+ ins_rotation:
+ lr: 0.00001
+ lr_final: 0.000005
+ ins_translation:
+ lr: 0.0005
+ lr_final: 0.0001
+ smpl_rotation:
+ lr: 0.00005
+ lr_final: 0.00001
+ w_dc_vox:
+ lr: 0.0001 # default 0.0001
+ lr_final: 0.00001 # default 0.00001
+ opt_after: 10000
+ ctrl:
+ sh_degree: 1
+ opacity_init_value: 0.99
+ ball_gaussians: false
+ constrain_xyz_offset: false
+ knn_update_interval: 100
+ knn_neighbors: 3
+ use_voxel_deformer: true
+ freeze_x: false
+ freeze_o: false
+ freeze_q: false
+ freeze_s: false
+ freeze_shs_dc: false
+ freeze_shs_rest: false
+ Sky:
+ type: models.modules.EnvLight
+ params:
+ resolution: 1024
+ optim:
+ all:
+ lr: 0.01
+
+# ------------- render ------------ #
+render:
+ fps: 16 # frames per second for the main rendered output
+ render_full: True # whether to render full resolution videos
+ render_test: True # whether to render test set
+ vis_lidar: False # whether to visualize lidar points on ground truth images
+ vis_sky: False # whether to include "rgb_sky" and "rgb_sky_blend" in rendered keys
+ vis_error: False # whether to include "rgb_error_map" in rendered keys
+
+# ------------- logging ------------ #
+logging:
+ vis_freq: 2000 # how often to visualize training stats
+ print_freq: 500 # how often to print training stats
+ saveckpt_freq: 15000 # how often to save checkpoints
+ save_seperate_video: True # whether to save seperate videos for each scene
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/paper_legacy/streetgs.yaml b/worldbench/third_party/drivestudio/configs/paper_legacy/streetgs.yaml
new file mode 100644
index 0000000..84dcef8
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/paper_legacy/streetgs.yaml
@@ -0,0 +1,139 @@
+seed: 0
+dataset: waymo/3cams
+
+# ------------- Trainer ------------ #
+trainer:
+ type: models.trainers.MultiTrainer
+ optim:
+ num_iters: 30000
+ use_grad_scaler: false
+ cache_buffer_freq: -1
+ render:
+ near_plane: 0.1
+ far_plane: 10000000000.0
+ antialiased: false
+ packed: false
+ absgrad: true
+ sparse_grad: false
+ batch_size: 1
+ losses:
+ rgb:
+ w: 0.8
+ ssim:
+ w: 0.2
+ mask:
+ w: 0.05
+ opacity_loss_type: bce # choose from [bce, safe_bce]
+ depth:
+ w: 0.1 # weight of depth loss
+ inverse_depth: True # whether to use inverse depth, NOTE that when set to True, must normalize=True
+ normalize: False # whether to normalize depth loss
+ loss_type: l1 # choose from ["l1", "l2"]
+ opacity_entropy:
+ w: 0.05
+ inverse_depth_smoothness:
+ w: 0.001
+ dynamic_region:
+ factor: 5
+ start_from: 20000
+ res_schedule:
+ double_steps: 250 # training starts at 1/d resolution, every n steps this is doubled
+ downscale_times: 2 # at the beginning, resolution is 1/2^d, where d is this number
+ gaussian_optim_general_cfg:
+ xyz:
+ lr: 1.6e-04
+ lr_final: 1.6e-06
+ scale_factor: scene_radius # str or float, if "scene_scale", scale the learning rate by the scene scale
+ sh_dc:
+ lr: 0.0025
+ sh_rest:
+ lr: 0.000125
+ opacity:
+ lr: 0.05
+ scaling:
+ lr: 0.005
+ rotation:
+ lr: 0.001
+ gaussian_ctrl_general_cfg:
+ warmup_steps: 500 # warmup steps for alpha
+ reset_alpha_interval: 3000 # reset alpha every n steps
+ refine_interval: 100 # refine gaussians every n steps
+ sh_degree_interval: 1000 # every n intervals turn on another sh degree
+ n_split_samples: 2 # number of samples to split gaussians into
+ # may differ in different models
+ reset_alpha_value: 0.01 # reset alpha to this value
+ densify_grad_thresh: 0.0003 # above this grad, gaussians are densified
+ densify_size_thresh: 0.003 # below this size, gaussians are *duplicated*, otherwise split
+ cull_alpha_thresh: 0.005 # threshold of opacity for culling gaussians
+ cull_scale_thresh: 0.5 # threshold of scale for culling gaussians
+ cull_screen_size: 0.15 # if a gaussian is more than this percent of screen space, cull it
+ split_screen_size: 0.05 # if a gaussian is more than this percent of screen space, split it
+ stop_screen_size_at: 4000 # stop culling/splitting at this step WRT screen size of gaussians
+ stop_split_at: 15000 # stop splitting at this step
+ sh_degree: 3 # sh degree for gaussians
+
+# ------------- Model ------------ #
+model:
+ Background:
+ type: models.gaussians.VanillaGaussians
+ init:
+ from_lidar:
+ num_samples: 600_000
+ return_color: True
+ near_randoms: 200_000
+ far_randoms: 200_000
+ reg:
+ sharp_shape_reg:
+ w: 1.
+ step_interval: 10
+ max_gauss_ratio: 10. # threshold of ratio of gaussian max to min scale before applying regularization loss from the PhysGaussian paper
+ RigidNodes:
+ type: models.nodes.RigidNodes
+ init:
+ instance_max_pts: 5000
+ only_moving: true
+ traj_length_thres: 1.0
+ ctrl:
+ cull_scale_thresh: 0.1
+ stop_screen_size_at: 30000
+ stop_split_at: 20000
+ cull_out_of_bound: true
+ reg:
+ sharp_shape_reg:
+ w: 1.
+ step_interval: 10
+ max_gauss_ratio: 10.
+ temporal_smooth_reg:
+ trans:
+ w: 0.01 # no ablation, maybe some other values will be better
+ smooth_range: 5 # no ablation
+ optim:
+ ins_rotation:
+ lr: 0.00001
+ lr_final: 0.000005
+ ins_translation:
+ lr: 0.0005
+ lr_final: 0.0001
+ Sky:
+ type: models.modules.EnvLight
+ params:
+ resolution: 1024
+ optim:
+ all:
+ lr: 0.01
+
+# ------------- render ------------ #
+render:
+ fps: 16 # frames per second for the main rendered output
+ render_full: True # whether to render full resolution videos
+ render_test: True # whether to render test set
+ vis_lidar: False # whether to visualize lidar points on ground truth images
+ vis_sky: False # whether to include "rgb_sky" and "rgb_sky_blend" in rendered keys
+ vis_error: False # whether to include "rgb_error_map" in rendered keys
+
+# ------------- logging ------------ #
+logging:
+ vis_freq: 2000 # how often to visualize training stats
+ print_freq: 500 # how often to print training stats
+ saveckpt_freq: 15000 # how often to save checkpoints
+ save_seperate_video: True # whether to save seperate videos for each scene
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/pvg.yaml b/worldbench/third_party/drivestudio/configs/pvg.yaml
new file mode 100644
index 0000000..77747a1
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/pvg.yaml
@@ -0,0 +1,156 @@
+seed: 0
+dataset: waymo/3cams
+
+# ------------- Trainer ------------ #
+trainer:
+ type: models.trainers.SingleTrainer
+ optim:
+ num_iters: 30000
+ use_grad_scaler: false
+ cache_buffer_freq: -1 # if > 0, use error based image sampler for training
+ render:
+ near_plane: 0.1 # near plane for rendering
+ far_plane: 10000000000.0 # far plane for rendering
+ antialiased: false # whether to use antialiasing for gaussian rendering, supported by gsplat kernel
+ packed: false # whether to use packed rendering, supported by gsplat kernel
+ absgrad: false # whether to use absolute gradient for rendering, supported by gsplat kernel
+ sparse_grad: false # whether to use sparse gradient for rendering, supported by gsplat kernel
+ batch_size: 1 # batch size for rendering, currently only support 1
+ losses:
+ rgb:
+ w: 0.8
+ ssim:
+ w: 0.2
+ mask:
+ w: 1.0
+ opacity_loss_type: safe_bce # choose from [bce, safe_bce]
+ depth:
+ w: 0.1 # weight of depth loss
+ inverse_depth: True # whether to use inverse depth, NOTE that when set to True, must normalize=True
+ normalize: False # whether to normalize depth loss
+ loss_type: l1 # choose from ["l1", "l2"]
+ reduction: mean_on_hw # choose from ["mean_on_hit", "mean_on_hw", "sum", "none"]
+ # refer to pvg codebase
+ weight_decay: 1.0
+ opacity_entropy:
+ w: 0.05
+ inverse_depth_smoothness:
+ w: 0.001
+ res_schedule:
+ double_steps: 5000 # training starts at 1/d resolution, every n steps this is doubled
+ downscale_times: 3 # at the beginning, resolution is 1/2^d, where d is this number
+ gaussian_optim_general_cfg:
+ xyz:
+ lr: 1.6e-04
+ lr_final: 1.6e-06
+ scale_factor: scene_radius # str or float, if "scene_scale", scale the learning rate by the scene scale
+ sh_dc:
+ lr: 0.0025
+ sh_rest:
+ lr: 0.000125
+ opacity:
+ lr: 0.005
+ scaling:
+ lr: 0.005
+ rotation:
+ lr: 0.001
+ gaussian_ctrl_general_cfg:
+ warmup_steps: 500 # warmup steps for alpha
+ reset_alpha_interval: 3000 # reset alpha every n steps
+ refine_interval: 100 # refine gaussians every n steps
+ sh_degree_interval: 2000 # every n intervals turn on another sh degree
+ n_split_samples: 2 # number of samples to split gaussians into
+ # may differ in different models
+ reset_alpha_value: 0.01 # reset alpha to this value
+ densify_grad_thresh: 0.0002 # above this grad, gaussians are densified
+ densify_size_thresh: 0.01 # below this size, gaussians are *duplicated*, otherwise split
+ cull_alpha_thresh: 0.005 # threshold of opacity for culling gaussians
+ cull_scale_thresh: 0.5 # threshold of scale for culling gaussians
+ cull_screen_size: 0.15 # if a gaussian is more than this percent of screen space, cull it
+ split_screen_size: 0.05 # if a gaussian is more than this percent of screen space, split it
+ stop_screen_size_at: 4000 # stop culling/splitting at this step WRT screen size of gaussians
+ stop_split_at: 20000 # stop splitting at this step
+ sh_degree: 3 # sh degree for gaussians
+
+# ------------- Model ------------ #
+model:
+ Background:
+ type: models.gaussians.PeriodicVibrationGaussians
+ init:
+ from_lidar:
+ num_samples: 800_000
+ return_color: True
+ return_normalized_time: True
+ near_randoms: 100_000
+ far_randoms: 100_000
+ reg:
+ sharp_shape_reg:
+ w: 1.
+ step_interval: 10
+ max_gauss_ratio: 10. # threshold of ratio of gaussian max to min scale before applying regularization loss from the PhysGaussian paper
+ velocity_reg:
+ w: 0.001
+ optim:
+ velocity:
+ lr: 1.0e-3
+ scale_factor: scene_radius
+ life_span:
+ lr: 2.0e-3
+ life_peak:
+ lr: 8.0e-4
+ lr_final: 8.0e-6
+ ctrl:
+ # pvg specific
+ cycle_length: 0.2
+ time_interval: 0.02
+ enable_temporal_smoothing: True
+ smooth_probability: 0.5
+ distribution_span: 1.5 # unit: frame interval
+ betas_init: 0.1
+ densify_until_num_points: 3_000_000
+ densify_t_grad_thresh: 0.002
+ densify_t_size_thresh: 0.01
+ no_time_split: true
+ Sky:
+ type: models.modules.EnvLight
+ params:
+ resolution: 1024
+ optim:
+ all:
+ lr: 0.01
+ Affine:
+ type: models.modules.AffineTransform
+ params:
+ embedding_dim: 4
+ base_mlp_layer_width: 64
+ pixel_affine: False
+ optim:
+ all:
+ lr: 1.0e-5
+ weight_decay: 1.0e-6
+ CamPose:
+ type: models.modules.CameraOptModule
+ optim:
+ all:
+ lr: 1.0e-5
+ weight_decay: 1.0e-6
+
+# ------------- render ------------ #
+render:
+ fps: 10 # frames per second for the main rendered output
+ render_full: True # whether to render full resolution videos
+ render_test: True # whether to render test set
+ render_novel:
+ traj_types:
+ - front_center_interp # type of trajectory for novel view synthesis
+ fps: 24 # frames per second for novel view rendering
+ vis_lidar: False # whether to visualize lidar points on ground truth images
+ vis_sky: False # whether to include "rgb_sky" and "rgb_sky_blend" in rendered keys
+ vis_error: False # whether to include "rgb_error_map" in rendered keys
+
+# ------------- logging ------------ #
+logging:
+ vis_freq: 2000 # how often to visualize training stats
+ print_freq: 500 # how often to print training stats
+ saveckpt_freq: 15000 # how often to save checkpoints
+ save_seperate_video: True # whether to save seperate videos for each scene
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/configs/streetgs.yaml b/worldbench/third_party/drivestudio/configs/streetgs.yaml
new file mode 100644
index 0000000..6ac03b5
--- /dev/null
+++ b/worldbench/third_party/drivestudio/configs/streetgs.yaml
@@ -0,0 +1,154 @@
+seed: 0
+dataset: waymo/3cams
+
+# ------------- Trainer ------------ #
+trainer:
+ type: models.trainers.MultiTrainer
+ optim:
+ num_iters: 30000
+ use_grad_scaler: false
+ cache_buffer_freq: -1 # if > 0, use error based image sampler for training
+ render:
+ near_plane: 0.1 # near plane for rendering
+ far_plane: 10000000000.0 # far plane for rendering
+ antialiased: false # whether to use antialiasing for gaussian rendering, supported by gsplat kernel
+ packed: false # whether to use packed rendering, supported by gsplat kernel
+ absgrad: true # whether to use absolute gradient for rendering, supported by gsplat kernel
+ sparse_grad: false # whether to use sparse gradient for rendering, supported by gsplat kernel
+ batch_size: 1 # batch size for rendering, currently only support 1
+ losses:
+ rgb:
+ w: 0.8
+ ssim:
+ w: 0.2
+ mask:
+ w: 0.05
+ opacity_loss_type: bce # choose from [bce, safe_bce]
+ depth:
+ w: 0.01 # weight of depth loss
+ inverse_depth: False # whether to use inverse depth, NOTE that when set to True, must normalize=True
+ normalize: False # whether to normalize depth loss
+ loss_type: l1 # choose from ["l1", "l2"]
+ affine:
+ w: 0.00001 # weight of affine regularization
+ res_schedule:
+ double_steps: 250 # training starts at 1/d resolution, every n steps this is doubled
+ downscale_times: 2 # at the beginning, resolution is 1/2^d, where d is this number
+ gaussian_optim_general_cfg:
+ xyz:
+ lr: 1.6e-04
+ lr_final: 1.6e-06
+ scale_factor: scene_radius # str or float, if "scene_scale", scale the learning rate by the scene scale
+ sh_dc:
+ lr: 0.0025
+ sh_rest:
+ lr: 0.000125
+ opacity:
+ lr: 0.05
+ scaling:
+ lr: 0.005
+ rotation:
+ lr: 0.001
+ gaussian_ctrl_general_cfg:
+ warmup_steps: 500 # warmup steps for alpha
+ reset_alpha_interval: 3000 # reset alpha every n steps
+ refine_interval: 100 # refine gaussians every n steps
+ sh_degree_interval: 1000 # every n intervals turn on another sh degree
+ n_split_samples: 2 # number of samples to split gaussians into
+ # may differ in different models
+ reset_alpha_value: 0.01 # reset alpha to this value
+ densify_grad_thresh: 0.0005 # above this grad, gaussians are densified
+ densify_size_thresh: 0.003 # below this size, gaussians are *duplicated*, otherwise split
+ cull_alpha_thresh: 0.005 # threshold of opacity for culling gaussians
+ cull_scale_thresh: 0.5 # threshold of scale for culling gaussians
+ cull_screen_size: 0.15 # if a gaussian is more than this percent of screen space, cull it
+ split_screen_size: 0.05 # if a gaussian is more than this percent of screen space, split it
+ stop_screen_size_at: 4000 # stop culling/splitting at this step WRT screen size of gaussians
+ stop_split_at: 15000 # stop splitting at this step
+ sh_degree: 3 # sh degree for gaussians
+
+# ------------- Model ------------ #
+model:
+ Background:
+ type: models.gaussians.VanillaGaussians
+ init:
+ from_lidar:
+ num_samples: 800_000
+ return_color: True
+ near_randoms: 100_000
+ far_randoms: 100_000
+ reg:
+ sharp_shape_reg:
+ w: 1.
+ step_interval: 10
+ max_gauss_ratio: 10. # threshold of ratio of gaussian max to min scale before applying regularization loss from the PhysGaussian paper
+ RigidNodes:
+ type: models.nodes.RigidNodes
+ init:
+ instance_max_pts: 5000 # max initial points for each instance
+ only_moving: true # only optimize moving instances
+ traj_length_thres: 1.0 # threshold of trajectory length for moving instances
+ ctrl:
+ cull_scale_thresh: 0.1
+ stop_screen_size_at: 30000
+ stop_split_at: 20000
+ cull_out_of_bound: true
+ reg:
+ sharp_shape_reg:
+ w: 1.
+ step_interval: 10
+ max_gauss_ratio: 10.
+ temporal_smooth_reg:
+ trans:
+ w: 0.01 # optimal value may vary
+ smooth_range: 5 # no ablation
+ optim:
+ ins_rotation:
+ lr: 0.00001
+ lr_final: 0.000005
+ ins_translation:
+ lr: 0.0005
+ lr_final: 0.0001
+ Sky:
+ type: models.modules.EnvLight
+ params:
+ resolution: 1024
+ optim:
+ all:
+ lr: 0.01
+ Affine:
+ type: models.modules.AffineTransform
+ params:
+ embedding_dim: 4
+ base_mlp_layer_width: 64
+ pixel_affine: False
+ optim:
+ all:
+ lr: 1.0e-5
+ weight_decay: 1.0e-6
+ CamPose:
+ type: models.modules.CameraOptModule
+ optim:
+ all:
+ lr: 1.0e-5
+ weight_decay: 1.0e-6
+
+# ------------- render ------------ #
+render:
+ fps: 10 # frames per second for the main rendered output
+ render_full: True # whether to render full resolution videos
+ render_test: True # whether to render test set
+ render_novel:
+ traj_types:
+ - front_center_interp # type of trajectory for novel view synthesis
+ fps: 24 # frames per second for novel view rendering
+ vis_lidar: False # whether to visualize lidar points on ground truth images
+ vis_sky: False # whether to include "rgb_sky" and "rgb_sky_blend" in rendered keys
+ vis_error: False # whether to include "rgb_error_map" in rendered keys
+
+# ------------- logging ------------ #
+logging:
+ vis_freq: 2000 # how often to visualize training stats
+ print_freq: 500 # how often to print training stats
+ saveckpt_freq: 15000 # how often to save checkpoints
+ save_seperate_video: True # whether to save seperate videos for each scene
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/__init__.py b/worldbench/third_party/drivestudio/datasets/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/worldbench/third_party/drivestudio/datasets/argoverse/argoverse_human_utils.py b/worldbench/third_party/drivestudio/datasets/argoverse/argoverse_human_utils.py
new file mode 100644
index 0000000..e853907
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/argoverse/argoverse_human_utils.py
@@ -0,0 +1,218 @@
+from typing import List
+import os
+import cv2
+import json
+import logging
+import numpy as np
+from tqdm import tqdm
+
+from utils.geometry import (
+ get_corners,
+ project_camera_points_to_image
+)
+from .argoverse_sourceloader import (
+ SMPLNODE_CLASSES,
+ OPENCV2DATASET,
+ AVAILABLE_CAM_LIST,
+)
+
+logger = logging.getLogger()
+
+CAMERA_LIST = AVAILABLE_CAM_LIST
+
+def project_human_boxes(
+ scene_dir: str,
+ camera_list: List[int],
+ save_temp=True,
+ verbose=False,
+ narrow_width_ratio=0.2,
+ fps=12
+):
+ """Project human boxes to 2D image space and save the results to pkl files
+
+ Args:
+ scene_dir: str, path to the scene directory
+ camera_list: List[int], a list of camera ids to be processed
+ save_temp: bool, whether to save the intermediate results
+ verbose: bool, whether to visualize the projected boxes
+ narrow_width_ratio: sometimes the projected boxes are too wide
+ we can narrow them with this ratio to get more accurate results
+ fps: int, FPS for the visualization video
+
+ Returns:
+ collector_all: dict, a dictionary containing the projected boxes for each camera
+ """
+ # check if the necessary directories exist
+ images_dir = f'{scene_dir}/images'
+ poses_dir = f'{scene_dir}/ego_pose'
+ extrinsics_dir = f'{scene_dir}/extrinsics'
+ intrinsics_dir = f'{scene_dir}/intrinsics'
+ instances_dir = f'{scene_dir}/instances'
+ valid_paths = [images_dir, poses_dir, extrinsics_dir, intrinsics_dir, instances_dir]
+ for path in valid_paths:
+ assert os.path.exists(path), \
+ f"Path {path} does not exist, you need to run waymo preprocess to generate the necessary files"
+
+ # create directories for saving the results
+ save_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks'
+ if not os.path.exists(save_dir):
+ os.makedirs(save_dir)
+
+ if verbose:
+ # create directories for saving the intermediate visualization results
+ video_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks/vis'
+ if not os.path.exists(video_dir):
+ os.makedirs(video_dir)
+ per_human_img_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks/vis/images'
+ if not os.path.exists(per_human_img_dir):
+ os.makedirs(per_human_img_dir)
+
+ # load instances and frame infos
+ frame_infos = json.load(
+ open(f'{instances_dir}/frame_instances.json')
+ )
+ instances_meta = json.load(
+ open(f'{instances_dir}/instances_info.json')
+ )
+
+ collector_all = {}
+ # iterate over each camera
+ for cam_id in camera_list:
+ # check if already processed
+ pkl_path = os.path.join(save_dir, f"{cam_id}.pkl")
+ if os.path.exists(pkl_path):
+ collector_all[cam_id] = json.load(open(pkl_path))
+ logger.info(f"Results for camera {cam_id} already exists at {pkl_path}")
+ continue
+
+ if verbose:
+ per_cam_vis_dir = os.path.join(per_human_img_dir, f"{cam_id}")
+ if not os.path.exists(per_cam_vis_dir):
+ os.makedirs(per_cam_vis_dir)
+
+ collector = {}
+ frames = []
+ for frame_id, frame_ins_list in frame_infos.items():
+ frame_id = int(frame_id)
+
+ # define empty instance collector for each frame
+ frame_collector = {
+ "gt_bbox": [],
+ "extra_data": {
+ "gt_track_id": [],
+ "gt_class": [],
+ }
+ }
+
+ # load extrinsic
+ cam_to_ego = np.loadtxt(os.path.join(extrinsics_dir, f"{cam_id}.txt"))
+ cam_to_ego = cam_to_ego @ OPENCV2DATASET
+ ego_to_world = np.loadtxt(os.path.join(poses_dir, f"{str(frame_id).zfill(3)}.txt"))
+ cam2world = ego_to_world @ cam_to_ego
+
+ # load intrinsic
+ Ks = np.loadtxt(os.path.join(intrinsics_dir, f"{cam_id}.txt"))
+ fx, fy, cx, cy = Ks[0], Ks[1], Ks[2], Ks[3]
+ # k1, k2, p1, p2, k3 = Ks[4], Ks[5], Ks[6], Ks[7], Ks[8]
+ intrinsic = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float32)
+
+ # load image
+ ori_image = cv2.imread(
+ os.path.join(images_dir, f"{str(frame_id).zfill(3)}_{cam_id}.jpg")
+ )
+ image = ori_image.copy()
+ H, W = image.shape[:2]
+
+ # iterate over each instance in the frame
+ if len(frame_ins_list) > 0:
+ # if there are pedestrians in this frame, project the boxes to the image
+ image_plotted = image.copy()
+ for instance_id in frame_ins_list:
+ ins = instances_meta[str(instance_id)]
+
+ if ins["class_name"] not in SMPLNODE_CLASSES:
+ continue
+
+ ins_anno = ins["frame_annotations"]
+ index = ins_anno['frame_idx'].index(frame_id)
+ obj_to_world = np.array(ins_anno['obj_to_world'][index])
+ l, w, h = ins_anno['box_size'][index]
+
+ # get box corners in object space
+ corners = get_corners(l, w, h)
+ # transform box corners to world space
+ corners_world = obj_to_world[:3, :3] @ corners + obj_to_world[:3, 3:4]
+ # transform box corners to image space
+ world2cam = np.linalg.inv(cam2world)
+ corners_cam = world2cam[:3, :3] @ corners_world + world2cam[:3, 3:4]
+ cam_points, depth = project_camera_points_to_image(corners_cam.T, intrinsic)
+
+ x_min, y_min = np.min(cam_points, axis=0)
+ x_max, y_max = np.max(cam_points, axis=0)
+ # clip left and right with this ratio
+ if narrow_width_ratio > 0.:
+ length = x_max - x_min
+ x_min += length * narrow_width_ratio
+ x_max -= length * narrow_width_ratio
+
+ # clip the box to the image
+ original_area = (x_max - x_min) * (y_max - y_min)
+ x_min, x_max = np.clip(x_min, 0, W), np.clip(x_max, 0, W)
+ y_min, y_max = np.clip(y_min, 0, H), np.clip(y_max, 0, H)
+ new_area = (x_max - x_min) * (y_max - y_min)
+
+ # filter out boxes that are too small or too large
+ behind = depth.max() < 0
+ too_small = new_area < W * H * (0.03)**2
+ too_large = new_area > W * H / 1.1
+ too_far = np.linalg.norm(obj_to_world[:3, 3] - cam2world[:3, 3]) > 40
+ clip_large = new_area / original_area < 1/3
+ if too_small or too_large or clip_large or behind or too_far:
+ continue
+
+ # gt box on image
+ gt_box = [x_min, y_min, x_max - x_min, y_max - y_min]
+
+ # save the projected box to the collector
+ frame_collector["gt_bbox"].append(gt_box)
+ frame_collector["extra_data"]["gt_track_id"].append(instance_id)
+ frame_collector["extra_data"]["gt_class"].append([0])
+
+ if verbose:
+ # visualize the projected boxes of ONE instance
+ raw_image = cv2.rectangle(
+ ori_image.copy(), (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 2
+ )
+ raw_image_path = os.path.join(per_cam_vis_dir, f"{frame_id}_{instance_id}.jpg")
+ cv2.imwrite(raw_image_path, raw_image)
+
+ # add this instance to the image
+ image_plotted = cv2.rectangle(
+ image, (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 2
+ )
+
+ if verbose:
+ frames.append(image_plotted)
+ else:
+ # if no instance in this frame, just save the original image
+ if verbose:
+ frames.append(ori_image)
+
+ collector[frame_id] = frame_collector
+
+ if verbose:
+ height, width = frames[0].shape[:2]
+ output_path = os.path.join(video_dir, f"cam_{cam_id}.mp4")
+ fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+ out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+ for frame in tqdm(frames, desc=f"Writing video for camera {cam_id}"):
+ out.write(frame)
+ out.release()
+
+ if save_temp:
+ # save collector to pkl
+ json.dump(collector, open(pkl_path, "w"))
+
+ collector_all[cam_id] = collector
+
+ return collector_all
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/argoverse/argoverse_preprocess.py b/worldbench/third_party/drivestudio/datasets/argoverse/argoverse_preprocess.py
new file mode 100644
index 0000000..e8774d3
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/argoverse/argoverse_preprocess.py
@@ -0,0 +1,571 @@
+import os
+import json
+from pathlib import Path
+from typing import List, Union
+
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+
+from av2.datasets.sensor.constants import RingCameras, StereoCameras
+from av2.datasets.sensor.sensor_dataloader import SensorDataloader, SynchronizedSensorData
+from datasets.tools.multiprocess_utils import track_parallel_progress
+from utils.visualization import dump_3d_bbox_on_image, color_mapper
+
+AV2_LABELS = [
+ "ANIMAL", "ARTICULATED_BUS", "BICYCLE", "BICYCLIST",
+ "BOLLARD", "BOX_TRUCK", "BUS", "CONSTRUCTION_BARREL",
+ "CONSTRUCTION_CONE", "DOG", "LARGE_VEHICLE", "MESSAGE_BOARD_TRAILER",
+ "MOBILE_PEDESTRIAN_CROSSING_SIGN", "MOTORCYCLE", "MOTORCYCLIST",
+ "OFFICIAL_SIGNALER", "PEDESTRIAN", "RAILED_VEHICLE", "REGULAR_VEHICLE",
+ "SCHOOL_BUS", "SIGN", "STOP_SIGN", "STROLLER", "TRAFFIC_LIGHT_TRAILER",
+ "TRUCK", "TRUCK_CAB", "VEHICULAR_TRAILER", "WHEELCHAIR",
+ "WHEELED_DEVICE", "WHEELED_RIDER"
+]
+
+AV2_NONRIGID_DYNAMIC_CLASSES = [
+ "BICYCLIST", "DOG", "MOTORCYCLIST", "PEDESTRIAN", "STROLLER",
+ "WHEELCHAIR", "WHEELED_DEVICE", "WHEELED_RIDER"
+]
+
+AV2_RIGID_DYNAMIC_CLASSES = [
+ "ARTICULATED_BUS", "BOX_TRUCK", "BUS", "LARGE_VEHICLE",
+ "MOTORCYCLE", "RAILED_VEHICLE", "REGULAR_VEHICLE", "SCHOOL_BUS",
+ "TRUCK", "TRUCK_CAB", "VEHICULAR_TRAILER"
+]
+
+AV2_DYNAMIC_CLASSES = AV2_NONRIGID_DYNAMIC_CLASSES + AV2_RIGID_DYNAMIC_CLASSES
+
+valid_ring_cams = set([x.value for x in RingCameras])
+valid_stereo_cams = set([x.value for x in StereoCameras])
+
+class ArgoVerseProcessor(object):
+ """Process ArgoVerse.
+
+ LiDAR: 10Hz, Camera: 20Hz
+ Since the LiDAR and Camera are not synchronized, we need to find the closest camera image for each LiDAR frame.
+ Thus the actual frame rate of the processed data is 10Hz, which is aligned with the LiDAR.
+
+ Args:
+ load_dir (str): Directory to load data.
+ save_dir (str): Directory to save data.
+ workers (int, optional): Number of workers for the parallel process.
+ Defaults to 64.
+ Defaults to False.
+ """
+ def __init__(
+ self,
+ load_dir,
+ save_dir,
+ process_keys=[
+ "images",
+ "lidar",
+ "calib",
+ "pose",
+ "dynamic_masks",
+ "objects"
+ ],
+ process_id_list=None,
+ workers=64,
+ ):
+ self.process_id_list = process_id_list
+ self.process_keys = process_keys
+ print("will process keys: ", self.process_keys)
+
+ # ArgoVerse Provides 7 cameras, we process 5 of them
+ self.cam_list = [ # {frame_idx}_{cam_id}.jpg
+ "ring_front_center", # "xxx_0.jpg"
+ "ring_front_left", # "xxx_1.jpg"
+ "ring_front_right", # "xxx_2.jpg"
+ "ring_side_left", # "xxx_3.jpg"
+ "ring_side_right", # "xxx_4.jpg"
+ "ring_rear_left", # "xxx_5.jpg"
+ "ring_rear_right", # "xxx_6.jpg"
+ ]
+ cam_enums: List[Union[RingCameras, StereoCameras]] = []
+ for cam_name in self.cam_list:
+ if cam_name in valid_ring_cams:
+ cam_enums.append(RingCameras(cam_name))
+ elif cam_name in valid_stereo_cams:
+ cam_enums.append(StereoCameras(cam_name))
+ else:
+ raise ValueError("Must provide _valid_ camera names!")
+
+ # Prepare dynamic objects' metadata
+ self.load_dir = Path(load_dir)
+ self.save_dir = f"{save_dir}"
+ self.workers = int(workers)
+ self.av2loader = SensorDataloader(
+ dataset_dir=self.load_dir,
+ with_annotations=True,
+ with_cache=True,
+ cam_names=tuple(cam_enums)
+ )
+ # a list of tfrecord pathnames
+ self.training_files = open("data/argoverse_train_list.txt").read().splitlines()
+ self.log_pathnames = [
+ f"{self.load_dir}/{f}" for f in self.training_files
+ ]
+ self.create_folder()
+
+ def convert(self):
+ """Convert action."""
+ print("Start converting ...")
+ if self.process_id_list is None:
+ id_list = range(len(self))
+ else:
+ id_list = self.process_id_list
+ track_parallel_progress(self.convert_one, id_list, self.workers)
+ print("\nFinished ...")
+
+ def get_lidar_indices(self, log_id: str):
+ sensor_cache = self.av2loader.sensor_cache
+ lidar_indices = sensor_cache.xs(key="lidar", level=2).index
+ lidar_indices_mask = lidar_indices.get_level_values('log_id') == log_id
+ # get the positions of the lidar indices, get nonzeros
+ lidar_indices = np.nonzero(lidar_indices_mask)[0]
+ return lidar_indices
+
+ def filter_lidar_indices(self, lidar_indices):
+ """
+ Filter lidar indices whose corresponding synchronized camera images are not complete.
+ These usually happen at the beginning and the end of the sequence.
+ """
+ valid_list = []
+ invalid_list = []
+ for idx in lidar_indices:
+ datum = self.av2loader[idx]
+ sweep = datum.sweep
+
+ timestamp_city_SE3_ego_dict = datum.timestamp_city_SE3_ego_dict
+ synchronized_imagery = datum.synchronized_imagery
+
+ cnt = 0
+ for _, cam in synchronized_imagery.items():
+ if (
+ cam.timestamp_ns in timestamp_city_SE3_ego_dict
+ and sweep.timestamp_ns in timestamp_city_SE3_ego_dict
+ ):
+ cnt += 1
+
+ if cnt != len(self.cam_list):
+ invalid_list.append(idx)
+ continue
+ valid_list.append(idx)
+ print(f"INFO: {len(invalid_list)} lidar indices filtered")
+
+ return valid_list
+
+ def convert_one(self, scene_idx):
+ """Convert action for single file.
+
+ Args:
+ scene_idx (str): Scene index.
+ """
+ lidar_indices = self.get_lidar_indices(
+ self.training_files[scene_idx]
+ )
+ lidar_indices = self.filter_lidar_indices(lidar_indices)
+
+ # process each frame
+ num_frames = len(lidar_indices)
+ for frame_idx, lidar_idx in tqdm(
+ enumerate(lidar_indices), desc=f"File {scene_idx}", total=num_frames, dynamic_ncols=True
+ ):
+ datum = self.av2loader[lidar_idx]
+ if "images" in self.process_keys:
+ self.save_image(datum, scene_idx, frame_idx)
+ if "calib" in self.process_keys:
+ self.save_calib(datum, scene_idx, frame_idx)
+ if "lidar" in self.process_keys:
+ self.save_lidar(datum, scene_idx, frame_idx)
+ if "pose" in self.process_keys:
+ self.save_pose(datum, scene_idx, frame_idx)
+ if "3dbox_vis" in self.process_keys:
+ # visualize 3d box, debug usage
+ self.visualize_3dbox(datum, scene_idx, frame_idx)
+ if "dynamic_masks" in self.process_keys:
+ self.save_dynamic_mask(datum, scene_idx, frame_idx, class_valid='all')
+ self.save_dynamic_mask(datum, scene_idx, frame_idx, class_valid='human')
+ self.save_dynamic_mask(datum, scene_idx, frame_idx, class_valid='vehicle')
+
+ # sort and save objects info
+ if "objects" in self.process_keys:
+ instances_info, frame_instances = self.save_objects(lidar_indices)
+ print(f"Processed instances info for {scene_idx}")
+
+ # Save instances info and frame instances
+ object_info_dir = f"{self.save_dir}/{str(scene_idx).zfill(3)}/instances"
+ with open(f"{object_info_dir}/instances_info.json", "w") as fp:
+ json.dump(instances_info, fp, indent=4)
+ with open(f"{object_info_dir}/frame_instances.json", "w") as fp:
+ json.dump(frame_instances, fp, indent=4)
+
+ # verbose: visualize the instances on the image (Debug Usage)
+ if "objects_vis" in self.process_keys:
+ self.visualize_dynamic_objects(
+ scene_idx, lidar_indices,
+ instances_info=instances_info,
+ frame_instances=frame_instances
+ )
+ print(f"Processed objects visualization for {scene_idx}")
+
+ def __len__(self):
+ """Length of the filename list."""
+ return len(self.process_id_list)
+
+ def save_image(self, datum: SynchronizedSensorData, scene_idx, frame_idx):
+ """Parse and save the images in jpg format.
+
+ Args:
+ datum (:obj:`SynchronizedSensorData`): ArgoVerse synchronized sensor data.
+ scene_idx (str): Current file index.
+ frame_idx (int): Current frame index.
+ """
+ synchronized_imagery = datum.synchronized_imagery
+ for cam_name, cam in synchronized_imagery.items():
+ idx = self.cam_list.index(cam_name)
+ img_path = (
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/images/"
+ + f"{str(frame_idx).zfill(3)}_{str(idx)}.jpg"
+ )
+ image = Image.fromarray(cam.img[:, :, [2, 1, 0]]) # BGR to RGB
+ image.save(img_path)
+
+ def save_calib(self, datum: SynchronizedSensorData, scene_idx, frame_idx):
+ """Parse and save the calibration data.
+
+ Args:
+ datum (:obj:`SynchronizedSensorData`): ArgoVerse synchronized sensor data.
+ scene_idx (str): Current file index.
+ frame_idx (int): Current frame index.
+ """
+ synchronized_imagery = datum.synchronized_imagery
+ for idx, cam_name in enumerate(self.cam_list):
+ cam = synchronized_imagery[cam_name]
+ c2v = cam.camera_model.ego_SE3_cam.transform_matrix
+ K = cam.camera_model.intrinsics
+ intrinsics = [K.fx_px, K.fy_px, K.cx_px, K.cy_px, 0.0, 0.0, 0.0, 0.0, 0.0]
+
+ np.savetxt(
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/extrinsics/"
+ + f"{str(idx)}.txt",
+ c2v,
+ )
+ np.savetxt(
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/intrinsics/"
+ + f"{str(idx)}.txt",
+ intrinsics,
+ )
+
+ def save_lidar(self, datum: SynchronizedSensorData, scene_idx, frame_idx):
+ """Parse and save the lidar data in psd format.
+
+ Args:
+ datum (:obj:`SynchronizedSensorData`): ArgoVerse synchronized sensor data.
+ scene_idx (str): Current file index.
+ frame_idx (int): Current frame index.
+ """
+ sweep = datum.sweep
+
+ point_cloud = np.column_stack(
+ (
+ sweep.xyz, # lidar points in ego frame
+ sweep.intensity, # intensity
+ )
+ )
+
+ pc_path = (
+ f"{self.save_dir}/"
+ + f"{str(scene_idx).zfill(3)}/lidar/{str(frame_idx).zfill(3)}.bin"
+ )
+ point_cloud.astype(np.float32).tofile(pc_path)
+
+ def save_pose(self, datum: SynchronizedSensorData, scene_idx, frame_idx):
+ """Parse and save the pose data.
+
+ Args:
+ datum (:obj:`SynchronizedSensorData`): ArgoVerse synchronized sensor data.
+ scene_idx (str): Current file index.
+ frame_idx (int): Current frame index.
+ """
+ sweep = datum.sweep
+ timestamp_city_SE3_ego_dict = datum.timestamp_city_SE3_ego_dict
+ synchronized_imagery = datum.synchronized_imagery
+ for cam_name, cam in synchronized_imagery.items():
+ if cam_name == "ring_front_center":
+ # we use the ego state at the time of the lidar as the ego pose
+ ego_pose_t_lidar = timestamp_city_SE3_ego_dict[sweep.timestamp_ns]
+
+ np.savetxt(
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/ego_pose/"
+ + f"{str(frame_idx).zfill(3)}.txt",
+ ego_pose_t_lidar.transform_matrix,
+ )
+
+ def visualize_3dbox(self, datum: SynchronizedSensorData, scene_idx, frame_idx):
+ """DEBUG: Visualize the 3D bounding box on the image.
+ Visualize the 3D bounding box all with the same COLOR.
+ If you want to visualize the 3D bounding box with different colors, please use the `visualize_dynamic_objects` function.
+
+ Args:
+ datum (:obj:`SynchronizedSensorData`): ArgoVerse synchronized sensor data.
+ scene_idx (str): Current file index.
+ frame_idx (int): Current frame index.
+ """
+ annotations = datum.annotations
+ synchronized_imagery = datum.synchronized_imagery
+ for idx, cam_name in enumerate(self.cam_list):
+ cam = synchronized_imagery[cam_name]
+ if annotations is not None:
+ img_plotted = annotations.project_to_cam(
+ img=cam.img,
+ cam_model=cam.camera_model,
+ )
+ else:
+ img_plotted = cam.img.copy()
+
+ # save
+ img_path = (
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/3dbox_vis/"
+ + f"{str(frame_idx).zfill(3)}_{str(idx)}.jpg"
+ )
+ Image.fromarray(
+ img_plotted[:, :, [2, 1, 0]]
+ ).save(img_path)
+
+ def visualize_dynamic_objects(
+ self, scene_idx, lidar_indices,
+ instances_info, frame_instances
+ ):
+ """DEBUG: Visualize the dynamic objects'box with different colors on the image.
+
+ Args:
+ scene_idx (str): Current file index.
+ lidar_indices (list): List of lidar indices.
+ instances_info (dict): Instances information.
+ frame_instances (dict): Frame instances.
+ """
+ output_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/instances/debug_vis"
+
+ print("Visualizing dynamic objects ...")
+ for frame_idx, lidar_idx in tqdm(
+ enumerate(lidar_indices), desc=f"Visualizing dynamic objects of scene {scene_idx} ...", total=len(lidar_indices), dynamic_ncols=True
+ ):
+ datum = self.av2loader[lidar_idx]
+ synchronized_imagery = datum.synchronized_imagery
+ for cam_name, cam in synchronized_imagery.items():
+ cam_idx = self.cam_list.index(cam_name)
+ img_path = (
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/images/"
+ + f"{str(frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+ )
+ canvas = np.array(Image.open(img_path))
+
+ if frame_idx in frame_instances:
+ objects = frame_instances[frame_idx]
+
+ if len(objects) == 0:
+ img_plotted = canvas
+ else:
+ lstProj2d = []
+ color_list = []
+ for obj_id in objects:
+ idx_in_obj = instances_info[obj_id]['frame_annotations']['frame_idx'].index(frame_idx)
+ o2w = np.array(
+ instances_info[obj_id]['frame_annotations']['obj_to_world'][idx_in_obj]
+ )
+ length, width, height = instances_info[obj_id]['frame_annotations']['box_size'][idx_in_obj]
+ half_dim_x, half_dim_y, half_dim_z = length/2.0, width/2.0, height/2.0
+ corners = np.array(
+ [[half_dim_x, half_dim_y, -half_dim_z],
+ [half_dim_x, -half_dim_y, -half_dim_z],
+ [-half_dim_x, -half_dim_y, -half_dim_z],
+ [-half_dim_x, half_dim_y, -half_dim_z],
+ [half_dim_x, half_dim_y, half_dim_z],
+ [half_dim_x, -half_dim_y, half_dim_z],
+ [-half_dim_x, -half_dim_y, half_dim_z],
+ [-half_dim_x, half_dim_y, half_dim_z]]
+ )
+ corners = (o2w[:3, :3] @ corners.T + o2w[:3, [3]]).T
+ v2w = datum.timestamp_city_SE3_ego_dict[cam.timestamp_ns]
+ w2v = v2w.inverse()
+ corners_in_ego = w2v.transform_point_cloud(corners)
+
+ projected_points2d, _, ok = cam.camera_model.project_ego_to_img(
+ corners_in_ego # cuboid corners in ego frame
+ )
+ projected_points2d = projected_points2d.tolist()
+ if all(ok):
+ lstProj2d.append(projected_points2d)
+ color_list.append(color_mapper(obj_id))
+
+ lstProj2d = np.asarray(lstProj2d)
+ img_plotted = dump_3d_bbox_on_image(coords=lstProj2d, img=canvas, color=color_list)
+
+ img_path = (
+ f"{output_path}/"
+ + f"{str(frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+ )
+ Image.fromarray(img_plotted).save(img_path)
+
+ def save_dynamic_mask(self, datum: SynchronizedSensorData, scene_idx, frame_idx, class_valid='all'):
+ """Parse and save the segmentation data.
+
+ Args:
+ datum (:obj:`SynchronizedSensorData`): ArgoVerse synchronized sensor data.
+ scene_idx (str): Current file index.
+ frame_idx (int): Current frame index.
+ class_valid (str): Class valid for dynamic mask.
+ """
+ assert class_valid in ['all', 'human', 'vehicle'], "Invalid class valid"
+ if class_valid == 'all':
+ VALID_CLASSES = AV2_DYNAMIC_CLASSES
+ elif class_valid == 'human':
+ VALID_CLASSES = AV2_NONRIGID_DYNAMIC_CLASSES
+ elif class_valid == 'vehicle':
+ VALID_CLASSES = AV2_RIGID_DYNAMIC_CLASSES
+ mask_dir = f"{self.save_dir}/{str(scene_idx).zfill(3)}/dynamic_masks/{class_valid}"
+ if not os.path.exists(mask_dir):
+ os.makedirs(mask_dir)
+
+ annotations = datum.annotations
+ synchronized_imagery = datum.synchronized_imagery
+
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ cam = synchronized_imagery[cam_name]
+ H, W = cam.img.shape[:2]
+ dynamic_mask = np.zeros_like(cam.img, dtype=np.float32)[..., 0]
+
+ for cuboid_idx in range(len(annotations)):
+ cuboid = annotations[cuboid_idx]
+ if cuboid.category not in VALID_CLASSES:
+ continue
+
+ uv, _, ok = cam.camera_model.project_ego_to_img(
+ cuboid.vertices_m # cuboid corners in ego frame
+ )
+
+ # Skip object if any corner projection failed. Note that this is very
+ # strict and can lead to exclusion of some partially visible objects.
+ if not all(ok):
+ continue
+ u, v= uv[:, 0], uv[:, 1]
+ u = u.astype(np.int32)
+ v = v.astype(np.int32)
+
+ # Clip box to image bounds.
+ u = np.clip(u, 0, W - 1)
+ v = np.clip(v, 0, H - 1)
+
+ if u.max() - u.min() == 0 or v.max() - v.min() == 0:
+ continue
+
+ # Draw projected 2D box onto the image.
+ xy = (u.min(), v.min())
+ width = u.max() - u.min()
+ height = v.max() - v.min()
+ # max pooling
+ dynamic_mask[
+ int(xy[1]) : int(xy[1] + height),
+ int(xy[0]) : int(xy[0] + width),
+ ] = np.maximum(
+ dynamic_mask[
+ int(xy[1]) : int(xy[1] + height),
+ int(xy[0]) : int(xy[0] + width),
+ ],
+ 1.,
+ )
+ dynamic_mask = np.clip((dynamic_mask > 0.) * 255, 0, 255).astype(np.uint8)
+ dynamic_mask = Image.fromarray(dynamic_mask, "L")
+ dynamic_mask_path = os.path.join(mask_dir, f"{str(frame_idx).zfill(3)}_{str(cam_idx)}.png")
+ dynamic_mask.save(dynamic_mask_path)
+
+ def save_objects(self, lidar_indices: List[int]):
+ """Parse and save the objects annotation data.
+
+ Args:
+ lidar_indices (list): List of lidar indices.
+ """
+ instances_info, frame_instances = {}, {}
+ for frame_idx, lidar_idx in enumerate(lidar_indices):
+ datum = self.av2loader[lidar_idx]
+ annotations = datum.annotations
+ sweep = datum.sweep
+ timestamp_city_SE3_ego_dict = datum.timestamp_city_SE3_ego_dict
+
+ frame_instances[frame_idx] = []
+ for cuboid_idx in range(len(annotations)):
+ cuboid = annotations[cuboid_idx]
+ track_id, label = cuboid.track_uuid, cuboid.category
+ if label not in AV2_DYNAMIC_CLASSES:
+ continue
+
+ if track_id not in instances_info:
+ instances_info[track_id] = dict(
+ id=track_id,
+ class_name=label,
+ frame_annotations={
+ "frame_idx": [],
+ "obj_to_world": [],
+ "box_size": [],
+ }
+ )
+
+ o2v = cuboid.dst_SE3_object.transform_matrix
+ v2w = timestamp_city_SE3_ego_dict[sweep.timestamp_ns].transform_matrix
+ # [object to world] transformation matrix
+ o2w = v2w @ o2v
+
+ # Dimensions of the box. length: dim x. width: dim y. height: dim z.
+ # length: dim_x: along heading; dim_y: verticle to heading; dim_z: verticle up
+ dimension = [cuboid.length_m, cuboid.width_m, cuboid.height_m]
+
+ instances_info[track_id]['frame_annotations']['frame_idx'].append(frame_idx)
+ instances_info[track_id]['frame_annotations']['obj_to_world'].append(o2w.tolist())
+ instances_info[track_id]['frame_annotations']['box_size'].append(dimension)
+
+ frame_instances[frame_idx].append(track_id)
+
+ # Correct ID mapping
+ id_map = {}
+ for i, (k, v) in enumerate(instances_info.items()):
+ id_map[v["id"]] = i
+
+ # Update keys in instances_info
+ new_instances_info = {}
+ for k, v in instances_info.items():
+ new_instances_info[id_map[v["id"]]] = v
+
+ # Update keys in frame_instances
+ new_frame_instances = {}
+ for k, v in frame_instances.items():
+ new_frame_instances[k] = [id_map[i] for i in v]
+
+ return new_instances_info, new_frame_instances
+
+ def create_folder(self):
+ """Create folder for data preprocessing."""
+ if self.process_id_list is None:
+ id_list = range(len(self))
+ else:
+ id_list = self.process_id_list
+ for i in id_list:
+ if "images" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/images", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/sky_masks", exist_ok=True)
+ if "calib" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/extrinsics", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/intrinsics", exist_ok=True)
+ if "pose" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/ego_pose", exist_ok=True)
+ if "lidar" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/lidar", exist_ok=True)
+ if "3dbox_vis" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/3dbox_vis", exist_ok=True)
+ if "dynamic_masks" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/dynamic_masks", exist_ok=True)
+ if "objects" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/instances", exist_ok=True)
+ if "objects_vis" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/instances/debug_vis", exist_ok=True)
diff --git a/worldbench/third_party/drivestudio/datasets/argoverse/argoverse_preview.py b/worldbench/third_party/drivestudio/datasets/argoverse/argoverse_preview.py
new file mode 100644
index 0000000..23b44f2
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/argoverse/argoverse_preview.py
@@ -0,0 +1,298 @@
+import os
+import cv2
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+from moviepy.editor import VideoFileClip, clips_array
+from typing import Mapping, Optional, Union
+import numpy.typing as npt
+
+NDArrayByte = npt.NDArray[np.uint8]
+
+def tile_cameras(
+ named_sensors: Mapping[str, Union[NDArrayByte, pd.DataFrame]],
+ bev_img: Optional[NDArrayByte] = None,
+) -> NDArrayByte:
+ """Combine ring cameras into a tiled image.
+
+ NOTE: Images are expected in BGR ordering.
+
+ Layout:
+
+ ##########################################################
+ # ring_front_left # ring_front_center # ring_front_right #
+ ##########################################################
+ # ring_side_left # # ring_side_right #
+ ##########################################################
+ ############ ring_rear_left # ring_rear_right ############
+ ##########################################################
+
+ Args:
+ named_sensors: Dictionary of camera names to the (width, height, 3) images.
+ bev_img: (H,W,3) Bird's-eye view image.
+
+ Returns:
+ Tiled image.
+ """
+ landscape_height = 2048
+ landscape_width = 1550
+ for _, v in named_sensors.items():
+ landscape_width = max(v.shape[0], v.shape[1])
+ landscape_height = min(v.shape[0], v.shape[1])
+ break
+
+ height = landscape_height + landscape_height
+ width = landscape_width + landscape_height + landscape_width
+ tiled_im_bgr: NDArrayByte = np.zeros((height, width, 3), dtype=np.uint8)
+
+ if "ring_front_left" in named_sensors:
+ ring_front_left = named_sensors["ring_front_left"]
+ tiled_im_bgr[:landscape_height, :landscape_width] = ring_front_left
+
+ if "ring_front_center" in named_sensors:
+ ring_front_center = named_sensors["ring_front_center"]
+ tiled_im_bgr[
+ :landscape_width, landscape_width : landscape_width + landscape_height
+ ] = ring_front_center
+
+ if "ring_front_right" in named_sensors:
+ ring_front_right = named_sensors["ring_front_right"]
+ tiled_im_bgr[:landscape_height, landscape_width + landscape_height :] = (
+ ring_front_right
+ )
+
+ if "ring_side_left" in named_sensors:
+ ring_side_left = named_sensors["ring_side_left"]
+ tiled_im_bgr[landscape_height : 2 * landscape_height, :landscape_width] = (
+ ring_side_left
+ )
+
+ if "ring_side_right" in named_sensors:
+ ring_side_right = named_sensors["ring_side_right"]
+ tiled_im_bgr[
+ landscape_height : 2 * landscape_height,
+ landscape_width + landscape_height :,
+ ] = ring_side_right
+
+ if bev_img is not None:
+ tiled_im_bgr[
+ landscape_width : 2 * landscape_width,
+ landscape_width : landscape_width + landscape_height,
+ ] = bev_img
+
+ if "ring_rear_left" in named_sensors:
+ ring_rear_left = named_sensors["ring_rear_left"]
+ tiled_im_bgr[2 * landscape_height : 3 * landscape_height, :landscape_width] = (
+ ring_rear_left
+ )
+
+ if "ring_rear_right" in named_sensors:
+ ring_rear_right = named_sensors["ring_rear_right"]
+ tiled_im_bgr[
+ 2 * landscape_height : 3 * landscape_height, width - landscape_width :
+ ] = ring_rear_right
+ return tiled_im_bgr
+
+def preview_camera_captures(
+ input_root: str = 'data/argoverse/processed/training',
+ output_root: str = 'data/argoverse/preview/camera_captures',
+ frame_width: int = int(5646 / 4), # Concatenate five camera images horizontally
+ frame_height: int = int(3100 / 4),
+ fps: int = 16
+):
+ if not os.path.exists(output_root):
+ os.makedirs(output_root)
+
+ # Iterate over each scene directory in the input root directory
+ for scene_id in tqdm(sorted(os.listdir(input_root))):
+ scene_folder = os.path.join(input_root, scene_id, "images")
+
+ # Define the output video file path
+ output_video_path = os.path.join(output_root, f'{scene_id}.mp4')
+
+ # Get paths to each camera image folder
+ images_all = sorted([f for f in sorted(os.listdir(scene_folder)) if f.lower().endswith('.jpg')])
+ num_images = len(images_all)
+ num_frames = num_images // 6
+
+ out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
+ for frame_id in range(num_frames):
+ front = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5]))
+ front_left = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5 + 1]))
+ front_right = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5 + 2]))
+ left = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5 + 3]))
+ right = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5 + 4]))
+
+ combined_img = tile_cameras({
+ "ring_front_left": front_left,
+ "ring_front_center": front,
+ "ring_front_right": front_right,
+ "ring_side_left": left,
+ "ring_side_right": right
+ })
+ combined_img = cv2.resize(combined_img, (int(frame_width), int(frame_height)))
+
+ out.write(combined_img)
+ out.release()
+
+def preview_3dbox_projection(
+ input_root: str = 'data/argoverse/processed/training',
+ output_root: str = 'data/argoverse/preview/3dbox_vis',
+ frame_width: int = int(5646 / 4),
+ frame_height: int = int(3100 / 4),
+ fps: int = 24
+):
+ if not os.path.exists(output_root):
+ os.makedirs(output_root)
+
+ # Iterate over each scene directory in the input root directory
+ for scene_id in tqdm(sorted(os.listdir(input_root))):
+ scene_folder = os.path.join(input_root, scene_id, "3dbox_vis")
+
+ # Define the output video file path
+ output_video_path = os.path.join(output_root, f'{scene_id}.mp4')
+
+ # Get paths to each camera image folder
+ images_all = sorted([f for f in sorted(os.listdir(scene_folder)) if f.lower().endswith('.jpg')])
+ num_images = len(images_all)
+ num_frames = num_images // 6
+
+ out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
+ for frame_id in range(num_frames):
+ front = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5]))
+ front_left = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5 + 1]))
+ front_right = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5 + 2]))
+ left = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5 + 3]))
+ right = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5 + 4]))
+
+ combined_img = tile_cameras({
+ "ring_front_left": front_left,
+ "ring_front_center": front,
+ "ring_front_right": front_right,
+ "ring_side_left": left,
+ "ring_side_right": right
+ })
+ combined_img = cv2.resize(combined_img, (int(frame_width), int(frame_height)))
+
+ out.write(combined_img)
+ out.release()
+
+def preview_dynamic_masks(
+ input_root: str = 'data/argoverse/processed/training',
+ output_root: str = 'data/argoverse/preview/dynamic_masks',
+ frame_width: int = int(5646 / 4),
+ frame_height: int = int(3100 / 4),
+ fps: int = 24,
+ dynamic_class: str = 'all' # 'all' or 'vehicle' or 'human'
+):
+ if not os.path.exists(output_root):
+ os.makedirs(output_root)
+
+ # Iterate over each scene directory in the input root directory
+ for scene_id in tqdm(sorted(os.listdir(input_root))):
+ scene_folder = os.path.join(input_root, scene_id, "dynamic_masks", dynamic_class)
+
+ # Define the output video file path
+ output_video_path = os.path.join(output_root, f'{scene_id}.mp4')
+
+ # Get paths to each camera image folder
+ images_all = sorted([f for f in sorted(os.listdir(scene_folder)) if f.lower().endswith('.png')])
+ num_images = len(images_all)
+ num_frames = num_images // 6
+
+ out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
+ for frame_id in range(num_frames):
+ front = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5]))
+ front_left = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5 + 1]))
+ front_right = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5 + 2]))
+ left = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5 + 3]))
+ right = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5 + 4]))
+
+ combined_img = tile_cameras({
+ "ring_front_left": front_left,
+ "ring_front_center": front,
+ "ring_front_right": front_right,
+ "ring_side_left": left,
+ "ring_side_right": right
+ })
+ combined_img = cv2.resize(combined_img, (int(frame_width), int(frame_height)))
+
+ out.write(combined_img)
+ out.release()
+
+def preview_instances(
+ input_root: str = 'data/argoverse/processed/training',
+ output_root: str = 'data/argoverse/preview/instances',
+ frame_width: int = int(5646 / 4),
+ frame_height: int = int(3100 / 4),
+ fps: int = 24,
+):
+ if not os.path.exists(output_root):
+ os.makedirs(output_root)
+
+ # Iterate over each scene directory in the input root directory
+ for scene_id in tqdm(sorted(os.listdir(input_root))):
+ scene_folder = os.path.join(input_root, scene_id, "instances", "debug_vis")
+
+ # Define the output video file path
+ output_video_path = os.path.join(output_root, f'{scene_id}.mp4')
+
+ # Get paths to each camera image folder
+ images_all = sorted([f for f in sorted(os.listdir(scene_folder)) if f.lower().endswith('.jpg')])
+ num_images = len(images_all)
+ num_frames = num_images // 6
+
+ out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
+ for frame_id in range(num_frames):
+ front = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5]))
+ front_left = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5 + 1]))
+ front_right = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5 + 2]))
+ left = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5 + 3]))
+ right = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 5 + 4]))
+
+ combined_img = tile_cameras({
+ "ring_front_left": front_left,
+ "ring_front_center": front,
+ "ring_front_right": front_right,
+ "ring_side_left": left,
+ "ring_side_right": right
+ })
+ combined_img = cv2.resize(combined_img, (int(frame_width), int(frame_height)))
+
+ out.write(combined_img)
+ out.release()
+
+def cat_all_videos(
+ keys: list,
+ output_root: str = 'data/argoverse/preview/catted_videos',
+ scene_id: str = 'all',
+ fps: int = 24
+):
+ if not os.path.exists(output_root):
+ os.makedirs(output_root)
+
+ if scene_id == 'all':
+ scene_ids = sorted(os.listdir('data/argoverse/processed/training'))
+ else:
+ scene_ids = [scene_id]
+
+ for scene_id in scene_ids:
+ video_paths = []
+ for key in keys:
+ video_paths.append(f'data/argoverse/preview/{key}/{scene_id}.mp4')
+
+ output_video_path = os.path.join(output_root, f'{scene_id}.mp4')
+
+ # use moviepy to concatenate videos
+ clips = [VideoFileClip(path) for path in video_paths]
+ # cat videos from top to bottom, vertically
+ final_clip = clips_array([[clip] for clip in clips])
+ final_clip.write_videofile(output_video_path, fps=fps)
+ final_clip.close()
+
+if __name__ == '__main__':
+ preview_camera_captures(fps=24)
+ preview_3dbox_projection(fps=24)
+ preview_dynamic_masks(fps=24)
+ preview_instances(fps=24)
+ cat_all_videos(keys=['camera_captures', '3dbox_vis', 'dynamic_masks', 'instances'], scene_id='all', fps=24)
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/argoverse/argoverse_sourceloader.py b/worldbench/third_party/drivestudio/datasets/argoverse/argoverse_sourceloader.py
new file mode 100644
index 0000000..581b36e
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/argoverse/argoverse_sourceloader.py
@@ -0,0 +1,524 @@
+import json
+import logging
+import os
+from typing import Dict
+
+import joblib
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from pytorch3d.transforms import matrix_to_quaternion
+from torch import Tensor
+from tqdm import tqdm, trange
+
+from datasets.base.scene_dataset import ModelType
+from datasets.base.lidar_source import SceneLidarSource
+from datasets.base.pixel_source import CameraData, ScenePixelSource
+
+logger = logging.getLogger()
+
+# define each class's node type
+OBJECT_CLASS_NODE_MAPPING = {
+ # Rigid objects (vehicles, infrastructure)
+ "ARTICULATED_BUS": ModelType.RigidNodes,
+ "BOLLARD": ModelType.RigidNodes,
+ "BOX_TRUCK": ModelType.RigidNodes,
+ "BUS": ModelType.RigidNodes,
+ "CONSTRUCTION_BARREL": ModelType.RigidNodes,
+ "CONSTRUCTION_CONE": ModelType.RigidNodes,
+ "LARGE_VEHICLE": ModelType.RigidNodes,
+ "MESSAGE_BOARD_TRAILER": ModelType.RigidNodes,
+ "MOBILE_PEDESTRIAN_CROSSING_SIGN": ModelType.RigidNodes,
+ "MOTORCYCLE": ModelType.RigidNodes,
+ "OFFICIAL_SIGNALER": ModelType.SMPLNodes,
+ "RAILED_VEHICLE": ModelType.RigidNodes,
+ "REGULAR_VEHICLE": ModelType.RigidNodes,
+ "SCHOOL_BUS": ModelType.RigidNodes,
+ "SIGN": ModelType.RigidNodes,
+ "STOP_SIGN": ModelType.RigidNodes,
+ "TRAFFIC_LIGHT_TRAILER": ModelType.RigidNodes,
+ "TRUCK": ModelType.RigidNodes,
+ "TRUCK_CAB": ModelType.RigidNodes,
+ "VEHICULAR_TRAILER": ModelType.RigidNodes,
+
+ # Humans
+ "PEDESTRIAN": ModelType.SMPLNodes,
+
+ # Deformable objects (animals, cyclists, etc.)
+ "ANIMAL": ModelType.DeformableNodes,
+ "BICYCLE": ModelType.DeformableNodes,
+ "BICYCLIST": ModelType.DeformableNodes,
+ "DOG": ModelType.DeformableNodes,
+ "MOTORCYCLIST": ModelType.DeformableNodes,
+ "STROLLER": ModelType.DeformableNodes,
+ "WHEELCHAIR": ModelType.DeformableNodes,
+ "WHEELED_DEVICE": ModelType.DeformableNodes,
+ "WHEELED_RIDER": ModelType.DeformableNodes
+}
+SMPLNODE_CLASSES = ["PEDESTRIAN"]
+
+# OpenCV to Dataset coordinate transformation
+# opencv coordinate system: x right, y down, z front
+OPENCV2DATASET = np.array(
+ [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]
+)
+
+# ArgoVerse Camera List:
+# 0: ring_front_center
+# 1: ring_front_left
+# 2: ring_front_right
+# 3: ring_side_left
+# 4: ring_side_right
+# 5: ring_rear_left
+# 6: ring_rear_right
+AVAILABLE_CAM_LIST = [0, 1, 2, 3, 4, 5, 6]
+
+class ArgoVerseCameraData(CameraData):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ def load_calibrations(self):
+ """
+ Load the camera intrinsics, extrinsics, timestamps, etc.
+ Compute the camera-to-world matrices, ego-to-world matrices, etc.
+ """
+ # load camera intrinsics
+ # 1d Array of [f_u, f_v, c_u, c_v, k{1, 2}, p{1, 2}, k{3}].
+ # ====!! we did not use distortion parameters for simplicity !!====
+ # to be improved!!
+ intrinsic = np.loadtxt(
+ os.path.join(self.data_path, "intrinsics", f"{self.cam_id}.txt")
+ )
+ fx, fy, cx, cy = intrinsic[0], intrinsic[1], intrinsic[2], intrinsic[3]
+ k1, k2, p1, p2, k3 = intrinsic[4], intrinsic[5], intrinsic[6], intrinsic[7], intrinsic[8]
+ # scale intrinsics w.r.t. load size
+ fx, fy = (
+ fx * self.load_size[1] / self.original_size[1],
+ fy * self.load_size[0] / self.original_size[0],
+ )
+ cx, cy = (
+ cx * self.load_size[1] / self.original_size[1],
+ cy * self.load_size[0] / self.original_size[0],
+ )
+ _intrinsics = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
+ _distortions = np.array([k1, k2, p1, p2, k3])
+
+ # load camera extrinsics
+ cam_to_ego = np.loadtxt(
+ os.path.join(self.data_path, "extrinsics", f"{self.cam_id}.txt")
+ )
+ # because we use opencv coordinate system to generate camera rays,
+ # opencv coordinate system: x right, y down, z front
+ # argoverse coordinate system: x right, y down, z front
+ cam_to_ego = cam_to_ego @ OPENCV2DATASET
+
+ # compute per-image poses and intrinsics
+ cam_to_worlds, ego_to_worlds = [], []
+ intrinsics, distortions = [], []
+
+ # we tranform the camera poses w.r.t. the first timestep to make the translation vector of
+ # the first ego pose as the origin of the world coordinate system.
+ ego_to_world_start = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{self.start_timestep:03d}.txt")
+ )
+ for t in range(self.start_timestep, self.end_timestep):
+ ego_to_world_current = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{t:03d}.txt")
+ )
+ # compute ego_to_world transformation
+ ego_to_world = np.linalg.inv(ego_to_world_start) @ ego_to_world_current
+ ego_to_worlds.append(ego_to_world)
+ cam2world = ego_to_world @ cam_to_ego
+ cam_to_worlds.append(cam2world)
+ intrinsics.append(_intrinsics)
+ distortions.append(_distortions)
+
+ self.intrinsics = torch.from_numpy(np.stack(intrinsics, axis=0)).float()
+ self.distortions = torch.from_numpy(np.stack(distortions, axis=0)).float()
+ self.cam_to_worlds = torch.from_numpy(np.stack(cam_to_worlds, axis=0)).float()
+
+ @classmethod
+ def get_camera2worlds(cls, data_path: str, cam_id: str, start_timestep: int, end_timestep: int) -> torch.Tensor:
+ """
+ Returns camera-to-world matrices for the specified camera and time range.
+
+ Args:
+ data_path (str): Path to the dataset.
+ cam_id (str): Camera ID.
+ start_timestep (int): Start timestep.
+ end_timestep (int): End timestep.
+
+ Returns:
+ torch.Tensor: Camera-to-world matrices of shape (num_frames, 4, 4).
+ """
+ # Load the ego pose for the start timestep to use as reference
+ ego_to_world_start = np.loadtxt(
+ os.path.join(data_path, "ego_pose", f"{start_timestep:03d}.txt")
+ )
+
+ # Load camera extrinsics (camera-to-ego transformation)
+ cam_to_ego = np.loadtxt(
+ os.path.join(data_path, "extrinsics", f"{cam_id}.txt")
+ )
+ cam_to_ego = cam_to_ego @ OPENCV2DATASET
+
+ cam_to_worlds = []
+
+ for t in range(start_timestep, end_timestep):
+ # Load ego pose for each timestep
+ ego_to_world_current = np.loadtxt(
+ os.path.join(data_path, "ego_pose", f"{t:03d}.txt")
+ )
+
+ # Compute ego_to_world transformation
+ ego_to_world = np.linalg.inv(ego_to_world_start) @ ego_to_world_current
+
+ # Compute camera-to-world transformation
+ cam2world = ego_to_world @ cam_to_ego
+
+ cam_to_worlds.append(cam2world)
+
+ return torch.from_numpy(np.stack(cam_to_worlds, axis=0)).float()
+
+class ArgoVersePixelSource(ScenePixelSource):
+ def __init__(
+ self,
+ dataset_name: str,
+ pixel_data_config: OmegaConf,
+ data_path: str,
+ start_timestep: int,
+ end_timestep: int,
+ device: torch.device = torch.device("cpu"),
+ ):
+ super().__init__(dataset_name, pixel_data_config, device=device)
+ self.data_path = data_path
+ self.start_timestep = start_timestep
+ self.end_timestep = end_timestep
+ self.load_data()
+
+ def load_cameras(self):
+ self._timesteps = torch.arange(self.start_timestep, self.end_timestep)
+ self.register_normalized_timestamps()
+
+ for idx, cam_id in enumerate(self.camera_list):
+ logger.info(f"Loading camera {cam_id}")
+ camera = ArgoVerseCameraData(
+ dataset_name=self.dataset_name,
+ data_path=self.data_path,
+ cam_id=cam_id,
+ start_timestep=self.start_timestep,
+ end_timestep=self.end_timestep,
+ load_dynamic_mask=self.data_cfg.load_dynamic_mask,
+ load_sky_mask=self.data_cfg.load_sky_mask,
+ downscale_when_loading=self.data_cfg.downscale_when_loading[idx],
+ undistort=self.data_cfg.undistort,
+ buffer_downscale=self.buffer_downscale,
+ device=self.device,
+ )
+ camera.load_time(self.normalized_time)
+ unique_img_idx = torch.arange(len(camera), device=self.device) * len(self.camera_list) + idx
+ camera.set_unique_ids(
+ unique_cam_idx = idx,
+ unique_img_idx = unique_img_idx
+ )
+ logger.info(f"Camera {camera.cam_name} loaded.")
+ self.camera_data[cam_id] = camera
+
+ def load_objects(self):
+ """
+ get ground truth bounding boxes of the dynamic objects
+
+ instances_info = {
+ "0": # simplified instance id
+ {
+ "id": str,
+ "class_name": str,
+ "frame_annotations": {
+ "frame_idx": List,
+ "obj_to_world": List,
+ "box_size": List,
+ },
+ ...
+ }
+ frame_instances = {
+ "0": # frame idx
+ List[int] # list of simplified instance ids
+ ...
+ }
+ """
+ instances_info_path = os.path.join(self.data_path, "instances", "instances_info.json")
+ frame_instances_path = os.path.join(self.data_path, "instances", "frame_instances.json")
+ with open(instances_info_path, "r") as f:
+ instances_info = json.load(f)
+ with open(frame_instances_path, "r") as f:
+ frame_instances = json.load(f)
+ # get pose of each instance at each frame
+ # shape (num_frames, num_instances, 4, 4)
+ num_instances = len(instances_info)
+ num_full_frames = len(frame_instances)
+ instances_pose = np.zeros((num_full_frames, num_instances, 4, 4))
+ instances_size = np.zeros((num_full_frames, num_instances, 3))
+ instances_true_id = np.arange(num_instances)
+ instances_model_types = np.ones(num_instances) * -1
+
+ ego_to_world_start = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{self.start_timestep:03d}.txt")
+ )
+ for k, v in instances_info.items():
+ instances_model_types[int(k)] = OBJECT_CLASS_NODE_MAPPING[v["class_name"]]
+ for frame_idx, obj_to_world, box_size in zip(v["frame_annotations"]["frame_idx"], v["frame_annotations"]["obj_to_world"], v["frame_annotations"]["box_size"]):
+ # the first ego pose as the origin of the world coordinate system.
+ obj_to_world = np.array(obj_to_world).reshape(4, 4)
+ obj_to_world = np.linalg.inv(ego_to_world_start) @ obj_to_world
+ instances_pose[frame_idx, int(k)] = np.array(obj_to_world)
+ instances_size[frame_idx, int(k)] = np.array(box_size)
+
+ # get frame valid instances
+ # shape (num_frames, num_instances)
+ per_frame_instance_mask = np.zeros((num_full_frames, num_instances))
+ for frame_idx, valid_instances in frame_instances.items():
+ per_frame_instance_mask[int(frame_idx), valid_instances] = 1
+
+ # select the frames that are in the range of start_timestep and end_timestep
+ instances_pose = torch.from_numpy(instances_pose[self.start_timestep:self.end_timestep]).float()
+ instances_size = torch.from_numpy(instances_size[self.start_timestep:self.end_timestep]).float()
+ instances_true_id = torch.from_numpy(instances_true_id).long()
+ instances_model_types = torch.from_numpy(instances_model_types).long()
+ per_frame_instance_mask = torch.from_numpy(per_frame_instance_mask[self.start_timestep:self.end_timestep]).bool()
+
+ # filter out the instances that are not visible in selected frames
+ ins_frame_cnt = per_frame_instance_mask.sum(dim=0)
+ instances_pose = instances_pose[:, ins_frame_cnt > 0]
+ instances_size = instances_size[:, ins_frame_cnt > 0]
+ instances_true_id = instances_true_id[ins_frame_cnt > 0]
+ instances_model_types = instances_model_types[ins_frame_cnt > 0]
+ per_frame_instance_mask = per_frame_instance_mask[:, ins_frame_cnt > 0]
+
+ # assign to the class
+ # (num_frames, num_instances, 4, 4)
+ self.instances_pose = instances_pose
+ # (num_instances, 3)
+ self.instances_size = instances_size.sum(0) / per_frame_instance_mask.sum(0).unsqueeze(-1)
+ # (num_frames, num_instances)
+ self.per_frame_instance_mask = per_frame_instance_mask
+ # (num_instances)
+ self.instances_true_id = instances_true_id
+ # (num_instances)
+ self.instances_model_types = instances_model_types
+
+ if self.data_cfg.load_smpl:
+ # Collect camera-to-world matrices for all available cameras
+ cam_to_worlds = {}
+ for cam_id in AVAILABLE_CAM_LIST:
+ cam_to_worlds[cam_id] = ArgoVerseCameraData.get_camera2worlds(
+ self.data_path,
+ str(cam_id),
+ self.start_timestep,
+ self.end_timestep
+ )
+
+ # load SMPL parameters
+ smpl_dict = joblib.load(os.path.join(self.data_path, "humanpose", "smpl.pkl"))
+ frame_num = self.end_timestep - self.start_timestep
+
+ smpl_human_all = {}
+ for fi in tqdm(range(self.start_timestep, self.end_timestep), desc="Loading SMPL"):
+ for instance_id, ins_smpl in smpl_dict.items():
+ if instance_id not in smpl_human_all:
+ smpl_human_all[instance_id] = {
+ "smpl_quats": torch.zeros((frame_num, 24, 4), dtype=torch.float32),
+ "smpl_trans": torch.zeros((frame_num, 3), dtype=torch.float32),
+ "smpl_betas": torch.zeros((frame_num, 10), dtype=torch.float32),
+ "frame_valid": torch.zeros((frame_num), dtype=torch.bool)
+ }
+ smpl_human_all[instance_id]["smpl_quats"][:, :, 0] = 1.0
+ if ins_smpl["valid_mask"][fi]:
+ betas = ins_smpl["smpl"]["betas"][fi]
+ smpl_human_all[instance_id]["smpl_betas"][fi - self.start_timestep] = betas
+
+ body_pose = ins_smpl["smpl"]["body_pose"][fi]
+ smpl_orient = ins_smpl["smpl"]["global_orient"][fi]
+ cam_depend = ins_smpl["selected_cam_idx"][fi].item()
+
+ c2w = cam_to_worlds[cam_depend][fi - self.start_timestep]
+ world_orient = c2w[:3, :3].to(smpl_orient.device) @ smpl_orient.squeeze()
+ smpl_quats = matrix_to_quaternion(
+ torch.cat([world_orient[None, ...], body_pose], dim=0)
+ )
+
+ ii = instances_info[str(instance_id)]['frame_annotations']["frame_idx"].index(fi)
+ o2w = np.array(
+ instances_info[str(instance_id)]['frame_annotations']["obj_to_world"][ii]
+ )
+ o2w = torch.from_numpy(
+ np.linalg.inv(ego_to_world_start) @ o2w
+ )
+ # box_size = instances_info[str(instance_id)]['frame_annotations']["box_size"][ii]
+
+ smpl_human_all[instance_id]["smpl_quats"][fi - self.start_timestep] = smpl_quats
+ smpl_human_all[instance_id]["smpl_trans"][fi - self.start_timestep] = o2w[:3, 3]
+ smpl_human_all[instance_id]["frame_valid"][fi - self.start_timestep] = True
+
+ self.smpl_human_all = smpl_human_all
+
+class ArgoVerseLiDARSource(SceneLidarSource):
+ def __init__(
+ self,
+ lidar_data_config: OmegaConf,
+ data_path: str,
+ start_timestep: int,
+ end_timestep: int,
+ device: torch.device = torch.device("cpu"),
+ ):
+ super().__init__(lidar_data_config, device=device)
+ self.data_path = data_path
+ self.start_timestep = start_timestep
+ self.end_timestep = end_timestep
+ self.create_all_filelist()
+ self.load_data()
+
+ def create_all_filelist(self):
+ """
+ Create a list of all the files in the dataset.
+ e.g., a list of all the lidar scans in the dataset.
+ """
+ lidar_filepaths = []
+ for t in range(self.start_timestep, self.end_timestep):
+ lidar_filepaths.append(
+ os.path.join(self.data_path, "lidar", f"{t:03d}.bin")
+ )
+ self.lidar_filepaths = np.array(lidar_filepaths)
+
+ def load_calibrations(self):
+ """
+ Load the calibration files of the dataset.
+ e.g., lidar to world transformation matrices.
+ """
+ # Note that in the Waymo Open Dataset, the lidar coordinate system is the same
+ # as the vehicle coordinate system
+ lidar_to_worlds = []
+
+ # we tranform the poses w.r.t. the first timestep to make the origin of the
+ # first ego pose as the origin of the world coordinate system.
+ ego_to_world_start = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{self.start_timestep:03d}.txt")
+ )
+ for t in range(self.start_timestep, self.end_timestep):
+ ego_to_world_current = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{t:03d}.txt")
+ )
+ # compute ego_to_world transformation
+ lidar_to_world = np.linalg.inv(ego_to_world_start) @ ego_to_world_current
+ lidar_to_worlds.append(lidar_to_world)
+
+ self.lidar_to_worlds = torch.from_numpy(
+ np.stack(lidar_to_worlds, axis=0)
+ ).float()
+
+ def load_lidar(self):
+ """
+ Load the lidar data of the dataset from the filelist.
+ """
+ origins, directions, ranges = [], [], []
+ # we simplify timestamps as the time indices
+ timesteps = []
+
+ accumulated_num_original_rays = 0
+ accumulated_num_rays = 0
+ for t in trange(
+ 0, len(self.lidar_filepaths), desc="Loading lidar", dynamic_ncols=True
+ ):
+ # each lidar_info contains an Nx4 array
+ # from left to right:
+ # points: 3d, intensity: 1d
+ lidar_info = np.memmap(
+ self.lidar_filepaths[t],
+ dtype=np.float32,
+ mode="r",
+ ).reshape(-1, 4)
+ original_length = len(lidar_info)
+ accumulated_num_original_rays += original_length
+
+ lidar_points = torch.from_numpy(lidar_info[:, :3]).float()
+ # convert the lidar points to opencv coordinate system
+ # since lidar points already include the information of origins
+ # NOTE: here we use the pesudo origins, TODO: consider a more accurate way
+ lidar_origins = torch.zeros_like(lidar_points)
+
+ # transform lidar points from lidar coordinate system to world coordinate system
+ lidar_origins = (
+ self.lidar_to_worlds[t][:3, :3] @ lidar_origins.T
+ + self.lidar_to_worlds[t][:3, 3:4]
+ ).T
+ lidar_points = (
+ self.lidar_to_worlds[t][:3, :3] @ lidar_points.T
+ + self.lidar_to_worlds[t][:3, 3:4]
+ ).T
+ # compute lidar directions
+ lidar_directions = lidar_points - lidar_origins
+ lidar_ranges = torch.norm(lidar_directions, dim=-1, keepdim=True)
+ lidar_directions = lidar_directions / lidar_ranges
+ # we use time indices as the timestamp
+ lidar_timestamp = torch.ones_like(lidar_ranges).squeeze(-1) * t
+ accumulated_num_rays += len(lidar_ranges)
+
+ origins.append(lidar_origins)
+ directions.append(lidar_directions)
+ ranges.append(lidar_ranges)
+ # we use time indices as the timestamp
+ timesteps.append(lidar_timestamp)
+
+ logger.info(
+ f"Number of lidar rays: {accumulated_num_rays} "
+ f"({accumulated_num_rays / accumulated_num_original_rays * 100:.2f}% of "
+ f"{accumulated_num_original_rays} original rays)"
+ )
+
+ self.origins = torch.cat(origins, dim=0)
+ self.directions = torch.cat(directions, dim=0)
+ self.ranges = torch.cat(ranges, dim=0)
+ self.visible_masks = torch.zeros_like(self.ranges).squeeze().bool()
+ self.colors = torch.ones_like(self.directions)
+
+ # the underscore here is important.
+ self._timesteps = torch.cat(timesteps, dim=0)
+ self.register_normalized_timestamps()
+
+ def get_lidar_rays(self, time_idx: int) -> Dict[str, Tensor]:
+ """
+ Get the of rays for rendering at the given timestep.
+ Args:
+ time_idx: the index of the lidar scan to render.
+ Returns:
+ a dict of the sampled rays.
+ """
+ origins = self.origins[self.timesteps == time_idx]
+ directions = self.directions[self.timesteps == time_idx]
+ ranges = self.ranges[self.timesteps == time_idx]
+ normalized_time = self.normalized_time[self.timesteps == time_idx]
+ return {
+ "lidar_origins": origins,
+ "lidar_viewdirs": directions,
+ "lidar_ranges": ranges,
+ "lidar_normed_time": normalized_time,
+ "lidar_mask": self.timesteps == time_idx,
+ }
+
+ def delete_invisible_pts(self) -> None:
+ """
+ Clear the unvisible points.
+ """
+ if self.visible_masks is not None:
+ num_bf = self.origins.shape[0]
+ self.origins = self.origins[self.visible_masks]
+ self.directions = self.directions[self.visible_masks]
+ self.ranges = self.ranges[self.visible_masks]
+ self._timesteps = self._timesteps[self.visible_masks]
+ self._normalized_time = self._normalized_time[self.visible_masks]
+ self.colors = self.colors[self.visible_masks]
+ logger.info(
+ f"[Lidar] {num_bf - self.visible_masks.sum()} out of {num_bf} points are cleared. {self.visible_masks.sum()} points left."
+ )
+ self.visible_masks = None
+ else:
+ logger.info("[Lidar] No unvisible points to clear.")
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/base/__init__.py b/worldbench/third_party/drivestudio/datasets/base/__init__.py
new file mode 100644
index 0000000..07fa987
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/base/__init__.py
@@ -0,0 +1,4 @@
+from .lidar_source import SceneLidarSource
+from .pixel_source import ScenePixelSource
+from .scene_dataset import SceneDataset
+from .split_wrapper import SplitWrapper
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/base/lidar_source.py b/worldbench/third_party/drivestudio/datasets/base/lidar_source.py
new file mode 100644
index 0000000..4cbc134
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/base/lidar_source.py
@@ -0,0 +1,260 @@
+import abc
+import logging
+from typing import Dict
+
+import torch
+import torch.nn.functional as F
+from omegaconf import OmegaConf
+from torch import Tensor
+
+logger = logging.getLogger()
+
+
+class SceneLidarSource(abc.ABC):
+ """
+ The base class for the lidar source of a scene.
+ """
+
+ data_cfg: OmegaConf = None
+ # the normalized time of all images (normalized to [0, 1]), shape: (num_frames,)
+ _normalized_time: Tensor = None
+ # timestep indices of frames, shape: (num_frames,)
+ _timesteps: Tensor = None
+ # origin of each lidar point, shape: (num_points, 3)
+ origins: Tensor = None
+ # unit direction of each lidar point, shape: (num_points, 3)
+ directions: Tensor = None
+ # range of each lidar point, shape: (num_points,)
+ ranges: Tensor = None
+ # the transformation matrices from lidar to world coordinate system,
+ lidar_to_worlds: Tensor = None
+ # indicate whether each lidar point is visible to the camera,
+ visible_masks: Tensor = None
+ # the color of each lidar point, shape: (num_points, 3)
+ colors: Tensor = None
+ # the indices of the lidar scans that are cached
+ cached_indices: Tensor = None
+ cached_origins: Tensor = None
+ cached_directions: Tensor = None
+ cached_ranges: Tensor = None
+ cached_normalized_timestamps: Tensor = None
+
+ def __init__(
+ self,
+ lidar_data_config: OmegaConf,
+ device: torch.device = torch.device("cpu"),
+ ) -> None:
+ # hold the config of the lidar data
+ self.data_cfg = lidar_data_config
+ self.device = device
+
+ @abc.abstractmethod
+ def create_all_filelist(self) -> None:
+ """
+ Create a list of all the files in the dataset.
+ e.g., a list of all the lidar scans in the dataset.
+ """
+ raise NotImplementedError
+
+ def load_data(self):
+ self.load_calibrations()
+ self.load_lidar()
+ logger.info("[Lidar] All Lidar Data loaded.")
+
+ def to(self, device: torch.device) -> "SceneLidarSource":
+ """
+ Move the dataset to the given device.
+ Args:
+ device: the device to move the dataset to.
+ """
+ self.device = device
+ if self.origins is not None:
+ self.origins = self.origins.to(device)
+ if self.directions is not None:
+ self.directions = self.directions.to(device)
+ if self.ranges is not None:
+ self.ranges = self.ranges.to(device)
+ if self._timesteps is not None:
+ self._timesteps = self._timesteps.to(device)
+ if self._normalized_time is not None:
+ self._normalized_time = self._normalized_time.to(device)
+ if self.lidar_to_worlds is not None:
+ self.lidar_to_worlds = self.lidar_to_worlds.to(device)
+ if self.visible_masks is not None:
+ self.visible_masks = self.visible_masks.to(device)
+ if self.colors is not None:
+ self.colors = self.colors.to(device)
+ return self
+
+ @abc.abstractmethod
+ def load_calibrations(self) -> None:
+ """
+ Load the calibration files of the dataset.
+ e.g., lidar to world transformation matrices.
+ """
+ raise NotImplementedError
+
+ @abc.abstractmethod
+ def load_lidar(self) -> None:
+ """
+ Load the lidar data of the dataset from the filelist.
+ """
+ raise NotImplementedError
+
+ def get_aabb(self) -> Tensor:
+ """
+ Returns:
+ aabb_min, aabb_max: the min and max of the axis-aligned bounding box of the scene
+ Note:
+ we assume the lidar points are already in the world coordinate system
+ we first downsample the lidar points, then compute the aabb by taking the
+ given percentiles of the lidar coordinates in each dimension.
+ """
+ assert (
+ self.origins is not None
+ and self.directions is not None
+ and self.ranges is not None
+ ), "Lidar points not loaded, cannot compute aabb."
+ logger.info("[Lidar] Computing auto AABB based on downsampled lidar points....")
+
+ lidar_pts = self.origins + self.directions * self.ranges
+
+ # downsample the lidar points by uniformly sampling a subset of them
+ lidar_pts = lidar_pts[
+ torch.randperm(len(lidar_pts))[
+ : int(len(lidar_pts) / self.data_cfg.lidar_downsample_factor)
+ ]
+ ]
+ # compute the aabb by taking the given percentiles of the lidar coordinates in each dimension
+ aabb_min = torch.quantile(lidar_pts, self.data_cfg.lidar_percentile, dim=0)
+ aabb_max = torch.quantile(lidar_pts, 1 - self.data_cfg.lidar_percentile, dim=0)
+ del lidar_pts
+ torch.cuda.empty_cache()
+
+ # usually the lidar's height is very small, so we slightly increase the height of the aabb
+ if aabb_max[-1] < 20:
+ aabb_max[-1] = 20.0
+ aabb = torch.tensor([*aabb_min, *aabb_max])
+ logger.info(f"[Lidar] Auto AABB from LiDAR: {aabb}")
+ return aabb
+
+ @property
+ def pts_xyz(self) -> Tensor:
+ """
+ Returns:
+ the xyz coordinates of the lidar points.
+ shape: (num_lidar_points, 3)
+ """
+ return self.origins + self.directions * self.ranges
+
+ @property
+ def num_points(self) -> int:
+ """
+ Returns:
+ the number of lidar points in the dataset.
+ """
+ return self.origins.size(0)
+
+ @property
+ def num_timesteps(self) -> int:
+ """
+ Returns:
+ the number of lidar timestamps in the dataset,
+ usually the number of captured lidar scans.
+ """
+ return len(self.timesteps.unique())
+
+ @property
+ def timesteps(self) -> Tensor:
+ """
+ Returns:
+ the integer timestep indices of each lidar timestamp,
+ shape: (num_lidar_points,)
+ Note:
+ the difference between timestamps and timesteps is that
+ timestamps are the actual timestamps (minus 1e9) of the lidar scans,
+ while timesteps are the integer timestep indices of the lidar scans.
+ """
+ return self._timesteps
+
+ @property
+ def normalized_time(self) -> Tensor:
+ """
+ Returns:
+ the normalized timestamps of the lidar scans
+ (normalized to the range [0, 1]).
+ shape: (num_lidar_points,)
+ """
+ return self._normalized_time
+
+ @property
+ def unique_normalized_timestamps(self) -> Tensor:
+ """
+ Returns:
+ the unique normalized timestamps of the lidar scans
+ (normalized to the range [0, 1]).
+ shape: (num_timesteps,)
+ """
+ return self._unique_normalized_timestamps
+
+ def register_normalized_timestamps(self) -> None:
+ # normalized timestamps are between 0 and 1
+ normalized_time = (self._timesteps - self._timesteps.min()) / (
+ self._timesteps.max() - self._timesteps.min()
+ )
+ self._normalized_time = normalized_time.to(self.device)
+ self._unique_normalized_timestamps = self._normalized_time.unique()
+
+ def find_closest_timestep(self, normed_timestamp: float) -> int:
+ """
+ Find the closest timestep to the given timestamp.
+ Args:
+ normed_timestamp: the normalized timestamp to find the closest timestep for.
+ Returns:
+ the closest timestep to the given timestamp.
+ """
+ return torch.argmin(
+ torch.abs(self.unique_normalized_timestamps - normed_timestamp)
+ )
+
+ def get_lidar_rays(self, time_idx: int) -> Dict[str, Tensor]:
+ """
+ Get the of rays for rendering at the given timestep.
+ Args:
+ time_idx: the index of the lidar scan to render.
+ Returns:
+ a dict of the sampled rays.
+ """
+ origins = self.origins[self.timesteps == time_idx]
+ directions = self.directions[self.timesteps == time_idx]
+ ranges = self.ranges[self.timesteps == time_idx]
+ normalized_time = self.normalized_time[self.timesteps == time_idx]
+ flows = self.flows[self.timesteps == time_idx]
+ return {
+ "lidar_origins": origins,
+ "lidar_viewdirs": directions,
+ "lidar_ranges": ranges,
+ "lidar_normed_time": normalized_time,
+ "lidar_mask": self.timesteps == time_idx,
+ "lidar_flows": flows,
+ }
+
+ def delete_invisible_pts(self) -> None:
+ """
+ Clear the unvisible points.
+ """
+ if self.visible_masks is not None:
+ num_bf = self.origins.shape[0]
+ self.origins = self.origins[self.visible_masks]
+ self.directions = self.directions[self.visible_masks]
+ self.ranges = self.ranges[self.visible_masks]
+ self.flows = self.flows[self.visible_masks]
+ self._timesteps = self._timesteps[self.visible_masks]
+ self._normalized_time = self._normalized_time[self.visible_masks]
+ self.colors = self.colors[self.visible_masks]
+ logger.info(
+ f"[Lidar] {num_bf - self.visible_masks.sum()} out of {num_bf} points are cleared. {self.visible_masks.sum()} points left."
+ )
+ self.visible_masks = None
+ else:
+ logger.info("[Lidar] No unvisible points to clear.")
diff --git a/worldbench/third_party/drivestudio/datasets/base/pixel_source.py b/worldbench/third_party/drivestudio/datasets/base/pixel_source.py
new file mode 100644
index 0000000..1237f24
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/base/pixel_source.py
@@ -0,0 +1,1139 @@
+from typing import Dict, Tuple, List, Callable
+from omegaconf import OmegaConf
+import os
+import abc
+import cv2
+import random
+import logging
+import numpy as np
+from tqdm import tqdm
+from PIL import Image
+
+import torch
+from torch import Tensor
+from datasets.dataset_meta import DATASETS_CONFIG
+
+logger = logging.getLogger()
+
+def idx_to_3d(idx, H, W):
+ """
+ Converts a 1D index to a 3D index (img_id, row_id, col_id)
+
+ Args:
+ idx (int): The 1D index to convert.
+ H (int): The height of the 3D space.
+ W (int): The width of the 3D space.
+
+ Returns:
+ tuple: A tuple containing the 3D index (i, j, k),
+ where i is the image index, j is the row index,
+ and k is the column index.
+ """
+ i = idx // (H * W)
+ j = (idx % (H * W)) // W
+ k = idx % W
+ return i, j, k
+
+
+def get_rays(
+ x: Tensor, y: Tensor, c2w: Tensor, intrinsic: Tensor
+) -> Tuple[Tensor, Tensor, Tensor]:
+ """
+ Args:
+ x: the horizontal coordinates of the pixels, shape: (num_rays,)
+ y: the vertical coordinates of the pixels, shape: (num_rays,)
+ c2w: the camera-to-world matrices, shape: (num_cams, 4, 4)
+ intrinsic: the camera intrinsic matrices, shape: (num_cams, 3, 3)
+ Returns:
+ origins: the ray origins, shape: (num_rays, 3)
+ viewdirs: the ray directions, shape: (num_rays, 3)
+ direction_norm: the norm of the ray directions, shape: (num_rays, 1)
+ """
+ if len(intrinsic.shape) == 2:
+ intrinsic = intrinsic[None, :, :]
+ if len(c2w.shape) == 2:
+ c2w = c2w[None, :, :]
+ camera_dirs = torch.nn.functional.pad(
+ torch.stack(
+ [
+ (x - intrinsic[:, 0, 2] + 0.5) / intrinsic[:, 0, 0],
+ (y - intrinsic[:, 1, 2] + 0.5) / intrinsic[:, 1, 1],
+ ],
+ dim=-1,
+ ),
+ (0, 1),
+ value=1.0,
+ ) # [num_rays, 3]
+
+ # rotate the camera rays w.r.t. the camera pose
+ directions = (camera_dirs[:, None, :] * c2w[:, :3, :3]).sum(dim=-1)
+ origins = torch.broadcast_to(c2w[:, :3, -1], directions.shape)
+ # TODO: not sure if we still need direction_norm
+ direction_norm = torch.linalg.norm(directions, dim=-1, keepdims=True)
+ # normalize the ray directions
+ viewdirs = directions / (direction_norm + 1e-8)
+ return origins, viewdirs, direction_norm
+
+def sparse_lidar_map_downsampler(lidar_depth_map, downscale_factor):
+ # NOTE(ziyu): Important!
+ # This is the correct way to downsample the sparse lidar depth map.
+ raw_avg = torch.nn.functional.interpolate(
+ lidar_depth_map.unsqueeze(0).unsqueeze(0),
+ scale_factor=downscale_factor,
+ mode="area",
+ ).squeeze(0).squeeze(0)
+ raw_mask = torch.nn.functional.interpolate(
+ (lidar_depth_map > 1e-3).float().unsqueeze(0).unsqueeze(0),
+ scale_factor=downscale_factor,
+ mode="area",
+ ).squeeze(0).squeeze(0)
+ downsampled_lidar_map = torch.zeros_like(raw_avg)
+ downsampled_lidar_map[raw_mask > 0] = raw_avg[raw_mask > 0] / raw_mask[raw_mask > 0]
+ return downsampled_lidar_map
+
+class CameraData(object):
+ def __init__(
+ self,
+ dataset_name: str,
+ data_path: str,
+ cam_id: int,
+ # the start timestep to load
+ start_timestep: int = 0,
+ # the end timestep to load
+ end_timestep: int = None,
+ # whether to load the dynamic masks
+ load_dynamic_mask: bool = False,
+ # whether to load the sky masks
+ load_sky_mask: bool = False,
+ # the size to load the images
+ downscale_when_loading: float = 1.0,
+ # whether to undistort the images
+ undistort: bool = False,
+ # whether to use buffer sampling
+ buffer_downscale: float = 1.0,
+ # the device to move the camera to
+ device: torch.device = torch.device("cpu"),
+ ):
+ self.dataset_name = dataset_name
+ self.cam_id = cam_id
+ self.data_path = data_path
+ self.start_timestep = start_timestep
+ self.end_timestep = end_timestep
+ self.undistort = undistort
+ self.buffer_downscale = buffer_downscale
+ self.device = device
+
+ self.cam_name = DATASETS_CONFIG[dataset_name][cam_id]["camera_name"]
+ self.original_size = DATASETS_CONFIG[dataset_name][cam_id]["original_size"]
+ self.load_size = [
+ int(self.original_size[0] / downscale_when_loading),
+ int(self.original_size[1] / downscale_when_loading),
+ ]
+
+ # Load the images, dynamic masks, sky masks, etc.
+ self.create_all_filelist()
+ self.load_calibrations()
+ self.load_images()
+ self.load_egocar_mask()
+ if load_dynamic_mask:
+ self.load_dynamic_masks()
+ if load_sky_mask:
+ self.load_sky_masks()
+ self.lidar_depth_maps = None # will be loaded by: self.load_depth()
+ self.image_error_maps = None # will be built by: self.build_image_error_buffer()
+ self.to(self.device)
+ self.downscale_factor = 1.0
+
+ @property
+ def num_frames(self) -> int:
+ return self.cam_to_worlds.shape[0]
+
+ @property
+ def HEIGHT(self) -> int:
+ return self.load_size[0]
+
+ @property
+ def WIDTH(self) -> int:
+ return self.load_size[1]
+
+ def __len__(self):
+ return self.num_frames
+
+ def set_downscale_factor(self, downscale_factor: float):
+ self.downscale_factor = downscale_factor
+
+ def set_unique_ids(self, unique_cam_idx: int, unique_img_idx: Tensor):
+ """
+ unique id is the compact order of the camera index and frame index
+ for example camera idx is [0, 2, 4]
+ the camera index is [0, 1, 2]
+ """
+ self.unique_cam_idx = unique_cam_idx
+ self.unique_img_idx = unique_img_idx.to(self.device)
+
+ def load_calibrations(
+ self,
+ unique_img_idx,
+ cam_to_worlds,
+ intrinsics,
+ distortions,
+ ):
+ self.unique_img_idx = unique_img_idx
+ # Load the camera intrinsics, extrinsics, timestamps, etc.
+ # Compute the camera-to-world matrices, ego-to-world matrices, etc.
+ self.cam_to_worlds = cam_to_worlds # (num_frames, 4, 4)
+ self.intrinsics = intrinsics # (num_frames, 3, 3)
+ self.distortions = distortions # (num_frames, 5)
+
+ def create_all_filelist(self):
+ """
+ Create file lists for all data files.
+ e.g., img files, feature files, etc.
+ """
+ # ---- define filepaths ---- #
+ img_filepaths = []
+ dynamic_mask_filepaths, sky_mask_filepaths = [], []
+ human_mask_filepaths, vehicle_mask_filepaths = [], []
+
+ fine_mask_path = os.path.join(self.data_path, "fine_dynamic_masks")
+ if os.path.exists(fine_mask_path):
+ dynamic_mask_dir = "fine_dynamic_masks"
+ logger.info("Using fine dynamic masks")
+ else:
+ dynamic_mask_dir = "dynamic_masks"
+ logger.info("Using coarse dynamic masks")
+
+ # Note: we assume all the files in waymo dataset are synchronized
+ for t in range(self.start_timestep, self.end_timestep):
+ img_filepaths.append(
+ os.path.join(self.data_path, "images", f"{t:03d}_{self.cam_id}.jpg")
+ )
+ dynamic_mask_filepaths.append(
+ os.path.join(
+ self.data_path, dynamic_mask_dir, "all", f"{t:03d}_{self.cam_id}.png"
+ )
+ )
+ human_mask_filepaths.append(
+ os.path.join(
+ self.data_path, dynamic_mask_dir, "human", f"{t:03d}_{self.cam_id}.png"
+ )
+ )
+ vehicle_mask_filepaths.append(
+ os.path.join(
+ self.data_path, dynamic_mask_dir, "vehicle", f"{t:03d}_{self.cam_id}.png"
+ )
+ )
+ sky_mask_filepaths.append(
+ os.path.join(self.data_path, "sky_masks", f"{t:03d}_{self.cam_id}.png")
+ )
+ self.img_filepaths = np.array(img_filepaths)
+ self.dynamic_mask_filepaths = np.array(dynamic_mask_filepaths)
+ self.human_mask_filepaths = np.array(human_mask_filepaths)
+ self.vehicle_mask_filepaths = np.array(vehicle_mask_filepaths)
+ self.sky_mask_filepaths = np.array(sky_mask_filepaths)
+
+ def load_images(self):
+ images = []
+ for ix, fname in tqdm(
+ enumerate(self.img_filepaths),
+ desc="Loading images",
+ dynamic_ncols=True,
+ total=len(self.img_filepaths),
+ ):
+ rgb = Image.open(fname).convert("RGB")
+ # resize them to the load_size
+ rgb = rgb.resize(
+ (self.load_size[1], self.load_size[0]), Image.BILINEAR
+ )
+ # undistort the images
+ if self.undistort:
+ if ix == 0:
+ print("undistorting rgb")
+ rgb = cv2.undistort(
+ np.array(rgb),
+ self.intrinsics[ix].numpy(),
+ self.distortions[ix].numpy(),
+ )
+ images.append(rgb)
+ # normalize the images to [0, 1]
+ self.images = images = torch.from_numpy(np.stack(images, axis=0)) / 255
+
+ def load_egocar_mask(self):
+ """
+ Since in some datasets, the ego car body is visible in the images,
+ we need to load the ego car mask to mask out the ego car body.
+ """
+ egocar_mask = os.path.join("data", "ego_masks", self.dataset_name, f"{self.cam_id}.png")
+ if os.path.exists(egocar_mask):
+ egocar_mask = Image.open(egocar_mask).convert("L")
+ # resize them to the load_size
+ egocar_mask = egocar_mask.resize(
+ (self.load_size[1], self.load_size[0]), Image.BILINEAR
+ )
+ if self.undistort:
+ egocar_mask = cv2.undistort(
+ np.array(egocar_mask),
+ self.intrinsics[0].numpy(),
+ self.distortions[0].numpy(),
+ )
+ self.egocar_mask = torch.from_numpy(np.array(egocar_mask) > 0).float()
+ else:
+ self.egocar_mask = None
+
+ def load_dynamic_masks(self):
+ dynamic_masks = []
+ for ix, fname in tqdm(
+ enumerate(self.dynamic_mask_filepaths),
+ desc="Loading dynamic masks",
+ dynamic_ncols=True,
+ total=len(self.dynamic_mask_filepaths),
+ ):
+ dyn_mask = Image.open(fname).convert("L")
+ # resize them to the load_size
+ dyn_mask = dyn_mask.resize(
+ (self.load_size[1], self.load_size[0]), Image.BILINEAR
+ )
+ if self.undistort:
+ if ix == 0:
+ print("undistorting dynamic mask")
+ dyn_mask = cv2.undistort(
+ np.array(dyn_mask),
+ self.intrinsics[ix].numpy(),
+ self.distortions[ix].numpy(),
+ )
+ dynamic_masks.append(np.array(dyn_mask) > 0)
+ self.dynamic_masks = torch.from_numpy(np.stack(dynamic_masks, axis=0)).float()
+
+ human_masks = []
+ for ix, fname in tqdm(
+ enumerate(self.human_mask_filepaths),
+ desc="Loading human masks",
+ dynamic_ncols=True,
+ total=len(self.human_mask_filepaths),
+ ):
+ human_mask = Image.open(fname).convert("L")
+ # resize them to the load_size
+ human_mask = human_mask.resize(
+ (self.load_size[1], self.load_size[0]), Image.BILINEAR
+ )
+ if self.undistort:
+ if ix == 0:
+ print("undistorting human mask")
+ human_mask = cv2.undistort(
+ np.array(human_mask),
+ self.intrinsics[ix].numpy(),
+ self.distortions[ix].numpy(),
+ )
+ human_masks.append(np.array(human_mask) > 0)
+ self.human_masks = torch.from_numpy(np.stack(human_masks, axis=0)).float()
+
+ vehicle_masks = []
+ for ix, fname in tqdm(
+ enumerate(self.vehicle_mask_filepaths),
+ desc="Loading vehicle masks",
+ dynamic_ncols=True,
+ total=len(self.vehicle_mask_filepaths),
+ ):
+ vehicle_mask = Image.open(fname).convert("L")
+ # resize them to the load_size
+ vehicle_mask = vehicle_mask.resize(
+ (self.load_size[1], self.load_size[0]), Image.BILINEAR
+ )
+ if self.undistort:
+ if ix == 0:
+ print("undistorting vehicle mask")
+ vehicle_mask = cv2.undistort(
+ np.array(vehicle_mask),
+ self.intrinsics[ix].numpy(),
+ self.distortions[ix].numpy(),
+ )
+ vehicle_masks.append(np.array(vehicle_mask) > 0)
+ self.vehicle_masks = torch.from_numpy(np.stack(vehicle_masks, axis=0)).float()
+
+ def load_sky_masks(self):
+ sky_masks = []
+ for ix, fname in tqdm(
+ enumerate(self.sky_mask_filepaths),
+ desc="Loading sky masks",
+ dynamic_ncols=True,
+ total=len(self.sky_mask_filepaths),
+ ):
+ sky_mask = Image.open(fname).convert("L")
+ # resize them to the load_size
+ sky_mask = sky_mask.resize(
+ (self.load_size[1], self.load_size[0]), Image.NEAREST
+ )
+ if self.undistort:
+ if ix == 0:
+ print("undistorting sky mask")
+ sky_mask = cv2.undistort(
+ np.array(sky_mask),
+ self.intrinsics[ix].numpy(),
+ self.distortions[ix].numpy(),
+ )
+ sky_masks.append(np.array(sky_mask) > 0)
+ self.sky_masks = torch.from_numpy(np.stack(sky_masks, axis=0)).float()
+
+ def load_depth(
+ self,
+ lidar_depth_maps: Tensor,
+ ):
+ self.lidar_depth_maps = lidar_depth_maps.to(self.device)
+
+ def load_time(
+ self,
+ normalized_time: Tensor,
+ ):
+ self.normalized_time = normalized_time.to(self.device)
+
+ def build_image_error_buffer(self) -> None:
+ """
+ Build the image error buffer.
+ """
+ # Tensor (num_frames, H // buffer_downscale, W // buffer_downscale)
+ self.image_error_maps = torch.ones(
+ (
+ self.num_frames,
+ self.HEIGHT // self.buffer_downscale,
+ self.WIDTH // self.buffer_downscale,
+ ),
+ dtype=torch.float32,
+ device=self.device,
+ )
+
+ def get_image_error_video(self) -> List[np.ndarray]:
+ """
+ Get the pixel sample weights video.
+ Returns:
+ frames: the pixel sample weights video.
+ shape: (num_frames, H, W, 3)
+ """
+ # normalize the image error buffer to [0, 1]
+ image_error_maps = (
+ self.image_error_maps - self.image_error_maps.min()
+ ) / (self.image_error_maps.max() - self.image_error_maps.min())
+
+ maps = []
+ loss_maps = (
+ image_error_maps.detach()
+ .cpu()
+ .numpy()
+ .reshape(
+ self.num_frames,
+ self.HEIGHT // self.buffer_downscale,
+ self.WIDTH // self.buffer_downscale,
+ )
+ )
+ for i in range(self.num_frames):
+ maps.append(loss_maps[i])
+ return maps
+
+ def update_image_error_maps(self, render_results: Dict[str, Tensor]) -> None:
+ """
+ Update the image error buffer with the given render results.
+ """
+ gt_rgbs = render_results["gt_rgbs"]
+ pred_rgbs = render_results["rgbs"]
+ gt_rgbs = torch.from_numpy(np.stack(gt_rgbs, axis=0))
+ pred_rgbs = torch.from_numpy(np.stack(pred_rgbs, axis=0))
+ image_error_maps = torch.abs(gt_rgbs - pred_rgbs).mean(dim=-1)
+ assert image_error_maps.shape == self.image_error_maps.shape
+ if "Dynamic_opacities" in render_results:
+ if len(render_results["Dynamic_opacities"]) > 0:
+ dynamic_opacity = render_results["Dynamic_opacities"]
+ dynamic_opacity = torch.from_numpy(np.stack(dynamic_opacity, axis=0))
+ # we prioritize the dynamic objects by multiplying the error by 5
+ image_error_maps[dynamic_opacity > 0.1] *= 5
+ # update the image error buffer
+ self.image_error_maps: Tensor = image_error_maps.to(self.device)
+ logger.info(f"Updated image error buffer for camera {self.cam_id}.")
+
+ def to(self, device: torch.device):
+ """
+ Move the camera to the given device.
+ Args:
+ device: the device to move the camera to.
+ """
+ self.cam_to_worlds = self.cam_to_worlds.to(device)
+ self.intrinsics = self.intrinsics.to(device)
+ if self.distortions is not None:
+ self.distortions = self.distortions.to(device)
+ self.images = self.images.to(device)
+ if self.egocar_mask is not None:
+ self.egocar_mask = self.egocar_mask.to(device)
+ if self.dynamic_masks is not None:
+ self.dynamic_masks = self.dynamic_masks.to(device)
+ if self.human_masks is not None:
+ self.human_masks = self.human_masks.to(device)
+ if self.vehicle_masks is not None:
+ self.vehicle_masks = self.vehicle_masks.to(device)
+ if self.sky_masks is not None:
+ self.sky_masks = self.sky_masks.to(device)
+ if self.lidar_depth_maps is not None:
+ self.lidar_depth_maps = self.lidar_depth_maps.to(device)
+ if self.image_error_maps is not None:
+ self.image_error_maps = self.image_error_maps.to(device)
+
+ def get_image(self, frame_idx: int) -> Dict[str, Tensor]:
+ """
+ Get the rays for rendering the given frame index.
+ Args:
+ frame_idx: the frame index.
+ Returns:
+ a dict containing the rays for rendering the given frame index.
+ """
+ rgb, sky_mask = None, None
+ dynamic_mask, human_mask, vehicle_mask = None, None, None
+ pixel_coords, normalized_time = None, None
+ egocar_mask = None
+
+ if self.images is not None:
+ rgb = self.images[frame_idx]
+ if self.downscale_factor != 1.0:
+ rgb = (
+ torch.nn.functional.interpolate(
+ rgb.unsqueeze(0).permute(0, 3, 1, 2),
+ scale_factor=self.downscale_factor,
+ mode="bicubic",
+ antialias=True,
+ )
+ .squeeze(0)
+ .permute(1, 2, 0)
+ )
+ img_height, img_width = rgb.shape[:2]
+ else:
+ img_height, img_width = self.HEIGHT, self.WIDTH
+
+ x, y = torch.meshgrid(
+ torch.arange(img_width),
+ torch.arange(img_height),
+ indexing="xy",
+ )
+ x, y = x.flatten(), y.flatten()
+ x, y = x.to(self.device), y.to(self.device)
+ # pixel coordinates
+ pixel_coords = (
+ torch.stack([y / img_height, x / img_width], dim=-1)
+ .float()
+ .reshape(img_height, img_width, 2)
+ )
+ if self.egocar_mask is not None:
+ egocar_mask = self.egocar_mask
+ if self.downscale_factor != 1.0:
+ egocar_mask = (
+ torch.nn.functional.interpolate(
+ egocar_mask.unsqueeze(0).unsqueeze(0),
+ scale_factor=self.downscale_factor,
+ mode="nearest",
+ )
+ .squeeze(0)
+ .squeeze(0)
+ )
+ if self.sky_masks is not None:
+ sky_mask = self.sky_masks[frame_idx]
+ if self.downscale_factor != 1.0:
+ sky_mask = (
+ torch.nn.functional.interpolate(
+ sky_mask.unsqueeze(0).unsqueeze(0),
+ scale_factor=self.downscale_factor,
+ mode="nearest",
+ )
+ .squeeze(0)
+ .squeeze(0)
+ )
+ if self.dynamic_masks is not None:
+ dynamic_mask = self.dynamic_masks[frame_idx]
+ if self.downscale_factor != 1.0:
+ dynamic_mask = (
+ torch.nn.functional.interpolate(
+ dynamic_mask.unsqueeze(0).unsqueeze(0),
+ scale_factor=self.downscale_factor,
+ mode="nearest",
+ )
+ .squeeze(0)
+ .squeeze(0)
+ )
+ if self.human_masks is not None:
+ human_mask = self.human_masks[frame_idx]
+ if self.downscale_factor != 1.0:
+ human_mask = (
+ torch.nn.functional.interpolate(
+ human_mask.unsqueeze(0).unsqueeze(0),
+ scale_factor=self.downscale_factor,
+ mode="nearest",
+ )
+ .squeeze(0)
+ .squeeze(0)
+ )
+ if self.vehicle_masks is not None:
+ vehicle_mask = self.vehicle_masks[frame_idx]
+ if self.downscale_factor != 1.0:
+ vehicle_mask = (
+ torch.nn.functional.interpolate(
+ vehicle_mask.unsqueeze(0).unsqueeze(0),
+ scale_factor=self.downscale_factor,
+ mode="nearest",
+ )
+ .squeeze(0)
+ .squeeze(0)
+ )
+
+ lidar_depth_map = None
+ if self.lidar_depth_maps is not None:
+ lidar_depth_map = self.lidar_depth_maps[frame_idx]
+ if self.downscale_factor != 1.0:
+ # BUG: cannot use, need futher investigation
+ # if self.data_cfg.denser_lidar_times > 1:
+ # lidar_supervision_downsample = 1 / self.data_cfg.denser_lidar_times * self.downscale_factor
+ # lidar_depth_map = sparse_lidar_map_downsampler(lidar_depth_map, lidar_supervision_downsample)
+
+ # # resize back
+ # lidar_depth_map = (
+ # torch.nn.functional.interpolate(
+ # lidar_depth_map.unsqueeze(0).unsqueeze(0),
+ # scale_factor=self.data_cfg.denser_lidar_times,
+ # mode="bicubic",
+ # )
+ # .squeeze(0)
+ # .squeeze(0)
+ # )
+ # else:
+ lidar_depth_map = sparse_lidar_map_downsampler(lidar_depth_map, self.downscale_factor)
+
+ if self.normalized_time is not None:
+ normalized_time = torch.full(
+ (img_height, img_width),
+ self.normalized_time[frame_idx],
+ dtype=torch.float32,
+ )
+ camera_id = torch.full(
+ (img_height, img_width),
+ self.cam_id,
+ dtype=torch.long,
+ )
+ image_id = torch.full(
+ (img_height, img_width),
+ self.unique_img_idx[frame_idx],
+ dtype=torch.long,
+ )
+ frame_id = torch.full(
+ (img_height, img_width),
+ frame_idx,
+ dtype=torch.long,
+ )
+ c2w = self.cam_to_worlds[frame_idx]
+ intrinsics = self.intrinsics[frame_idx] * self.downscale_factor
+ intrinsics[2, 2] = 1.0
+ origins, viewdirs, direction_norm = get_rays(x, y, c2w, intrinsics)
+ origins = origins.reshape(img_height, img_width, 3)
+ viewdirs = viewdirs.reshape(img_height, img_width, 3)
+ direction_norm = direction_norm.reshape(img_height, img_width, 1)
+ _image_infos = {
+ "origins": origins,
+ "viewdirs": viewdirs,
+ "direction_norm": direction_norm,
+ "pixel_coords": pixel_coords,
+ "normed_time": normalized_time,
+ "img_idx": image_id,
+ "frame_idx": frame_id,
+ "pixels": rgb,
+ "sky_masks": sky_mask,
+ "dynamic_masks": dynamic_mask,
+ "human_masks": human_mask,
+ "vehicle_masks": vehicle_mask,
+ "egocar_masks": egocar_mask,
+ "lidar_depth_map": lidar_depth_map,
+ }
+ image_infos = {k: v for k, v in _image_infos.items() if v is not None}
+
+ cam_infos = {
+ "cam_id": camera_id,
+ "cam_name": self.cam_name,
+ "camera_to_world": c2w,
+ "height": torch.tensor(img_height, dtype=torch.long, device=c2w.device),
+ "width": torch.tensor(img_width, dtype=torch.long, device=c2w.device),
+ "intrinsics": intrinsics,
+ }
+ return image_infos, cam_infos
+
+class ScenePixelSource(abc.ABC):
+ """
+ The base class for all pixel sources of a scene.
+ """
+ # define a transformation matrix to convert the opencv camera coordinate system to the dataset camera coordinate system
+ data_cfg: OmegaConf = None
+ # the dataset name, choose from ["waymo", "kitti", "nuscenes", "pandaset", "argoverse"]
+ dataset_name: str = None
+ # the dict of camera data
+ camera_data: Dict[int, CameraData] = {}
+ # the normalized time of all images (normalized to [0, 1]), shape: (num_frames,)
+ _normalized_time: Tensor = None
+ # timestep indices of frames, shape: (num_frames,)
+ _timesteps: Tensor = None
+ # image error buffer, (num_images, )
+ image_error_buffer: Tensor = None
+ # whether the buffer is computed
+ image_error_buffered: bool = False
+ # the downscale factor of the error buffer
+ buffer_downscale: float = 1.0
+
+ # -------------- object annotations
+ # (num_frame, num_instances, 4, 4)
+ instances_pose: Tensor = None
+ # (num_instances, 3)
+ instances_size: Tensor = None
+ # (num_instances, )
+ instances_true_id: Tensor = None
+ # (num_instances, )
+ instances_model_types: Tensor = None
+ # (num_frame, num_instances)
+ per_frame_instance_mask: Tensor = None
+
+ def __init__(
+ self, dataset_name, pixel_data_config: OmegaConf, device: torch.device = torch.device("cpu")
+ ) -> None:
+ # hold the config of the pixel data
+ self.dataset_name = dataset_name
+ self.data_cfg = pixel_data_config
+ self.device = device
+ self._downscale_factor = 1 / pixel_data_config.downscale
+ self._old_downscale_factor = []
+
+ @abc.abstractmethod
+ def load_cameras(self) -> None:
+ """
+ Load the camera intrinsics, extrinsics, timestamps, etc.
+ Load the images, dynamic masks, sky masks, etc.
+ """
+ raise NotImplementedError
+
+ @abc.abstractmethod
+ def load_objects(self) -> None:
+ """
+ Load the object annotations.
+ """
+ raise NotImplementedError
+
+ def load_data(self) -> None:
+ """
+ A general function to load all data.
+ """
+ self.load_cameras()
+ self.build_image_error_buffer()
+ logger.info("[Pixel] All Pixel Data loaded.")
+
+ if self.data_cfg.load_objects:
+ self.load_objects()
+ logger.info("[Pixel] All Object Annotations loaded.")
+
+ # set initial downscale factor
+ for cam_id in self.camera_list:
+ self.camera_data[cam_id].set_downscale_factor(self._downscale_factor)
+
+ def to(self, device: torch.device) -> "ScenePixelSource":
+ """
+ Move the dataset to the given device.
+ Args:
+ device: the device to move the dataset to.
+ """
+ self.device = device
+ if self._timesteps is not None:
+ self._timesteps = self._timesteps.to(device)
+ if self._normalized_time is not None:
+ self._normalized_time = self._normalized_time.to(device)
+ if self.instances_pose is not None:
+ self.instances_pose = self.instances_pose.to(device)
+ if self.instances_size is not None:
+ self.instances_size = self.instances_size.to(device)
+ if self.per_frame_instance_mask is not None:
+ self.per_frame_instance_mask = self.per_frame_instance_mask.to(device)
+ if self.instances_model_types is not None:
+ self.instances_model_types = self.instances_model_types.to(device)
+ return self
+
+ def get_aabb(self) -> Tensor:
+ """
+ Returns:
+ aabb_min, aabb_max: the min and max of the axis-aligned bounding box of the scene
+ Note:
+ We compute the coarse aabb by using the front camera positions / trajectories. We then
+ extend this aabb by 40 meters along horizontal directions and 20 meters up and 5 meters
+ down along vertical directions.
+ """
+ front_camera_trajectory = self.front_camera_trajectory
+
+ # compute the aabb
+ aabb_min = front_camera_trajectory.min(dim=0)[0]
+ aabb_max = front_camera_trajectory.max(dim=0)[0]
+
+ # extend aabb by 40 meters along forward direction and 40 meters along the left/right direction
+ # aabb direction: x, y, z: front, left, up
+ aabb_max[0] += 40
+ aabb_max[1] += 40
+ # when the car is driving uphills
+ aabb_max[2] = min(aabb_max[2] + 20, 20)
+
+ # for waymo, there will be a lot of waste of space because we don't have images in the back,
+ # it's more reasonable to extend the aabb only by a small amount, e.g., 5 meters
+ # we use 40 meters here for a more general case
+ aabb_min[0] -= 40
+ aabb_min[1] -= 40
+ # when a car is driving downhills
+ aabb_min[2] = max(aabb_min[2] - 5, -5)
+ aabb = torch.tensor([*aabb_min, *aabb_max])
+ logger.info(f"[Pixel] Auto AABB from camera: {aabb}")
+ return aabb
+
+ @property
+ def front_camera_trajectory(self) -> Tensor:
+ """
+ Returns:
+ the front camera trajectory.
+ """
+ front_camera = self.camera_data[0]
+ assert (
+ front_camera.cam_to_worlds is not None
+ ), "Camera poses not loaded, cannot compute front camera trajectory."
+ return front_camera.cam_to_worlds[:, :3, 3]
+
+ def parse_img_idx(self, img_idx: int) -> Tuple[int, int]:
+ """
+ Parse the image index to the camera index and frame index.
+ Args:
+ img_idx: the image index.
+ Returns:
+ cam_idx: the camera index.
+ frame_idx: the frame index.
+ """
+ unique_cam_idx = img_idx % self.num_cams
+ frame_idx = img_idx // self.num_cams
+ return unique_cam_idx, frame_idx
+
+ def get_image(self, img_idx: int) -> Dict[str, Tensor]:
+ """
+ Get the rays for rendering the given image index.
+ Args:
+ img_idx: the image index.
+ Returns:
+ a dict containing the rays for rendering the given image index.
+ """
+ unique_cam_idx, frame_idx = self.parse_img_idx(img_idx)
+ for cam_id in self.camera_list:
+ if unique_cam_idx == self.camera_data[cam_id].unique_cam_idx:
+ return self.camera_data[cam_id].get_image(frame_idx)
+
+ @property
+ def camera_list(self) -> List[int]:
+ """
+ Returns:
+ the list of camera indices
+ """
+ return self.data_cfg.cameras
+
+ @property
+ def num_cams(self) -> int:
+ """
+ Returns:
+ the number of cameras in the dataset
+ """
+ return len(self.data_cfg.cameras)
+
+ @property
+ def num_frames(self) -> int:
+ """
+ Returns:
+ the number of frames in the dataset
+ """
+ return len(self._timesteps)
+
+ @property
+ def num_timesteps(self) -> int:
+ """
+ Returns:
+ the number of image timesteps in the dataset
+ """
+ return len(self._timesteps)
+
+ @property
+ def num_imgs(self) -> int:
+ """
+ Returns:
+ the number of images in the dataset
+ """
+ return self.num_cams * self.num_frames
+
+ @property
+ def timesteps(self) -> Tensor:
+ """
+ Returns:
+ the integer timestep indices of all images,
+ shape: (num_imgs,)
+ Note:
+ the difference between timestamps and timesteps is that
+ timestamps are the actual timestamps (minus 1e9) of images
+ while timesteps are the integer timestep indices of images.
+ """
+ return self._timesteps
+
+ @property
+ def normalized_time(self) -> Tensor:
+ """
+ Returns:
+ the normalized timestamps of all images
+ (normalized to the range [0, 1]),
+ shape: (num_imgs,)
+ """
+ return self._normalized_time
+
+ def register_normalized_timestamps(self) -> None:
+ # normalized timestamps are between 0 and 1
+ normalized_time = (self._timesteps - self._timesteps.min()) / (
+ self._timesteps.max() - self._timesteps.min()
+ )
+
+ self._normalized_time = normalized_time.to(self.device)
+ self._unique_normalized_timestamps = self._normalized_time.unique()
+
+ def find_closest_timestep(self, normed_timestamp: float) -> int:
+ """
+ Find the closest timestep to the given timestamp.
+ Args:
+ normed_timestamp: the normalized timestamp to find the closest timestep for.
+ Returns:
+ the closest timestep to the given timestamp.
+ """
+ return torch.argmin(
+ torch.abs(self._normalized_time - normed_timestamp)
+ )
+
+ def propose_training_image(
+ self,
+ candidate_indices: Tensor = None,
+ ) -> Dict[str, Tensor]:
+ if random.random() < self.buffer_ratio and self.image_error_buffered:
+ # sample according to the image error buffer
+ image_mean_error = self.image_error_buffer[candidate_indices]
+ start_enhance_weight = self.data_cfg.sampler.get('start_enhance_weight', 1)
+ if start_enhance_weight > 1:
+ # increase the error of the first 10% frames
+ frame_num = int(self.num_imgs / self.num_cams)
+ error_weight = torch.cat((
+ torch.linspace(start_enhance_weight, 1, int(frame_num * 0.1)),
+ torch.ones(frame_num - int(frame_num * 0.1))
+ ))
+ error_weight = error_weight[..., None].repeat(1, self.num_cams).reshape(-1)
+ error_weight = error_weight[candidate_indices].to(self.device)
+
+ image_mean_error = image_mean_error * error_weight
+ idx = torch.multinomial(
+ image_mean_error, 1, replacement=False
+ ).item()
+ img_idx = candidate_indices[idx]
+ else:
+ # random sample one from candidate_indices
+ img_idx = random.choice(candidate_indices)
+
+ return img_idx
+
+ def build_image_error_buffer(self) -> None:
+ """
+ Build the image error buffer.
+ """
+ if self.buffer_ratio > 0:
+ for cam_id in self.camera_list:
+ self.camera_data[cam_id].build_image_error_buffer()
+ else:
+ logger.info("Not building image error buffer because buffer_ratio <= 0.")
+
+ def update_image_error_maps(self, render_results: Dict[str, Tensor]) -> None:
+ """
+ Update the image error buffer with the given render results for each camera.
+ """
+ # (img_num, )
+ image_error_buffer = torch.zeros(self.num_imgs, device=self.device)
+ image_cam_id = torch.from_numpy(np.stack(render_results["cam_ids"], axis=0))
+ for cam_id in self.camera_list:
+ cam_name = self.camera_data[cam_id].cam_name
+ gt_rgbs, pred_rgbs = [], []
+ Dynamic_opacities = []
+ for img_idx, img_cam in enumerate(render_results["cam_names"]):
+ if img_cam == cam_name:
+ gt_rgbs.append(render_results["gt_rgbs"][img_idx])
+ pred_rgbs.append(render_results["rgbs"][img_idx])
+ if "Dynamic_opacities" in render_results:
+ Dynamic_opacities.append(render_results["Dynamic_opacities"][img_idx])
+
+ camera_results = {
+ "gt_rgbs": gt_rgbs,
+ "rgbs": pred_rgbs,
+ }
+ if len(Dynamic_opacities) > 0:
+ camera_results["Dynamic_opacities"] = Dynamic_opacities
+ self.camera_data[cam_id].update_image_error_maps(
+ camera_results
+ )
+
+ # update the image error buffer
+ image_error_buffer[image_cam_id == cam_id] = \
+ self.camera_data[cam_id].image_error_maps.mean(dim=(1, 2))
+
+
+ self.image_error_buffer = image_error_buffer
+ self.image_error_buffered = True
+ logger.info("Successfully updated image error buffer")
+
+ def get_image_error_video(self, layout: Callable) -> List[np.ndarray]:
+ """
+ Get the image error buffer video.
+ Returns:
+ frames: the pixel sample weights video.
+ """
+ per_cam_video = {}
+ for cam_id in self.camera_list:
+ per_cam_video[cam_id] = self.camera_data[cam_id].get_image_error_video()
+
+ all_error_images = []
+ all_cam_names = []
+ for frame_id in range(self.num_frames):
+ for cam_id in self.camera_list:
+ all_error_images.append(per_cam_video[cam_id][frame_id])
+ all_cam_names.append(self.camera_data[cam_id].cam_name)
+
+ merged_list = []
+ for i in range(len(all_error_images) // self.num_cams):
+ frames = all_error_images[
+ i
+ * self.num_cams : (i + 1)
+ * self.num_cams
+ ]
+ frames = [
+ np.stack([frame, frame, frame], axis=-1) for frame in frames
+ ]
+ cam_names = all_cam_names[
+ i
+ * self.num_cams : (i + 1)
+ * self.num_cams
+ ]
+ tiled_img = layout(frames, cam_names)
+ merged_list.append(tiled_img)
+
+ merged_video = np.stack(merged_list, axis=0)
+ merged_video -= merged_video.min()
+ merged_video /= merged_video.max()
+ merged_video = np.clip(merged_video * 255, 0, 255).astype(np.uint8)
+ return merged_video
+
+ @property
+ def downscale_factor(self) -> float:
+ """
+ Returns:
+ downscale_factor: the downscale factor of the images
+ """
+ return self._downscale_factor
+
+ def update_downscale_factor(self, downscale: float) -> None:
+ """
+ Args:
+ downscale: the new downscale factor
+ Updates the downscale factor
+ """
+ self._old_downscale_factor.append(self._downscale_factor)
+ self._downscale_factor = downscale
+ for cam_id in self.camera_list:
+ self.camera_data[cam_id].set_downscale_factor(self._downscale_factor)
+
+ def reset_downscale_factor(self) -> None:
+ """
+ Resets the downscale factor to the original value
+ """
+ assert len(self._old_downscale_factor) > 0, "No downscale factor to reset to"
+ self._downscale_factor = self._old_downscale_factor.pop()
+ for cam_id in self.camera_list:
+ self.camera_data[cam_id].set_downscale_factor(self._downscale_factor)
+
+ @property
+ def buffer_ratio(self) -> float:
+ """
+ Returns:
+ buffer_ratio: the ratio of the rays sampled from the image error buffer
+ """
+ return self.data_cfg.sampler.buffer_ratio
+
+ @property
+ def buffer_downscale(self) -> float:
+ """
+ Returns:
+ buffer_downscale: the downscale factor of the image error buffer
+ """
+ return self.data_cfg.sampler.buffer_downscale
+
+ def prepare_novel_view_render_data(self, dataset_type: str, traj: torch.Tensor, cam_id: int = None) -> list:
+ """
+ Prepare all necessary elements for novel view rendering.
+
+ Args:
+ dataset_type (str): Type of dataset
+ traj (torch.Tensor): Novel view trajectory, shape (N, 4, 4)
+ cam_id (int, optional): Specific camera ID to use for intrinsics and image size.
+ If None, keep dataset default behavior (Argoverse uses 1, others use 0).
+
+ Returns:
+ list: List of dicts, each containing elements required for rendering a single frame:
+ - cam_infos: Camera information (extrinsics, intrinsics, image dimensions)
+ - image_infos: Image-related information (indices, normalized time, viewdirs, etc.)
+ """
+ if cam_id is None:
+ if dataset_type == "argoverse":
+ cam_id = 1 # default for Argoverse
+ else:
+ cam_id = 0 # default for other datasets
+ # ensure the selected camera exists in current configuration
+ if cam_id not in self.camera_list:
+ raise ValueError(f"Requested camera id {cam_id} is not in active camera list {self.camera_list}")
+
+ intrinsics = self.camera_data[cam_id].intrinsics[0] # Assume intrinsics are constant across frames
+ H, W = self.camera_data[cam_id].HEIGHT, self.camera_data[cam_id].WIDTH
+
+ original_frame_count = self.num_frames
+ scaled_indices = torch.linspace(0, original_frame_count - 1, len(traj))
+ normed_time = torch.linspace(0, 1, len(traj))
+
+ render_data = []
+ for i in range(len(traj)):
+ c2w = traj[i]
+
+ # Generate ray origins and directions
+ x, y = torch.meshgrid(torch.arange(W), torch.arange(H), indexing='xy')
+ x, y = x.to(self.device), y.to(self.device)
+
+ origins, viewdirs, direction_norm = get_rays(x.flatten(), y.flatten(), c2w, intrinsics)
+ origins = origins.reshape(H, W, 3)
+ viewdirs = viewdirs.reshape(H, W, 3)
+ direction_norm = direction_norm.reshape(H, W, 1)
+
+ cam_infos = {
+ "camera_to_world": c2w,
+ "intrinsics": intrinsics,
+ "height": torch.tensor([H], dtype=torch.long, device=self.device),
+ "width": torch.tensor([W], dtype=torch.long, device=self.device),
+ }
+
+ image_infos = {
+ "origins": origins,
+ "viewdirs": viewdirs,
+ "direction_norm": direction_norm,
+ # keep per-video local indexing for appearance embeddings; videos are saved per-cam
+ "img_idx": torch.full((H, W), i, dtype=torch.long, device=self.device),
+ "frame_idx": torch.full((H, W), scaled_indices[i].round().long(), device=self.device),
+ "normed_time": torch.full((H, W), normed_time[i], dtype=torch.float32, device=self.device),
+ "pixel_coords": torch.stack(
+ [y.float() / H, x.float() / W], dim=-1
+ ), # [H, W, 2]
+ }
+
+ render_data.append({
+ "cam_infos": cam_infos,
+ "image_infos": image_infos,
+ })
+
+ return render_data
diff --git a/worldbench/third_party/drivestudio/datasets/base/scene_dataset.py b/worldbench/third_party/drivestudio/datasets/base/scene_dataset.py
new file mode 100644
index 0000000..d9ac480
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/base/scene_dataset.py
@@ -0,0 +1,119 @@
+import abc
+import logging
+from enum import IntEnum
+from typing import List, Callable
+
+import torch
+from omegaconf import OmegaConf
+from torch import Tensor
+
+from .lidar_source import SceneLidarSource
+from .pixel_source import ScenePixelSource
+from .split_wrapper import SplitWrapper
+
+logger = logging.getLogger()
+
+class ModelType(IntEnum):
+ RigidNodes = 0
+ SMPLNodes = 1
+ DeformableNodes = 2
+
+class SceneDataset(abc.ABC):
+ """
+ Base class for scene dataset.
+ """
+
+ data_cfg: OmegaConf = None
+ pixel_source: ScenePixelSource = None
+ lidar_source: SceneLidarSource = None
+ # training and testing indices are indices into the full dataset
+ # train_indices are img indices, so the length is num_cams * num_timesteps
+ train_indices: List[int] = None
+ test_indices: List[int] = None
+ # train_timesteps are timesteps, so the length is num_timesteps (len(unique_timesteps))
+ train_timesteps: Tensor = None
+ test_timesteps: Tensor = None
+
+ # dataset layout for multi-camera visualization
+ layout: Callable = None
+
+ # dataset wrappers
+ # full: includes all data
+ full_image_set: SplitWrapper = None
+ full_lidar_set: SplitWrapper = None
+ # train: includes only training data
+ train_image_set: SplitWrapper = None
+ train_lidar_set: SplitWrapper = None
+ # test: includes only testing data
+ test_image_set: SplitWrapper = None
+ test_lidar_set: SplitWrapper = None
+
+ def __init__(
+ self,
+ data_config: OmegaConf,
+ ):
+ super().__init__()
+ self.data_cfg = data_config
+
+ @abc.abstractmethod
+ def build_data_source(self):
+ """
+ Create the data source for the dataset.
+ """
+ raise NotImplementedError
+
+ @abc.abstractmethod
+ def build_split_wrapper(self):
+ """
+ Makes each data source as a Pytorch Dataset.
+ """
+ raise NotImplementedError
+
+ @abc.abstractmethod
+ def split_train_test(self):
+ raise NotImplementedError
+
+ def get_aabb(self) -> Tensor:
+ if self.lidar_source is not None:
+ aabb = self.lidar_source.get_aabb()
+ else:
+ aabb = self.pixel_source.get_aabb()
+ return aabb
+
+ def get_init_pcd(self) -> Tensor:
+ raise NotImplementedError
+
+ @property
+ def num_cams(self) -> int:
+ return self.pixel_source.num_cams
+
+ @property
+ def scene_idx(self) -> int:
+ return self.data_cfg.scene_idx
+
+ @property
+ def num_img_timesteps(self) -> int:
+ return self.pixel_source.num_timesteps
+
+ @property
+ def num_lidar_timesteps(self) -> int:
+ if self.lidar_source is None:
+ logger.warning("No lidar source, returning num_img_timesteps")
+ return self.num_img_timesteps
+ return self.lidar_source.num_timesteps
+
+ @property
+ def num_train_timesteps(self) -> int:
+ return len(self.train_timesteps)
+
+ @property
+ def num_test_timesteps(self) -> int:
+ return len(self.test_timesteps)
+
+ @property
+ def unique_normalized_training_timestamps(self) -> Tensor:
+ return self.pixel_source.unique_normalized_timestamps[self.train_timesteps]
+
+ @property
+ def device(self):
+ return self.data_cfg.preload_device
diff --git a/worldbench/third_party/drivestudio/datasets/base/split_wrapper.py b/worldbench/third_party/drivestudio/datasets/base/split_wrapper.py
new file mode 100644
index 0000000..f7601a7
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/base/split_wrapper.py
@@ -0,0 +1,54 @@
+from typing import List, Tuple
+import torch
+
+from .pixel_source import ScenePixelSource
+
+class SplitWrapper(torch.utils.data.Dataset):
+
+ # a sufficiently large number to make sure we don't run out of data
+ _num_iters = 1000000
+
+ def __init__(
+ self,
+ datasource: ScenePixelSource,
+ split_indices: List[int] = None,
+ split: str = "train",
+ ):
+ super().__init__()
+ self.datasource = datasource
+ self.split_indices = split_indices
+ self.split = split
+
+ def get_image(self, idx, camera_downscale) -> dict:
+ downscale_factor = 1 / camera_downscale * self.datasource.downscale_factor
+ self.datasource.update_downscale_factor(downscale_factor)
+ image_infos, cam_infos = self.datasource.get_image(self.split_indices[idx])
+ self.datasource.reset_downscale_factor()
+ return image_infos, cam_infos
+
+ def next(self, camera_downscale) -> Tuple[dict, dict]:
+ assert self.split == "train", "Only train split supports next()"
+
+ img_idx = self.datasource.propose_training_image(
+ candidate_indices=self.split_indices
+ )
+
+ downscale_factor = 1 / camera_downscale * self.datasource.downscale_factor
+ self.datasource.update_downscale_factor(downscale_factor)
+ image_infos, cam_infos = self.datasource.get_image(img_idx)
+ self.datasource.reset_downscale_factor()
+
+ return image_infos, cam_infos
+
+ def __getitem__(self, idx) -> dict:
+ return self.get_image(idx, camera_downscale=1.0)
+
+ def __len__(self) -> int:
+ return len(self.split_indices)
+
+ @property
+ def num_iters(self) -> int:
+ return self._num_iters
+
+ def set_num_iters(self, num_iters) -> None:
+ self._num_iters = num_iters
diff --git a/worldbench/third_party/drivestudio/datasets/base/utils.py b/worldbench/third_party/drivestudio/datasets/base/utils.py
new file mode 100644
index 0000000..f02978d
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/base/utils.py
@@ -0,0 +1,85 @@
+from typing import List, Tuple, Union
+import numpy as np
+import imageio
+import numbers
+
+import torch
+import skimage
+from skimage.transform import resize as cpu_resize
+from torchvision.transforms.functional import resize as gpu_resize
+
+def load_rgb(path: str, downscale: numbers.Number = 1) -> np.ndarray:
+ """ Load image
+
+ Args:
+ path (str): Given image file path
+ downscale (numbers.Number, optional): Optional downscale ratio. Defaults to 1.
+
+ Returns:
+ np.ndarray: [H, W, 3], in range [0,1]
+ """
+ img = imageio.imread(path)
+ img = skimage.img_as_float32(img)
+ if downscale != 1:
+ H, W, _ = img.shape
+ img = cpu_resize(img, (int(H // downscale), int(W // downscale)), anti_aliasing=False)
+ # [H, W, 3]
+ return img
+
+def img_to_torch_and_downscale(
+ x: Union[np.ndarray, torch.Tensor], hw: Tuple[int, int],
+ use_cpu_downscale=False, antialias=False,
+ dtype=None, device=None):
+ """ Check, convert and apply downscale to input image `x`
+
+ Args:
+ x (Union[np.ndarray, torch.Tensor]): [H, W, (...)] Input image
+ downscale (float, optional): Downscaling ratio. Defaults to 1.
+ use_cpu_downscale (bool, optional): Whether use CPU downscaling algo (T), or use GPU (F). Defaults to False.
+ antialias (bool, optional): Whether use anti-aliasing. Defaults to False.
+ dtype (torch.dtype, optional): Output torch.dtype. Defaults to None.
+ device (torch.device, optional): Output torch.device. Defaults to None.
+
+ Returns:
+ torch.Tensor: [new_H, new_W, (...)] Converted and downscaled torch.Tensor image
+ """
+ H_, W_ = hw
+ if use_cpu_downscale:
+ x_np = x if isinstance(x, np.ndarray) else x.data.cpu().numpy()
+ x = torch.tensor(cpu_resize(x_np, (H_, W_), anti_aliasing=antialias),
+ dtype=dtype, device=device)
+ else:
+ x = check_to_torch(x, dtype=dtype, device=device)
+ x = x.cuda() if not x.is_cuda else x
+ if x.dim() == 2:
+ x = gpu_resize(x.unsqueeze(0), (H_, W_), antialias=antialias).squeeze(0)
+ else:
+ x = gpu_resize(x.movedim(-1, 0), (H_, W_), antialias=antialias).movedim(0, -1)
+ assert [H_, W_] == [*x.shape[:2]]
+ return check_to_torch(x, dtype=dtype, device=device)
+
+def check_to_torch(
+ x: Union[np.ndarray, torch.Tensor, List, Tuple],
+ ref: torch.Tensor=None, dtype: torch.dtype=None, device: torch.device=None) -> torch.Tensor:
+ """ Check and convert input `x` to torch.Tensor
+
+ Args:
+ x (Union[np.ndarray, torch.Tensor, List, Tuple]): Input
+ ref (torch.Tensor, optional): Reference tensor for dtype and device. Defaults to None.
+ dtype (torch.dtype, optional): Target torch.dtype. Defaults to None.
+ device (torch.device, optional): Target torch.device. Defaults to None.
+
+ Returns:
+ torch.Tensor: Converted torch.Tensor
+ """
+ if ref is not None:
+ if dtype is None:
+ dtype = ref.dtype
+ if device is None:
+ device = ref.device
+ if x is None:
+ return x
+ elif isinstance(x, torch.Tensor):
+ return x.to(dtype=dtype or x.dtype, device=device or x.device)
+ else:
+ return torch.tensor(x, dtype=dtype, device=device)
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/dataset_meta.py b/worldbench/third_party/drivestudio/datasets/dataset_meta.py
new file mode 100644
index 0000000..fb28d27
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/dataset_meta.py
@@ -0,0 +1,184 @@
+DATASETS_CONFIG = {
+ "waymo": {
+ 0: {
+ "camera_name": "front_camera",
+ "original_size": (1280, 1920),
+ "egocar_visible": False
+ },
+ 1: {
+ "camera_name": "front_left_camera",
+ "original_size": (1280, 1920),
+ "egocar_visible": False
+ },
+ 2: {
+ "camera_name": "front_right_camera",
+ "original_size": (1280, 1920),
+ "egocar_visible": False
+ },
+ 3: {
+ "camera_name": "left_camera",
+ "original_size": (866, 1920),
+ "egocar_visible": False
+ },
+ 4: {
+ "camera_name": "right_camera",
+ "original_size": (866, 1920),
+ "egocar_visible": False
+ },
+ },
+ "pandaset": {
+ 0: {
+ "camera_name": "front_camera",
+ "original_size": (1080, 1920),
+ "egocar_visible": False
+ },
+ 1: {
+ "camera_name": "front_left_camera",
+ "original_size": (1080, 1920),
+ "egocar_visible": False
+ },
+ 2: {
+ "camera_name": "front_right_camera",
+ "original_size": (1080, 1920),
+ "egocar_visible": False
+ },
+ 3: {
+ "camera_name": "left_camera",
+ "original_size": (1080, 1920),
+ "egocar_visible": False
+ },
+ 4: {
+ "camera_name": "right_camera",
+ "original_size": (1080, 1920),
+ "egocar_visible": False
+ },
+ 5: {
+ "camera_name": "back_camera",
+ "original_size": (1080, 1920),
+ "egocar_visible": True
+ },
+ },
+ "argoverse": {
+ 0: {
+ "camera_name": "ring_front_center",
+ "original_size": (2048, 1550),
+ "egocar_visible": True
+ },
+ 1: {
+ "camera_name": "ring_front_left",
+ "original_size": (1550, 2048),
+ "egocar_visible": False
+ },
+ 2: {
+ "camera_name": "ring_front_right",
+ "original_size": (1550, 2048),
+ "egocar_visible": False
+ },
+ 3: {
+ "camera_name": "ring_side_left",
+ "original_size": (1550, 2048),
+ "egocar_visible": False
+ },
+ 4: {
+ "camera_name": "ring_side_right",
+ "original_size": (1550, 2048),
+ "egocar_visible": False
+ },
+ 5: {
+ "camera_name": "ring_rear_left",
+ "original_size": (1550, 2048),
+ "egocar_visible": True
+ },
+ 6: {
+ "camera_name": "ring_rear_right",
+ "original_size": (1550, 2048),
+ "egocar_visible": True
+ },
+ },
+ "nuscenes": {
+ 0: {
+ "camera_name": "CAM_FRONT",
+ "original_size": (900, 1600),
+ "egocar_visible": False
+ },
+ 1: {
+ "camera_name": "CAM_FRONT_LEFT",
+ "original_size": (900, 1600),
+ "egocar_visible": False
+ },
+ 2: {
+ "camera_name": "CAM_FRONT_RIGHT",
+ "original_size": (900, 1600),
+ "egocar_visible": False
+ },
+ 3: {
+ "camera_name": "CAM_BACK_LEFT",
+ "original_size": (900, 1600),
+ "egocar_visible": False
+ },
+ 4: {
+ "camera_name": "CAM_BACK_RIGHT",
+ "original_size": (900, 1600),
+ "egocar_visible": False
+ },
+ 5: {
+ "camera_name": "CAM_BACK",
+ "original_size": (900, 1600),
+ "egocar_visible": True
+ }
+ },
+ "kitti": {
+ 0: {
+ "camera_name": "CAM_LEFT",
+ "original_size": (375, 1242),
+ "egocar_visible": False
+ },
+ 1: {
+ "camera_name": "CAM_RIGHT",
+ "original_size": (375, 1242),
+ "egocar_visible": False
+ }
+ },
+ "nuplan": {
+ 0: {
+ "camera_name": "CAM_F0",
+ "original_size": (1080, 1920),
+ "egocar_visible": False
+ },
+ 1: {
+ "camera_name": "CAM_L0",
+ "original_size": (1080, 1920),
+ "egocar_visible": True
+ },
+ 2: {
+ "camera_name": "CAM_R0",
+ "original_size": (1080, 1920),
+ "egocar_visible": True
+ },
+ 3: {
+ "camera_name": "CAM_L1",
+ "original_size": (1080, 1920),
+ "egocar_visible": False
+ },
+ 4: {
+ "camera_name": "CAM_R1",
+ "original_size": (1080, 1920),
+ "egocar_visible": False
+ },
+ 5: {
+ "camera_name": "CAM_L2",
+ "original_size": (1080, 1920),
+ "egocar_visible": True
+ },
+ 6: {
+ "camera_name": "CAM_R2",
+ "original_size": (1080, 1920),
+ "egocar_visible": True
+ },
+ 7: {
+ "camera_name": "CAM_B0",
+ "original_size": (1080, 1920),
+ "egocar_visible": False
+ },
+ },
+}
diff --git a/worldbench/third_party/drivestudio/datasets/driving_dataset.py b/worldbench/third_party/drivestudio/datasets/driving_dataset.py
new file mode 100644
index 0000000..635bf3e
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/driving_dataset.py
@@ -0,0 +1,760 @@
+from typing import Dict, Union, Literal
+import logging
+import os
+import cv2
+import numpy as np
+from tqdm import trange, tqdm
+from omegaconf import OmegaConf
+
+import torch
+from torch import Tensor
+
+from models.gaussians.basics import *
+from datasets.base.scene_dataset import ModelType
+from datasets.base.scene_dataset import SceneDataset
+from datasets.base.split_wrapper import SplitWrapper
+from utils.visualization import get_layout
+from utils.geometry import transform_points
+from utils.camera import get_interp_novel_trajectories
+from utils.misc import export_points_to_ply, import_str
+
+logger = logging.getLogger()
+
+DEBUG_PCD=False
+if DEBUG_PCD:
+ DEBUG_OUTPUT_DIR="debug"
+ os.makedirs(DEBUG_OUTPUT_DIR, exist_ok=True)
+
+NAME_TO_NODE = {
+ "RigidNodes": ModelType.RigidNodes,
+ "SMPLNodes": ModelType.SMPLNodes,
+ "DeformableNodes": ModelType.DeformableNodes
+}
+
+class DrivingDataset(SceneDataset):
+ def __init__(
+ self,
+ data_cfg: OmegaConf,
+ ) -> None:
+ super().__init__(data_cfg)
+
+ # AVAILABLE DATASETS:
+ # Waymo: 5 Cameras
+ # KITTI: 2 Cameras
+ # NuScenes: 6 Cameras
+ # ArgoVerse:7 Cameras
+ # PandaSet: 6 Cameras
+ # NuPlan: 8 Cameras
+ self.type = self.data_cfg.dataset
+ try: # For Waymo, NuScenes, ArgoVerse, PandaSet
+ self.data_path = os.path.join(
+ self.data_cfg.data_root,
+ f"{int(self.scene_idx):03d}"
+ )
+ except: # For KITTI, NuPlan
+ self.data_path = os.path.join(self.data_cfg.data_root, self.scene_idx)
+
+ assert os.path.exists(self.data_path), f"{self.data_path} does not exist"
+ if os.path.exists(os.path.join(self.data_path, "ego_pose")):
+ total_frames = len(os.listdir(os.path.join(self.data_path, "ego_pose")))
+ elif os.path.exists(os.path.join(self.data_path, "lidar_pose")):
+ total_frames = len(os.listdir(os.path.join(self.data_path, "lidar_pose")))
+ else:
+ raise ValueError("Unable to determine the total number of frames. Neither 'ego_pose' nor 'lidar_pose' directories found.")
+
+ # ---- find the number of synchronized frames ---- #
+ if self.data_cfg.end_timestep == -1:
+ end_timestep = total_frames - 1
+ else:
+ end_timestep = self.data_cfg.end_timestep
+ # to make sure the last timestep is included
+ self.end_timestep = end_timestep + 1
+ self.start_timestep = self.data_cfg.start_timestep
+
+ # ---- create layout for visualization ---- #
+ self.layout = get_layout(self.type)
+
+ # ---- create data source ---- #
+ self.pixel_source, self.lidar_source = self.build_data_source()
+ assert self.pixel_source is not None and self.lidar_source is not None, \
+ "Must have both pixel source and lidar source"
+ self.project_lidar_pts_on_images(
+ delete_out_of_view_points=True
+ )
+ self.aabb = self.get_aabb()
+
+ # ---- define train and test indices ---- #
+ # note that the timestamps of the pixel source and the lidar source are the same in waymo dataset
+ (
+ self.train_timesteps,
+ self.test_timesteps,
+ self.train_indices,
+ self.test_indices,
+ ) = self.split_train_test()
+
+ # ---- create split wrappers ---- #
+ image_sets = self.build_split_wrapper()
+ self.train_image_set, self.test_image_set, self.full_image_set = image_sets
+
+ # debug use
+ # self.seg_dynamic_instances_in_lidar_frame(-1, frame_idx=0)
+ # self.get_init_objects()
+
+ @property
+ def instance_num(self):
+ return len(self.pixel_source.instances_pose[0])
+
+ @property
+ def frame_num(self):
+ return self.pixel_source.num_frames
+
+ def get_instance_infos(self):
+ return (
+ self.pixel_source.instances_pose.clone(),
+ self.pixel_source.instances_size.clone(),
+ self.pixel_source.instances_model_types.clone(),
+ self.pixel_source.per_frame_instance_mask.clone()
+ )
+
+ def build_split_wrapper(self):
+ train_image_set = SplitWrapper(
+ datasource=self.pixel_source,
+ # train_indices are img indices, so the length is num_cams * num_train_timesteps
+ split_indices=self.train_indices,
+ split="train",
+ )
+ full_image_set = SplitWrapper(
+ datasource=self.pixel_source,
+ # cover all the images
+ split_indices=np.arange(self.pixel_source.num_imgs).tolist(),
+ split="full",
+ )
+ test_image_set = None
+ if len(self.test_indices) > 0:
+ test_image_set = SplitWrapper(
+ datasource=self.pixel_source,
+ # test_indices are img indices, so the length is num_cams * num_test_timesteps
+ split_indices=self.test_indices,
+ split="test",
+ )
+ image_sets = (train_image_set, test_image_set, full_image_set)
+ return image_sets
+
+ def build_data_source(self):
+ """
+ Create the data source for the dataset.
+ """
+ # ---- create pixel source ---- #
+ pixel_source = import_str(self.data_cfg.pixel_source.type)(
+ self.data_cfg.dataset,
+ self.data_cfg.pixel_source,
+ self.data_path,
+ self.start_timestep,
+ self.end_timestep,
+ device=self.device,
+ )
+ pixel_source.to(self.device)
+
+ # ---- create lidar source ---- #
+ lidar_source = None
+ if self.data_cfg.lidar_source.load_lidar:
+ lidar_source = import_str(self.data_cfg.lidar_source.type)(
+ self.data_cfg.lidar_source,
+ self.data_path,
+ self.start_timestep,
+ self.end_timestep,
+ device=self.device,
+ )
+ lidar_source.to(self.device)
+ assert (pixel_source._unique_normalized_timestamps - lidar_source._unique_normalized_timestamps).abs().sum().item() == 0., \
+ "The timestamps of the pixel source and the lidar source are not synchronized"
+ return pixel_source, lidar_source
+
+ def get_lidar_samples(
+ self,
+ num_samples: float = None,
+ downsample_factor: float = None,
+ return_color=False,
+ return_normalized_time=False,
+ device: torch.device = torch.device("cpu")
+ ) -> Tensor:
+ assert self.lidar_source is not None, "Must have lidar source if you want to get init pcd"
+ assert (num_samples is None) != (downsample_factor is None), \
+ "Must provide either num_samples or downsample_factor, but not both"
+ if downsample_factor is not None:
+ num_samples = int(len(self.lidar_source.pts_xyz) / downsample_factor)
+ if num_samples > len(self.lidar_source.pts_xyz):
+ logger.warning(f"num_samples {num_samples} is larger than the number of points {len(self.lidar_source.pts_xyz)}")
+ num_samples = len(self.lidar_source.pts_xyz)
+
+ # randomly sample points
+ sampled_idx = torch.randperm(len(self.lidar_source.pts_xyz))[:num_samples]
+ sampled_pts = self.lidar_source.pts_xyz[sampled_idx].to(device)
+
+ # get color if needed
+ sampled_color = None
+ if return_color:
+ sampled_color = self.lidar_source.colors[sampled_idx].to(device)
+
+ sampled_time = None
+ if return_normalized_time:
+ sampled_time = self.lidar_source._normalized_time[sampled_idx].to(device)
+ sampled_time = sampled_time[..., None]
+
+ return sampled_pts, sampled_color, sampled_time
+
+ def seg_dynamic_instances_in_lidar_frame(
+ self,
+ instance_ids: Union[int, list],
+ frame_idx: int
+ ):
+ if isinstance(instance_ids, int):
+ instance_num = len(self.pixel_source.instances_pose[frame_idx])
+ assert instance_ids < instance_num, f"instance_id {instance_ids} is larger than the number of instances {instance_num}"
+ if instance_ids == -1:
+ instance_ids = list(range(instance_num))
+ else:
+ instance_ids = [instance_ids]
+ elif isinstance(instance_ids, list):
+ instance_ids = instance_ids
+
+ # get the lidar points
+ lidar_dict = self.lidar_source.get_lidar_rays(frame_idx)
+ lidar_pts = lidar_dict["lidar_origins"] + lidar_dict["lidar_viewdirs"] * lidar_dict["lidar_ranges"]
+ valid_mask = torch.zeros_like(lidar_pts[:, 0]).bool()
+ for instance_id in instance_ids:
+ is_valid_instance = self.pixel_source.per_frame_instance_mask[frame_idx, instance_id]
+ if not is_valid_instance:
+ continue
+ # get the pose of the instance at the given frame
+ o2w = self.pixel_source.instances_pose[frame_idx, instance_id]
+ o_size = self.pixel_source.instances_size[instance_id]
+
+ # transform the lidar points to the instance's coordinate system
+ # instance_pose [4, 4], pts [N, 3]
+ w2o = torch.inverse(o2w)
+ o_pts = transform_points(lidar_pts, w2o)
+ # get the mask of the points that are inside the instance's bounding box
+ mask = (
+ (o_pts[:, 0] > -o_size[0] / 2)
+ & (o_pts[:, 0] < o_size[0] / 2)
+ & (o_pts[:, 1] > -o_size[1] / 2)
+ & (o_pts[:, 1] < o_size[1] / 2)
+ & (o_pts[:, 2] > -o_size[2] / 2)
+ & (o_pts[:, 2] < o_size[2] / 2)
+ )
+ valid_mask = valid_mask | mask
+
+ valid_points = lidar_pts[valid_mask]
+ valid_colors = self.lidar_source.colors[lidar_dict["lidar_mask"]][valid_mask]
+
+ if DEBUG_PCD:
+ export_points_to_ply(
+ valid_points,
+ valid_colors,
+ save_path=os.path.join(DEBUG_OUTPUT_DIR, "vehicle_lidar_pts.ply")
+ )
+ export_points_to_ply(
+ lidar_pts,
+ self.lidar_source.colors[lidar_dict["lidar_mask"]],
+ save_path=os.path.join(DEBUG_OUTPUT_DIR, "lidar_pts.ply")
+ )
+
+ def get_init_objects(
+ self,
+ cur_node_type: Literal["RigidNodes", "DeformableNodes"],
+ instance_max_pts: int = 5000,
+ only_moving: bool = True,
+ traj_length_thres: float = 0.5,
+ exclude_smpl: bool = False,
+ ):
+ """
+ return:
+ instances_dict: Dict[int, Dict[str, Tensor]]
+ keys: instance_id
+ values: Dict[str, Tensor]
+ keys: "pts", "colors", "num_pts", "flows"(Optional)
+ values: Tensor
+
+ NOTE: pts are in object coordinate system
+ """
+ if self.type == "KITTI":
+ traj_length_thres = 5.0
+ logger.info(f"For KITTI dataset, the trajectory length threshold is set \
+ to {traj_length_thres} to filter out noisy short trajectories of static objects")
+
+ instance_dict = {}
+ for fi in range(self.frame_num):
+ lidar_dict = self.lidar_source.get_lidar_rays(fi)
+ lidar_pts = lidar_dict["lidar_origins"] + lidar_dict["lidar_viewdirs"] * lidar_dict["lidar_ranges"]
+ for ins_id in range(self.instance_num):
+ instance_active = self.pixel_source.per_frame_instance_mask[fi, ins_id]
+ o_type = self.pixel_source.instances_model_types[ins_id].item()
+
+ if not instance_active:
+ continue
+
+ if cur_node_type == "DeformableNodes":
+ if not (
+ o_type == ModelType.DeformableNodes or
+ o_type == ModelType.SMPLNodes
+ ):
+ continue
+ elif cur_node_type == "RigidNodes":
+ if not o_type == ModelType.RigidNodes:
+ continue
+
+ if exclude_smpl:
+ # objects with smpl pose will be modeled by SMPLNodes
+ assert cur_node_type == "DeformableNodes", \
+ "Only exclude SMPL for DeformableNodes"
+ true_id = self.pixel_source.instances_true_id[ins_id].item()
+ if true_id in self.pixel_source.smpl_human_all.keys():
+ continue
+
+ if ins_id not in instance_dict:
+ instance_dict[ins_id] = {
+ "node_type": cur_node_type,
+ "pts": [],
+ "colors": [],
+ # "flows": [],
+ }
+ # get the pose of the instance at the given frame
+ o2w = self.pixel_source.instances_pose[fi, ins_id]
+ o_size = self.pixel_source.instances_size[ins_id]
+ # convert the lidar points to the instance's coordinate system
+ w2o = torch.inverse(o2w)
+ o_pts = transform_points(lidar_pts, w2o)
+ # get the mask of the points that are inside the instance's bounding box
+ mask = (
+ (o_pts[:, 0] > -o_size[0] / 2)
+ & (o_pts[:, 0] < o_size[0] / 2)
+ & (o_pts[:, 1] > -o_size[1] / 2)
+ & (o_pts[:, 1] < o_size[1] / 2)
+ & (o_pts[:, 2] > -o_size[2] / 2)
+ & (o_pts[:, 2] < o_size[2] / 2)
+ )
+ valid_pts = o_pts[mask]
+ valid_colors = self.lidar_source.colors[lidar_dict["lidar_mask"]][mask]
+ # valid_flows = lidar_dict["lidar_flows"][mask]
+ instance_dict[ins_id]["pts"].append(valid_pts)
+ instance_dict[ins_id]["colors"].append(valid_colors)
+ # instance_dict[ins_id]["flows"].append(valid_flows)
+
+ logger.info(f"Aggregating lidar points across {self.frame_num} frames")
+ for ins_id in instance_dict:
+ instance_dict[ins_id]["pts"] = torch.cat(instance_dict[ins_id]["pts"], dim=0)
+ instance_dict[ins_id]["colors"] = torch.cat(instance_dict[ins_id]["colors"], dim=0)
+ # instance_dict[ins_id]["flows"] = torch.cat(instance_dict[ins_id]["flows"], dim=0)
+ instance_dict[ins_id]["num_pts"] = instance_dict[ins_id]["pts"].shape[0]
+ if instance_dict[ins_id]["num_pts"] > instance_max_pts:
+ # randomly sample points
+ sampled_idx = torch.randperm(instance_dict[ins_id]["num_pts"])[:instance_max_pts]
+ instance_dict[ins_id]["pts"] = instance_dict[ins_id]["pts"][sampled_idx]
+ instance_dict[ins_id]["colors"] = instance_dict[ins_id]["colors"][sampled_idx]
+ # instance_dict[ins_id]["flows"] = instance_dict[ins_id]["flows"][sampled_idx]
+ instance_dict[ins_id]["num_pts"] = instance_max_pts
+ logger.info(f"Instance {ins_id} has {instance_dict[ins_id]['num_pts']} lidar sample points")
+
+ if only_moving:
+ # consider only the instances with non-zero flows
+ logger.info(f"Filtering out the instances with non-moving trajectories")
+ new_instance_dict = {}
+ for k, v in instance_dict.items():
+ if v["num_pts"] > 0:
+ # flows = v["flows"]
+ # if flows.norm(dim=-1).mean() > moving_thres:
+ # v.pop("flows")
+ # new_instance_dict[k] = v
+ # logger.info(f"Instance {k} has {v['num_pts']} lidar sample points")
+ frame_info = self.pixel_source.per_frame_instance_mask[:, k]
+ instances_pose = self.pixel_source.instances_pose[:, k]
+ instances_trans = instances_pose[:, :3, 3]
+ valid_trans = instances_trans[frame_info]
+ traj_length = valid_trans[1:] - valid_trans[:-1]
+ traj_length = torch.norm(traj_length, dim=-1).sum()
+ if traj_length > traj_length_thres:
+ new_instance_dict[k] = v
+ logger.info(f"Instance {k} has {v['num_pts']} lidar sample points")
+ instance_dict = new_instance_dict
+
+ # get instance info
+ for ins_id in instance_dict:
+ instance_dict[ins_id]["poses"] = self.pixel_source.instances_pose[:, ins_id]
+ instance_dict[ins_id]["size"] = self.pixel_source.instances_size[ins_id]
+ instance_dict[ins_id]["frame_info"] = self.pixel_source.per_frame_instance_mask[:, ins_id]
+
+ if DEBUG_PCD:
+ output_dir = os.path.join(DEBUG_OUTPUT_DIR, "aggregated_instance_lidar_pts")
+ os.makedirs(output_dir, exist_ok=True)
+ for ins_id in instance_dict:
+ export_points_to_ply(
+ instance_dict[ins_id]["pts"],
+ instance_dict[ins_id]["colors"],
+ save_path=os.path.join(output_dir, f"ID={ins_id}.ply")
+ )
+ return instance_dict
+
+ def get_init_smpl_objects(self, only_moving: bool = False, traj_length_thres: float = 0.5):
+ instance_dict = {}
+ """
+ instance_dict = {
+ ins_id: {
+ "node_type": str,
+ "pts": Tensor, [frame_num, num_pts, 3]
+ "colors": Tensor, [frame_num, num_pts, 3]
+ "quats": Tensor, [frame_num, 4]
+ "trans": Tensor, [frame_num, 3]
+ "size": Tensor, [3]
+ "frame_info": Tensor, [frame_num]
+ }
+ """
+
+ for ins_id in range(self.instance_num):
+ true_id = self.pixel_source.instances_true_id[ins_id].item()
+ if true_id in self.pixel_source.smpl_human_all.keys():
+ if self.pixel_source.smpl_human_all[true_id]["frame_valid"].sum() == 0:
+ continue
+ smpl_trans = self.pixel_source.smpl_human_all[true_id]["smpl_trans"]
+ frame_info = self.pixel_source.smpl_human_all[true_id]["frame_valid"]
+ if only_moving and traj_length_thres > 0:
+ # compute the distance between two consecutive frames
+ traj_length = smpl_trans[frame_info][1:] - smpl_trans[frame_info][:-1]
+ traj_length = torch.norm(traj_length, dim=-1).sum()
+ if traj_length < traj_length_thres:
+ continue
+ smpl_quats = self.pixel_source.smpl_human_all[true_id]["smpl_quats"]
+ smpl_betas = self.pixel_source.smpl_human_all[true_id]["smpl_betas"]
+ size = self.pixel_source.instances_size[ins_id]
+ # NOTE: set the first frame's betas as the betas of the instance
+ first_frame_betas = smpl_betas[frame_info][0]
+
+ collected_lidar_pts = []
+ collected_lidar_colors = []
+ for fi in range(self.frame_num):
+ lidar_dict = self.lidar_source.get_lidar_rays(fi)
+ lidar_pts = lidar_dict["lidar_origins"] + lidar_dict["lidar_viewdirs"] * lidar_dict["lidar_ranges"]
+ instance_active = self.pixel_source.per_frame_instance_mask[fi, ins_id]
+ if not instance_active:
+ continue
+
+ # get the pose of the instance at the given frame
+ o2w = self.pixel_source.instances_pose[fi, ins_id]
+ o_size = self.pixel_source.instances_size[ins_id]
+ # convert the lidar points to the instance's coordinate system
+ w2o = torch.inverse(o2w)
+ o_pts = transform_points(lidar_pts, w2o)
+ # get the mask of the points that are inside the instance's bounding box
+ mask = (
+ (o_pts[:, 0] > -o_size[0] / 2)
+ & (o_pts[:, 0] < o_size[0] / 2)
+ & (o_pts[:, 1] > -o_size[1] / 2)
+ & (o_pts[:, 1] < o_size[1] / 2)
+ & (o_pts[:, 2] > -o_size[2] / 2)
+ & (o_pts[:, 2] < o_size[2] / 2)
+ )
+ valid_pts = o_pts[mask]
+ valid_colors = self.lidar_source.colors[lidar_dict["lidar_mask"]][mask]
+ # valid_flows = lidar_dict["lidar_flows"][mask]
+ collected_lidar_pts.append(valid_pts)
+ collected_lidar_colors.append(valid_colors)
+
+ instance_dict[ins_id] = {
+ "node_type": "SMPLNodes",
+ "smpl_quats": smpl_quats, # [frame_num, 24, 4]
+ "smpl_trans": smpl_trans, # [frame_num, 3]
+ "smpl_betas": first_frame_betas, # [10]
+ "size": size, # [3]
+ "frame_info": frame_info, # [frame_num]
+ "pts": torch.cat(collected_lidar_pts, dim=0),
+ "colors": torch.cat(collected_lidar_colors, dim=0),
+ }
+
+ return instance_dict
+
+ def filter_pts_in_boxes(
+ self,
+ seed_pts: Tensor,
+ valid_instances_dict: Dict[int, Dict[str, Tensor]],
+ seed_colors: Tensor = None,
+ seed_time: Tensor = None,
+ ):
+ """
+ This function is used to filter out the points that are inside the bounding boxes of the instances
+ """
+ if DEBUG_PCD:
+ os.makedirs(DEBUG_OUTPUT_DIR, exist_ok=True)
+ export_points_to_ply(
+ seed_pts,
+ seed_colors,
+ save_path=os.path.join(DEBUG_OUTPUT_DIR, "original_seed_pts.ply")
+ )
+ valid_instance_keys = valid_instances_dict.keys()
+
+ inside_mask = torch.zeros_like(seed_pts[:, 0]).bool()
+ for fi in range(self.frame_num):
+ for ins_id in valid_instance_keys:
+ instance_active = self.pixel_source.per_frame_instance_mask[fi, ins_id]
+ if not instance_active:
+ continue
+ # get the pose of the instance at the given frame
+ o2w = self.pixel_source.instances_pose[fi, ins_id].to(seed_pts.device)
+ o_size = self.pixel_source.instances_size[ins_id].to(seed_pts.device)
+ # convert the lidar points to the instance's coordinate system
+ w2o = torch.inverse(o2w)
+ o_pts = transform_points(seed_pts, w2o)
+ # get the mask of the points that are inside the instance's bounding box
+ mask = (
+ (o_pts[:, 0] > -o_size[0] / 2)
+ & (o_pts[:, 0] < o_size[0] / 2)
+ & (o_pts[:, 1] > -o_size[1] / 2)
+ & (o_pts[:, 1] < o_size[1] / 2)
+ & (o_pts[:, 2] > -o_size[2] / 2)
+ & (o_pts[:, 2] < o_size[2] / 2)
+ )
+ inside_mask = inside_mask | mask
+
+ # filter out the points that are inside the bounding boxes
+ seed_pts = seed_pts[~inside_mask]
+ if seed_colors is not None:
+ seed_colors = seed_colors[~inside_mask]
+ if seed_time is not None:
+ seed_time = seed_time[~inside_mask]
+
+ if DEBUG_PCD:
+ export_points_to_ply(
+ seed_pts,
+ seed_colors,
+ save_path=os.path.join(DEBUG_OUTPUT_DIR, "filtered_seed_pts.ply")
+ )
+
+ for fi in range(self.frame_num):
+ if fi % 10 != 0:
+ continue
+ frame_save_dir = os.path.join(DEBUG_OUTPUT_DIR, f"frame_{fi}")
+ os.makedirs(frame_save_dir, exist_ok=True)
+ for ins_id in valid_instances_dict:
+ # print number of points
+ # print(f"Frame {fi}, Instance {ins_id} has {valid_instances_dict[ins_id]['pts'].shape[0]} points")
+ o2w = self.pixel_source.instances_pose[fi, ins_id]
+ pts_in_obj = valid_instances_dict[ins_id]["pts"]
+ # rotate the points back to the world coordinate system
+ pts_in_world = transform_points(pts_in_obj, o2w)
+ export_points_to_ply(
+ pts_in_world,
+ valid_instances_dict[ins_id]["colors"],
+ save_path=os.path.join(frame_save_dir, f"ID={ins_id}.ply")
+ )
+
+ return {
+ "pts": seed_pts,
+ "colors": seed_colors,
+ "time": seed_time
+ }
+
+ def check_pts_visibility(self, pts_xyz):
+ # filter out the lidar points that are not visible from the camera
+ pts_xyz = pts_xyz.to(self.device)
+ valid_mask = torch.zeros_like(pts_xyz[:, 0]).bool()
+ # project lidar points to the image plane
+ for cam in self.pixel_source.camera_data.values():
+ for frame_idx in range(len(cam)):
+ intrinsic_4x4 = torch.nn.functional.pad(
+ cam.intrinsics[frame_idx], (0, 1, 0, 1)
+ )
+ intrinsic_4x4[3, 3] = 1.0
+ lidar2img = (
+ intrinsic_4x4 @ cam.cam_to_worlds[frame_idx].inverse()
+ )
+ projected_points = (
+ lidar2img[:3, :3] @ pts_xyz.T + lidar2img[:3, 3:4]
+ ).T
+ depth = projected_points[:, 2]
+ cam_points = projected_points[:, :2] / (depth.unsqueeze(-1) + 1e-6)
+ current_valid_mask = (
+ (cam_points[:, 0] >= 0)
+ & (cam_points[:, 0] < cam.WIDTH)
+ & (cam_points[:, 1] >= 0)
+ & (cam_points[:, 1] < cam.HEIGHT)
+ & (depth > 0)
+ )
+ valid_mask = valid_mask | current_valid_mask
+ return valid_mask
+
+ def split_train_test(self):
+ if self.data_cfg.pixel_source.test_image_stride != 0:
+ test_timesteps = np.arange(
+ # it makes no sense to have test timesteps before the start timestep
+ self.data_cfg.pixel_source.test_image_stride,
+ self.num_img_timesteps,
+ self.data_cfg.pixel_source.test_image_stride,
+ )
+ else:
+ test_timesteps = []
+ train_timesteps = np.array(
+ [i for i in range(self.num_img_timesteps) if i not in test_timesteps]
+ )
+ logger.info(
+ f"Train timesteps: \n{np.arange(self.start_timestep, self.end_timestep)[train_timesteps]}"
+ )
+ logger.info(
+ f"Test timesteps: \n{np.arange(self.start_timestep, self.end_timestep)[test_timesteps]}"
+ )
+
+ # propagate the train and test timesteps to the train and test indices
+ train_indices, test_indices = [], []
+ for t in range(self.num_img_timesteps):
+ if t in train_timesteps:
+ for cam in range(self.pixel_source.num_cams):
+ train_indices.append(t * self.pixel_source.num_cams + cam)
+ elif t in test_timesteps:
+ for cam in range(self.pixel_source.num_cams):
+ test_indices.append(t * self.pixel_source.num_cams + cam)
+ logger.info(f"Number of train indices: {len(train_indices)}")
+ logger.info(f"Train indices: {train_indices}")
+ logger.info(f"Number of test indices: {len(test_indices)}")
+ logger.info(f"Test indices: {test_indices}")
+
+ # Again, training and testing indices are indices into the full dataset
+ # train_indices are img indices, so the length is num_cams * num_train_timesteps
+ # but train_timesteps are timesteps, so the length is num_train_timesteps (len(unique_train_timestamps))
+ return train_timesteps, test_timesteps, train_indices, test_indices
+
+ def project_lidar_pts_on_images(self, delete_out_of_view_points=True):
+ """
+ Project the lidar points on the images and attribute the color of the nearest pixel to the lidar point.
+
+ Args:
+ delete_out_of_view_points: bool
+ If True, the lidar points that are not visible from the camera will be removed.
+ """
+ for cam in self.pixel_source.camera_data.values():
+ lidar_depth_maps = []
+ for frame_idx in tqdm(
+ range(len(cam)),
+ desc="Projecting lidar pts on images for camera {}".format(cam.cam_name),
+ dynamic_ncols=True
+ ):
+ normed_time = self.pixel_source.normalized_time[frame_idx]
+
+ # get lidar depth on image plane
+ closest_lidar_idx = self.lidar_source.find_closest_timestep(normed_time)
+ lidar_infos = self.lidar_source.get_lidar_rays(closest_lidar_idx)
+ lidar_points = (
+ lidar_infos["lidar_origins"]
+ + lidar_infos["lidar_viewdirs"] * lidar_infos["lidar_ranges"]
+ )
+
+ # project lidar points to the image plane
+ if cam.undistort:
+ new_camera_matrix, _ = cv2.getOptimalNewCameraMatrix(
+ cam.intrinsics[frame_idx].cpu().numpy(),
+ cam.distortions[frame_idx].cpu().numpy(),
+ (cam.WIDTH, cam.HEIGHT),
+ alpha=1,
+ )
+ intrinsic_4x4 = torch.nn.functional.pad(
+ torch.from_numpy(new_camera_matrix), (0, 1, 0, 1)
+ ).to(self.device)
+ else:
+ intrinsic_4x4 = torch.nn.functional.pad(
+ cam.intrinsics[frame_idx], (0, 1, 0, 1)
+ )
+ intrinsic_4x4[3, 3] = 1.0
+ lidar2img = intrinsic_4x4 @ cam.cam_to_worlds[frame_idx].inverse()
+ lidar_points = (
+ lidar2img[:3, :3] @ lidar_points.T + lidar2img[:3, 3:4]
+ ).T # (num_pts, 3)
+
+ depth = lidar_points[:, 2]
+ cam_points = lidar_points[:, :2] / (depth.unsqueeze(-1) + 1e-6) # (num_pts, 2)
+ valid_mask = (
+ (cam_points[:, 0] >= 0)
+ & (cam_points[:, 0] < cam.WIDTH)
+ & (cam_points[:, 1] >= 0)
+ & (cam_points[:, 1] < cam.HEIGHT)
+ & (depth > 0)
+ ) # (num_pts, )
+ depth = depth[valid_mask]
+ _cam_points = cam_points[valid_mask]
+ depth_map = torch.zeros(
+ cam.HEIGHT, cam.WIDTH
+ ).to(self.device)
+ depth_map[
+ _cam_points[:, 1].long(), _cam_points[:, 0].long()
+ ] = depth.squeeze(-1)
+ lidar_depth_maps.append(depth_map)
+
+ # used to filter out the lidar points that are visible from the camera
+ visible_indices = torch.arange(
+ self.lidar_source.num_points, device=self.device
+ )[lidar_infos["lidar_mask"]][valid_mask]
+
+ self.lidar_source.visible_masks[visible_indices] = True
+
+ # attribute the color of the nearest pixel to the lidar point
+ points_color = cam.images[frame_idx][
+ _cam_points[:, 1].long(), _cam_points[:, 0].long()
+ ]
+ self.lidar_source.colors[visible_indices] = points_color
+
+ cam.load_depth(
+ torch.stack(lidar_depth_maps, dim=0).to(self.device).float()
+ )
+
+ if delete_out_of_view_points:
+ self.lidar_source.delete_invisible_pts()
+
+ def get_novel_render_traj(
+ self,
+ traj_types: List[str] = ["front_center_interp"],
+ target_frames: int = 100
+ ) -> Dict[str, torch.Tensor]:
+ """
+ Get multiple novel trajectories of the scene for rendering.
+
+ Args:
+ traj_types: List[str]
+ A list of trajectory types to generate. Options for each type include:
+ - "front_center_interp": Interpolate key frames from the front center camera
+ - "s_curve": S-shaped trajectory using the front three cameras
+ - "three_key_poses": Creates a trajectory using three key poses from different cameras
+ - "lateral_offset": Shift the reference trajectory laterally by a fixed offset
+ target_frames: int
+ The total number of frames for each novel trajectory
+
+ Returns:
+ Dict[str, torch.Tensor]: A dictionary where keys are trajectory types and values
+ are the generated novel trajectories, each of shape (target_frames, 4, 4)
+ """
+ per_cam_poses = {}
+ for cam_id in self.pixel_source.camera_list:
+ per_cam_poses[cam_id] = self.pixel_source.camera_data[cam_id].cam_to_worlds
+
+ novel_trajs = {}
+ for traj_type in traj_types:
+ novel_trajs[traj_type] = get_interp_novel_trajectories(
+ self.type,
+ self.scene_idx,
+ per_cam_poses,
+ traj_type,
+ target_frames
+ )
+
+ return novel_trajs
+
+ def prepare_novel_view_render_data(self, traj: torch.Tensor, cam_id: int = None) -> list:
+ """
+ Prepare all necessary elements for novel view rendering.
+
+ Args:
+ traj (torch.Tensor): Novel view trajectory, shape (N, 4, 4)
+ cam_id (int, optional): Specific camera ID to use for intrinsics and image size.
+
+ Returns:
+ list: List of dicts, each containing elements required for rendering a single frame:
+ - cam_infos: Camera information (extrinsics, intrinsics, image dimensions)
+ - image_infos: Image-related information (indices, normalized time, viewdirs, etc.)
+ """
+ # Call the PixelSource's method
+ return self.pixel_source.prepare_novel_view_render_data(self.type, traj, cam_id)
diff --git a/worldbench/third_party/drivestudio/datasets/kitti/kitti_human_utils.py b/worldbench/third_party/drivestudio/datasets/kitti/kitti_human_utils.py
new file mode 100644
index 0000000..3bc1439
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/kitti/kitti_human_utils.py
@@ -0,0 +1,218 @@
+from typing import List
+import os
+import cv2
+import json
+import logging
+import numpy as np
+from tqdm import tqdm
+
+from utils.geometry import (
+ get_corners,
+ project_camera_points_to_image
+)
+from .kitti_sourceloader import (
+ SMPLNODE_CLASSES,
+ OPENCV2DATASET,
+ AVAILABLE_CAM_LIST,
+)
+
+logger = logging.getLogger()
+
+CAMERA_LIST = AVAILABLE_CAM_LIST
+
+def project_human_boxes(
+ scene_dir: str,
+ camera_list: List[int],
+ save_temp=True,
+ verbose=False,
+ narrow_width_ratio=0.2,
+ fps=12
+):
+ """Project human boxes to 2D image space and save the results to pkl files
+
+ Args:
+ scene_dir: str, path to the scene directory
+ camera_list: List[int], a list of camera ids to be processed
+ save_temp: bool, whether to save the intermediate results
+ verbose: bool, whether to visualize the projected boxes
+ narrow_width_ratio: sometimes the projected boxes are too wide
+ we can narrow them with this ratio to get more accurate results
+ fps: int, FPS for the visualization video
+
+ Returns:
+ collector_all: dict, a dictionary containing the projected boxes for each camera
+ """
+ # check if the necessary directories exist
+ images_dir = f'{scene_dir}/images'
+ poses_dir = f'{scene_dir}/ego_pose'
+ extrinsics_dir = f'{scene_dir}/extrinsics'
+ intrinsics_dir = f'{scene_dir}/intrinsics'
+ instances_dir = f'{scene_dir}/instances'
+ valid_paths = [images_dir, poses_dir, extrinsics_dir, intrinsics_dir, instances_dir]
+ for path in valid_paths:
+ assert os.path.exists(path), \
+ f"Path {path} does not exist, you need to run waymo preprocess to generate the necessary files"
+
+ # create directories for saving the results
+ save_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks'
+ if not os.path.exists(save_dir):
+ os.makedirs(save_dir)
+
+ if verbose:
+ # create directories for saving the intermediate visualization results
+ video_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks/vis'
+ if not os.path.exists(video_dir):
+ os.makedirs(video_dir)
+ per_human_img_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks/vis/images'
+ if not os.path.exists(per_human_img_dir):
+ os.makedirs(per_human_img_dir)
+
+ # load instances and frame infos
+ frame_infos = json.load(
+ open(f'{instances_dir}/frame_instances.json')
+ )
+ instances_meta = json.load(
+ open(f'{instances_dir}/instances_info.json')
+ )
+
+ collector_all = {}
+ # iterate over each camera
+ for cam_id in camera_list:
+ # check if already processed
+ pkl_path = os.path.join(save_dir, f"{cam_id}.pkl")
+ if os.path.exists(pkl_path):
+ collector_all[cam_id] = json.load(open(pkl_path))
+ logger.info(f"Results for camera {cam_id} already exists at {pkl_path}")
+ continue
+
+ if verbose:
+ per_cam_vis_dir = os.path.join(per_human_img_dir, f"{cam_id}")
+ if not os.path.exists(per_cam_vis_dir):
+ os.makedirs(per_cam_vis_dir)
+
+ collector = {}
+ frames = []
+ for frame_id, frame_ins_list in frame_infos.items():
+ frame_id = int(frame_id)
+
+ # define empty instance collector for each frame
+ frame_collector = {
+ "gt_bbox": [],
+ "extra_data": {
+ "gt_track_id": [],
+ "gt_class": [],
+ }
+ }
+
+ # load extrinsic
+ cam_to_ego = np.loadtxt(os.path.join(extrinsics_dir, f"{cam_id}.txt"))
+ cam_to_ego = cam_to_ego @ OPENCV2DATASET
+ ego_to_world = np.loadtxt(os.path.join(poses_dir, f"{str(frame_id).zfill(3)}.txt"))
+ cam2world = ego_to_world @ cam_to_ego
+
+ # load intrinsic
+ Ks = np.loadtxt(os.path.join(intrinsics_dir, f"{cam_id}.txt"))
+ fx, fy, cx, cy = Ks[0], Ks[1], Ks[2], Ks[3]
+ # k1, k2, p1, p2, k3 = Ks[4], Ks[5], Ks[6], Ks[7], Ks[8]
+ intrinsic = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float32)
+
+ # load image
+ ori_image = cv2.imread(
+ os.path.join(images_dir, f"{str(frame_id).zfill(3)}_{cam_id}.jpg")
+ )
+ image = ori_image.copy()
+ H, W = image.shape[:2]
+
+ # iterate over each instance in the frame
+ if len(frame_ins_list) > 0:
+ # if there are pedestrians in this frame, project the boxes to the image
+ image_plotted = image.copy()
+ for instance_id in frame_ins_list:
+ ins = instances_meta[str(instance_id)]
+
+ if ins["class_name"] not in SMPLNODE_CLASSES:
+ continue
+
+ ins_anno = ins["frame_annotations"]
+ index = ins_anno['frame_idx'].index(frame_id)
+ obj_to_world = np.array(ins_anno['obj_to_world'][index])
+ l, w, h = ins_anno['box_size'][index]
+
+ # get box corners in object space
+ corners = get_corners(l, w, h)
+ # transform box corners to world space
+ corners_world = obj_to_world[:3, :3] @ corners + obj_to_world[:3, 3:4]
+ # transform box corners to image space
+ world2cam = np.linalg.inv(cam2world)
+ corners_cam = world2cam[:3, :3] @ corners_world + world2cam[:3, 3:4]
+ cam_points, depth = project_camera_points_to_image(corners_cam.T, intrinsic)
+
+ x_min, y_min = np.min(cam_points, axis=0)
+ x_max, y_max = np.max(cam_points, axis=0)
+ # clip left and right with this ratio
+ if narrow_width_ratio > 0.:
+ length = x_max - x_min
+ x_min += length * narrow_width_ratio
+ x_max -= length * narrow_width_ratio
+
+ # clip the box to the image
+ original_area = (x_max - x_min) * (y_max - y_min)
+ x_min, x_max = np.clip(x_min, 0, W), np.clip(x_max, 0, W)
+ y_min, y_max = np.clip(y_min, 0, H), np.clip(y_max, 0, H)
+ new_area = (x_max - x_min) * (y_max - y_min)
+
+ # filter out boxes that are too small or too large
+ behind = depth.max() < 0
+ too_small = new_area < W * H * (0.03)**2
+ too_large = new_area > W * H / 1.1
+ too_far = np.linalg.norm(obj_to_world[:3, 3] - cam2world[:3, 3]) > 40
+ clip_large = new_area / original_area < 1/3
+ if too_small or too_large or clip_large or behind or too_far:
+ continue
+
+ # gt box on image
+ gt_box = [x_min, y_min, x_max - x_min, y_max - y_min]
+
+ # save the projected box to the collector
+ frame_collector["gt_bbox"].append(gt_box)
+ frame_collector["extra_data"]["gt_track_id"].append(instance_id)
+ frame_collector["extra_data"]["gt_class"].append([0])
+
+ if verbose:
+ # visualize the projected boxes of ONE instance
+ raw_image = cv2.rectangle(
+ ori_image.copy(), (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 2
+ )
+ raw_image_path = os.path.join(per_cam_vis_dir, f"{frame_id}_{instance_id}.jpg")
+ cv2.imwrite(raw_image_path, raw_image)
+
+ # add this instance to the image
+ image_plotted = cv2.rectangle(
+ image, (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 2
+ )
+
+ if verbose:
+ frames.append(image_plotted)
+ else:
+ # if no instance in this frame, just save the original image
+ if verbose:
+ frames.append(ori_image)
+
+ collector[frame_id] = frame_collector
+
+ if verbose:
+ height, width = frames[0].shape[:2]
+ output_path = os.path.join(video_dir, f"cam_{cam_id}.mp4")
+ fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+ out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+ for frame in tqdm(frames, desc=f"Writing video for camera {cam_id}"):
+ out.write(frame)
+ out.release()
+
+ if save_temp:
+ # save collector to pkl
+ json.dump(collector, open(pkl_path, "w"))
+
+ collector_all[cam_id] = collector
+
+ return collector_all
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/kitti/kitti_preprocess.py b/worldbench/third_party/drivestudio/datasets/kitti/kitti_preprocess.py
new file mode 100644
index 0000000..a80ffe1
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/kitti/kitti_preprocess.py
@@ -0,0 +1,411 @@
+import json
+import os
+from typing import Dict, List, Tuple
+
+import cv2
+import numpy as np
+import pykitti
+from PIL import Image
+
+from datasets.kitti.trackletparser import parseXML
+from datasets.tools.multiprocess_utils import track_parallel_progress
+from utils.geometry import get_corners, project_camera_points_to_image
+from utils.visualization import color_mapper, dump_3d_bbox_on_image
+
+KITTI_LABELS = [
+ 'Car', 'Van', 'Truck', 'Pedestrian', 'Person_sitting', 'Cyclist', 'Tram', 'Misc'
+]
+
+KITTI_NONRIGID_DYNAMIC_CLASSES = [
+ 'Pedestrian', 'Person_sitting', 'Cyclist'
+]
+
+KITTI_RIGID_DYNAMIC_CLASSES = [
+ 'Car', 'Van', 'Truck', 'Tram'
+]
+
+KITTI_DYNAMIC_CLASSES = KITTI_NONRIGID_DYNAMIC_CLASSES + KITTI_RIGID_DYNAMIC_CLASSES
+
+class KittiProcessor(object):
+ """Process KITTI dataset."""
+
+ def __init__(
+ self,
+ load_dir: str,
+ save_dir: str,
+ process_keys: List[str] = [
+ "images",
+ "lidar",
+ "calib",
+ "dynamic_masks",
+ "objects"
+ ],
+ prefix: str = "2011_09_26",
+ process_id_list: List[str] = None,
+ workers: int = 64,
+ ):
+ self.process_id_list = process_id_list
+ self.process_keys = process_keys
+ self.HW = (375, 1242)
+ print("will process keys: ", self.process_keys)
+
+ self.cam_list = [
+ "CAM_LEFT", # "xxx_0.jpg"
+ "CAM_RIGHT" # "xxx_1.jpg"
+ ]
+
+ self.split_dir = os.path.join(load_dir, prefix)
+ self.load_dir = load_dir
+ self.save_dir = save_dir
+ self.workers = int(workers)
+ self.create_folder()
+
+ def convert(self):
+ """Convert action."""
+ print("Start converting ...")
+ if self.process_id_list is None:
+ id_list = range(len(self))
+ else:
+ id_list = self.process_id_list
+ track_parallel_progress(self.convert_one, id_list, self.workers)
+ print("\nFinished ...")
+
+ def get_kitti_data(self, basedir):
+ date = basedir.split("/")[-2]
+ drive = basedir.split("/")[-1].split('_')[-2]
+ data = pykitti.raw(self.load_dir, date, drive)
+ return data
+
+ def convert_one(self, scene_name: str):
+ """Convert action for single file."""
+ basedir = os.path.join(self.split_dir, scene_name)
+ kitti_data = self.get_kitti_data(basedir)
+ if "images" in self.process_keys:
+ self.save_image(kitti_data, scene_name)
+ if "calib" in self.process_keys:
+ self.save_calib(kitti_data, scene_name)
+ if "pose" in self.process_keys:
+ self.save_pose(kitti_data, scene_name)
+ if "lidar" in self.process_keys:
+ self.save_lidar(kitti_data, scene_name)
+
+ tracklet_file = os.path.join(basedir, 'tracklet_labels.xml')
+ tracklets = parseXML(tracklet_file)
+ if "dynamic_masks" in self.process_keys:
+ self.save_dynamic_mask(kitti_data, tracklets, scene_name, class_valid='all')
+ self.save_dynamic_mask(kitti_data, tracklets, scene_name, class_valid='human')
+ self.save_dynamic_mask(kitti_data, tracklets, scene_name, class_valid='vehicle')
+
+ # process annotated objects
+ if "objects" in self.process_keys:
+ instances_info, frame_instances = self.save_objects(kitti_data, tracklets)
+
+ # Save instances info and frame instances
+ instances_info_save_path = f"{self.save_dir}/{str(scene_name).zfill(3)}/instances"
+ os.makedirs(instances_info_save_path, exist_ok=True)
+ with open(f"{instances_info_save_path}/instances_info.json", "w") as fp:
+ json.dump(instances_info, fp, indent=4)
+ with open(f"{instances_info_save_path}/frame_instances.json", "w") as fp:
+ json.dump(frame_instances, fp, indent=4)
+
+ if "objects_vis" in self.process_keys:
+ objects_vis_path = f"{self.save_dir}/{str(scene_name).zfill(3)}/instances/debug_vis"
+ if not os.path.exists(objects_vis_path):
+ os.makedirs(objects_vis_path)
+ self.visualize_dynamic_objects(kitti_data, scene_name, objects_vis_path, instances_info, frame_instances)
+
+ def __len__(self):
+ """Length of the filename list."""
+ return len(self.process_id_list) if self.process_id_list else 1
+
+ def save_image(self, kitti_data, scene_name: str):
+ for frame_idx in range(len(kitti_data)):
+ images = kitti_data.get_rgb(frame_idx)
+ # Save left and right images
+ images[0].save(f"{self.save_dir}/{scene_name}/images/{str(frame_idx).zfill(3)}_0.jpg")
+ images[1].save(f"{self.save_dir}/{scene_name}/images/{str(frame_idx).zfill(3)}_1.jpg")
+
+ def save_calib(self, kitti_data, scene_name: str):
+ # NOTE: assume ego vehicle is the same as velodyne frame
+ np.savetxt(
+ f"{self.save_dir}/{scene_name}/extrinsics/0.txt",
+ np.linalg.inv(kitti_data.calib.T_cam2_velo)
+ )
+ np.savetxt(
+ f"{self.save_dir}/{scene_name}/extrinsics/1.txt",
+ np.linalg.inv(kitti_data.calib.T_cam3_velo)
+ )
+
+ cam2_Ks = kitti_data.calib.K_cam2
+ cam3_Ks = kitti_data.calib.K_cam3
+ # fx, fy, cx, cy, p1, p2, k1, k2, k3
+ Ks_left = np.array([cam2_Ks[0, 0], cam2_Ks[1, 1], cam2_Ks[0, 2], cam2_Ks[1, 2], 0, 0, 0, 0, 0])
+ Ks_right = np.array([cam3_Ks[0, 0], cam3_Ks[1, 1], cam3_Ks[0, 2], cam3_Ks[1, 2], 0, 0, 0, 0, 0])
+ np.savetxt(
+ f"{self.save_dir}/{scene_name}/intrinsics/0.txt",
+ Ks_left
+ )
+ np.savetxt(
+ f"{self.save_dir}/{scene_name}/intrinsics/1.txt",
+ Ks_right
+ )
+
+ def save_pose(self, kitti_data, scene_name: str):
+ # NOTE: we assume the ego pose is the same as the velodyne pose
+ for frame_idx in range(len(kitti_data)):
+ imu2world = kitti_data.oxts[frame_idx].T_w_imu
+ imu2velo = kitti_data.calib.T_velo_imu
+ velo2world = imu2world @ imu2velo
+ np.savetxt(
+ f"{self.save_dir}/{scene_name}/ego_pose/{str(frame_idx).zfill(3)}.txt",
+ velo2world
+ )
+
+ def save_lidar(self, kitti_data, scene_name: str):
+ for frame_idx in range(len(kitti_data)):
+ # Points are already in ego frame (velodyne frame), so we don't need to transform them
+ points = kitti_data.get_velo(frame_idx)
+
+ # Save lidar points
+ lidar_save_path = f"{self.save_dir}/{scene_name}/lidar/{str(frame_idx).zfill(3)}.bin"
+ points.astype(np.float32).tofile(lidar_save_path)
+
+ def save_dynamic_mask(self, kitti_data, tracklets, scene_name: str, class_valid: str):
+ # Define valid classes based on class_valid parameter
+ if class_valid == 'all':
+ valid_classes = KITTI_DYNAMIC_CLASSES
+ elif class_valid == 'human':
+ valid_classes = KITTI_NONRIGID_DYNAMIC_CLASSES
+ elif class_valid == 'vehicle':
+ valid_classes = KITTI_RIGID_DYNAMIC_CLASSES
+ else:
+ raise ValueError("Invalid class_valid parameter")
+
+ for frame_idx in range(len(kitti_data)):
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ img_path = f"{self.save_dir}/{scene_name}/images/{str(frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+
+ img = cv2.imread(img_path)
+ dynamic_mask = np.zeros(img.shape[:2], dtype=np.float32)
+
+ for object in tracklets:
+ if frame_idx < object.firstFrame or frame_idx >= object.firstFrame + object.nFrames:
+ continue
+
+ if object.objectType not in valid_classes:
+ continue
+
+ obj_step = frame_idx - object.firstFrame
+ t_obj = np.array(object.trans[obj_step])
+ roty = np.array(object.rots[obj_step][2])
+
+ # Create 3D bounding box
+ h, w, l = object.size
+ corners_3d_local = get_corners(l=l, w=w, h=h)
+
+ # Rotate and translate 3D bounding box
+ c = np.cos(roty)
+ s = np.sin(roty)
+ tx, ty, tz = t_obj
+ tz += h / 2 # NOTE: objects are annotated at the bottom center in KITTI dataset
+ o2v = np.array([
+ [ c, -s, 0, tx],
+ [ s, c, 0, ty],
+ [ 0, 0, 1, tz],
+ [ 0, 0, 0, 1]
+ ])
+ corners_3d_velo = o2v[:3, :3] @ corners_3d_local + o2v[:3, 3:4]
+
+ # Project 3D bounding box to 2D
+ if cam_name == "CAM_LEFT":
+ velo2cam = kitti_data.calib.T_cam2_velo
+ intrinsics_3X3 = kitti_data.calib.K_cam2
+ elif cam_name == "CAM_RIGHT":
+ velo2cam = kitti_data.calib.T_cam3_velo
+ intrinsics_3X3 = kitti_data.calib.K_cam3
+ corners_3d_cam = velo2cam[:3, :3] @ corners_3d_velo + velo2cam[:3, 3:4]
+ corners_2d, _ = project_camera_points_to_image(corners_3d_cam.T, intrinsics_3X3)
+
+ # Check if the all object's corners are in the image
+ # NOTE: we use strict visibility check here, requiring all corners to be visible
+ in_image = np.all(corners_2d[0, :] >= 0) & np.all(corners_2d[0, :] < img.shape[1]) & \
+ np.all(corners_2d[1, :] >= 0) & np.all(corners_2d[1, :] < img.shape[0])
+ if not in_image:
+ continue
+
+ # Fill the mask
+ u, v = corners_2d[0, :].astype(np.int32), corners_2d[1, :].astype(np.int32)
+ u = np.clip(u, 0, img.shape[1] - 1)
+ v = np.clip(v, 0, img.shape[0] - 1)
+
+ if u.max() - u.min() == 0 or v.max() - v.min() == 0:
+ continue
+
+ xy = (u.min(), v.min())
+ width = u.max() - u.min()
+ height = v.max() - v.min()
+
+ dynamic_mask[
+ int(xy[1]): int(xy[1] + height),
+ int(xy[0]): int(xy[0] + width)
+ ] = np.maximum(
+ dynamic_mask[
+ int(xy[1]): int(xy[1] + height),
+ int(xy[0]): int(xy[0] + width)
+ ],
+ 1
+ )
+
+ # Save the mask
+ dynamic_mask = np.clip((dynamic_mask > 0.) * 255, 0, 255).astype(np.uint8)
+ dynamic_mask = Image.fromarray(dynamic_mask, "L")
+ dynamic_mask_path = f"{self.save_dir}/{scene_name}/dynamic_masks/{class_valid}/{str(frame_idx).zfill(3)}_{cam_idx}.png"
+ dynamic_mask.save(dynamic_mask_path)
+
+ def save_objects(self, kitti_data, tracklets) -> Tuple[Dict, Dict]:
+ instances_info = {}
+ frame_instances = {i: [] for i in range(len(kitti_data))}
+ for obj_id, object in enumerate(tracklets):
+ if object.objectType not in KITTI_DYNAMIC_CLASSES:
+ continue
+
+ instances_info[obj_id] = {
+ "id": obj_id,
+ "class_name": object.objectType,
+ "frame_annotations": {
+ "frame_idx": [],
+ "obj_to_world": [],
+ "box_size": []
+ }
+ }
+
+ for frame_idx in range(object.firstFrame, object.firstFrame + object.nFrames):
+ if frame_idx >= len(kitti_data):
+ break
+
+ obj_step = frame_idx - object.firstFrame
+ t_obj = np.array(object.trans[obj_step])
+ roty = np.array(object.rots[obj_step][2])
+
+ # Create transformation matrix
+ c = np.cos(roty)
+ s = np.sin(roty)
+ tx, ty, tz = t_obj
+ h, _, _ = object.size
+ tz += h / 2 # NOTE: objects are annotated at the bottom center in KITTI dataset
+ o2v = np.array([
+ [ c, -s, 0, tx],
+ [ s, c, 0, ty],
+ [ 0, 0, 1, tz],
+ [ 0, 0, 0, 1]])
+
+ # Transform to world coordinates
+ imu2world = kitti_data.oxts[frame_idx].T_w_imu
+ imu2velo = kitti_data.calib.T_velo_imu
+ velo2world = imu2world @ imu2velo
+ obj_to_world = velo2world @ o2v
+
+ # convert hwl to lwh
+ lwh = [object.size[2], object.size[1], object.size[0]]
+
+ instances_info[obj_id]["frame_annotations"]["frame_idx"].append(frame_idx)
+ instances_info[obj_id]["frame_annotations"]["obj_to_world"].append(obj_to_world.tolist())
+ instances_info[obj_id]["frame_annotations"]["box_size"].append(lwh)
+
+ frame_instances[frame_idx].append(obj_id)
+
+ # Correct ID mapping
+ id_map = {}
+ for i, (k, v) in enumerate(instances_info.items()):
+ id_map[v["id"]] = i
+
+ # Update keys in instances_info
+ new_instances_info = {}
+ for k, v in instances_info.items():
+ new_instances_info[id_map[v["id"]]] = v
+
+ # Update keys in frame_instances
+ new_frame_instances = {}
+ for k, v in frame_instances.items():
+ new_frame_instances[k] = [id_map[i] for i in v]
+
+ return new_instances_info, new_frame_instances
+
+ def visualize_dynamic_objects(self, kitti_data, scene_name:str, objects_vis_path: str, instances_info: Dict, frame_instances: Dict):
+ for frame_idx, obj_ids in frame_instances.items():
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ # Load the image
+ img_path = os.path.join(self.save_dir, scene_name, 'images', f'{frame_idx:03d}_{cam_idx}.jpg')
+ canvas = np.array(Image.open(img_path))
+
+ if len(obj_ids) == 0:
+ img_plotted = canvas
+ else:
+ lstProj2d = []
+ color_list = []
+ for obj_id in obj_ids:
+ obj_info = instances_info[obj_id]
+ frame_ann_idx = obj_info['frame_annotations']['frame_idx'].index(frame_idx)
+
+ # Get object pose and size
+ obj_to_world = np.array(obj_info['frame_annotations']['obj_to_world'][frame_ann_idx])
+ box_size = np.array(obj_info['frame_annotations']['box_size'][frame_ann_idx])
+
+ # Create 3D bounding box in object coordinates
+ l, w, h = box_size
+ corners = get_corners(l=l, w=w, h=h)
+
+ # Transform box to world coordinates
+ corners_3d_world = obj_to_world[:3, :3] @ corners + obj_to_world[:3, 3:4]
+
+ # world to velo
+ imu2world = kitti_data.oxts[frame_idx].T_w_imu
+ imu2velo = kitti_data.calib.T_velo_imu
+ world2velo = np.linalg.inv(imu2world @ imu2velo)
+ corners_3d_velo = world2velo[:3, :3] @ corners_3d_world + world2velo[:3, 3:4]
+
+ # velo to cam
+ if cam_name == "CAM_LEFT":
+ velo2cam = kitti_data.calib.T_cam2_velo
+ intrinsics_3X3 = kitti_data.calib.K_cam2
+ elif cam_name == "CAM_RIGHT":
+ velo2cam = kitti_data.calib.T_cam3_velo
+ intrinsics_3X3 = kitti_data.calib.K_cam3
+ corners_3d_cam = velo2cam[:3, :3] @ corners_3d_velo + velo2cam[:3, 3:4]
+
+ # Project to image plane
+ corners_2d, _ = project_camera_points_to_image(corners_3d_cam.T, intrinsics_3X3)
+
+ # Check if the all object's corners are in the image
+ in_image = np.all(corners_2d[0, :] >= 0) & np.all(corners_2d[0, :] < canvas.shape[1]) & \
+ np.all(corners_2d[1, :] >= 0) & np.all(corners_2d[1, :] < canvas.shape[0])
+ if in_image:
+ projected_points2d = corners_2d[:2, :].T
+ lstProj2d.append(projected_points2d)
+ color_list.append(color_mapper(str(obj_id)))
+
+ # Draw all bounding boxes at once
+ img_plotted = dump_3d_bbox_on_image(coords=np.array(lstProj2d), img=canvas, color=color_list)
+
+ # Save the visualized image
+ out_path = os.path.join(objects_vis_path, f'{frame_idx:06d}_{cam_idx}.jpg')
+ Image.fromarray(img_plotted).save(out_path)
+
+ def create_folder(self):
+ """Create folder for data preprocessing."""
+ if self.process_id_list is None:
+ id_list = range(len(self))
+ else:
+ id_list = self.process_id_list
+ for scene_name in id_list:
+ os.makedirs(f"{self.save_dir}/{scene_name}/images", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{scene_name}/extrinsics", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{scene_name}/intrinsics", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{scene_name}/sky_masks", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{scene_name}/ego_pose", exist_ok=True)
+ if "lidar" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{scene_name}/lidar", exist_ok=True)
+ if "dynamic_masks" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{scene_name}/dynamic_masks/all", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{scene_name}/dynamic_masks/human", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{scene_name}/dynamic_masks/vehicle", exist_ok=True)
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/kitti/kitti_sourceloader.py b/worldbench/third_party/drivestudio/datasets/kitti/kitti_sourceloader.py
new file mode 100644
index 0000000..d4396a9
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/kitti/kitti_sourceloader.py
@@ -0,0 +1,469 @@
+import json
+import logging
+import os
+from typing import Dict
+
+import joblib
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from pytorch3d.transforms import matrix_to_quaternion
+from torch import Tensor
+from tqdm import tqdm, trange
+
+from datasets.base.scene_dataset import ModelType
+from datasets.base.lidar_source import SceneLidarSource
+from datasets.base.pixel_source import CameraData, ScenePixelSource
+
+logger = logging.getLogger()
+
+# define each class's node type
+OBJECT_CLASS_NODE_MAPPING = {
+ # Rigid objects (vehicles)
+ 'Car': ModelType.RigidNodes,
+ 'Van': ModelType.RigidNodes,
+ 'Truck': ModelType.RigidNodes,
+ 'Tram': ModelType.RigidNodes,
+
+ # Humans (SMPL model)
+ 'Pedestrian': ModelType.SMPLNodes,
+ 'Person_sitting': ModelType.SMPLNodes,
+
+ # Potentially deformable objects
+ 'Cyclist': ModelType.DeformableNodes,
+}
+SMPLNODE_CLASSES = ["Pedestrian", "Person_sitting"]
+
+# OpenCV to Dataset coordinate transformation
+# opencv coordinate system: x right, y down, z front
+# KITTI coordinate system: x right, y down, z front (same as OpenCV)
+OPENCV2DATASET = np.array(
+ [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]
+)
+
+# KITTI Camera List:
+# 0: CAM_LEFT
+# 1: CAM_RIGHT
+AVAILABLE_CAM_LIST = [0, 1]
+
+class KITTICameraData(CameraData):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ def load_calibrations(self):
+ """
+ Load the camera intrinsics, extrinsics, timestamps, etc.
+ Compute the camera-to-world matrices, ego-to-world matrices, etc.
+ """
+ # load camera intrinsics
+ # 1d Array of [f_u, f_v, c_u, c_v, k{1, 2}, p{1, 2}, k{3}].
+ # ====!! we did not use distortion parameters for simplicity !!====
+ # to be improved!!
+ intrinsic = np.loadtxt(
+ os.path.join(self.data_path, "intrinsics", f"{self.cam_id}.txt")
+ )
+ fx, fy, cx, cy = intrinsic[0], intrinsic[1], intrinsic[2], intrinsic[3]
+ k1, k2, p1, p2, k3 = intrinsic[4], intrinsic[5], intrinsic[6], intrinsic[7], intrinsic[8]
+ # scale intrinsics w.r.t. load size
+ fx, fy = (
+ fx * self.load_size[1] / self.original_size[1],
+ fy * self.load_size[0] / self.original_size[0],
+ )
+ cx, cy = (
+ cx * self.load_size[1] / self.original_size[1],
+ cy * self.load_size[0] / self.original_size[0],
+ )
+ _intrinsics = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
+ _distortions = np.array([k1, k2, p1, p2, k3])
+
+ # load camera extrinsics
+ cam_to_ego = np.loadtxt(
+ os.path.join(self.data_path, "extrinsics", f"{self.cam_id}.txt")
+ )
+ cam_to_ego = cam_to_ego @ OPENCV2DATASET
+
+ # compute per-image poses and intrinsics
+ cam_to_worlds, ego_to_worlds = [], []
+ intrinsics, distortions = [], []
+
+ # we tranform the camera poses w.r.t. the first timestep to make the translation vector of
+ # the first ego pose as the origin of the world coordinate system.
+ ego_to_world_start = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{self.start_timestep:03d}.txt")
+ )
+ for t in range(self.start_timestep, self.end_timestep):
+ ego_to_world_current = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{t:03d}.txt")
+ )
+ # compute ego_to_world transformation
+ ego_to_world = np.linalg.inv(ego_to_world_start) @ ego_to_world_current
+ ego_to_worlds.append(ego_to_world)
+ cam2world = ego_to_world @ cam_to_ego
+ cam_to_worlds.append(cam2world)
+ intrinsics.append(_intrinsics)
+ distortions.append(_distortions)
+
+ self.intrinsics = torch.from_numpy(np.stack(intrinsics, axis=0)).float()
+ self.distortions = torch.from_numpy(np.stack(distortions, axis=0)).float()
+ self.cam_to_worlds = torch.from_numpy(np.stack(cam_to_worlds, axis=0)).float()
+
+ @classmethod
+ def get_camera2worlds(cls, data_path: str, cam_id: str, start_timestep: int, end_timestep: int) -> torch.Tensor:
+ """
+ Returns camera-to-world matrices for the specified camera and time range.
+
+ Args:
+ data_path (str): Path to the dataset.
+ cam_id (str): Camera ID.
+ start_timestep (int): Start timestep.
+ end_timestep (int): End timestep.
+
+ Returns:
+ torch.Tensor: Camera-to-world matrices of shape (num_frames, 4, 4).
+ """
+ # Load the camera-to-ego transformation (this is constant for KITTI)
+ cam_to_ego = np.loadtxt(
+ os.path.join(data_path, "extrinsics", f"{cam_id}.txt")
+ )
+ cam_to_ego = cam_to_ego @ OPENCV2DATASET
+
+ # Load the initial ego pose to align the world coordinate system
+ ego_to_world_start = np.loadtxt(
+ os.path.join(data_path, "ego_pose", f"{start_timestep:03d}.txt")
+ )
+
+ cam_to_worlds = []
+
+ for t in range(start_timestep, end_timestep):
+ # Load ego pose for each timestep
+ ego_to_world_current = np.loadtxt(
+ os.path.join(data_path, "ego_pose", f"{t:03d}.txt")
+ )
+
+ # Compute ego_to_world transformation relative to the start
+ ego_to_world = np.linalg.inv(ego_to_world_start) @ ego_to_world_current
+
+ # Compute camera-to-world transformation
+ cam2world = ego_to_world @ cam_to_ego
+
+ # No need to apply OPENCV2DATASET as KITTI uses the same coordinate system
+ cam_to_worlds.append(cam2world)
+
+ return torch.from_numpy(np.stack(cam_to_worlds, axis=0)).float()
+
+class KITTIPixelSource(ScenePixelSource):
+ def __init__(
+ self,
+ dataset_name: str,
+ pixel_data_config: OmegaConf,
+ data_path: str,
+ start_timestep: int,
+ end_timestep: int,
+ device: torch.device = torch.device("cpu"),
+ ):
+ super().__init__(dataset_name, pixel_data_config, device=device)
+ self.data_path = data_path
+ self.start_timestep = start_timestep
+ self.end_timestep = end_timestep
+ self.load_data()
+
+ def load_cameras(self):
+ self._timesteps = torch.arange(self.start_timestep, self.end_timestep)
+ self.register_normalized_timestamps()
+
+ for idx, cam_id in enumerate(self.camera_list):
+ logger.info(f"Loading camera {cam_id}")
+ camera = KITTICameraData(
+ dataset_name=self.dataset_name,
+ data_path=self.data_path,
+ cam_id=cam_id,
+ start_timestep=self.start_timestep,
+ end_timestep=self.end_timestep,
+ load_dynamic_mask=self.data_cfg.load_dynamic_mask,
+ load_sky_mask=self.data_cfg.load_sky_mask,
+ downscale_when_loading=self.data_cfg.downscale_when_loading[idx],
+ undistort=self.data_cfg.undistort,
+ buffer_downscale=self.buffer_downscale,
+ device=self.device,
+ )
+ camera.load_time(self.normalized_time)
+ unique_img_idx = torch.arange(len(camera), device=self.device) * len(self.camera_list) + idx
+ camera.set_unique_ids(
+ unique_cam_idx = idx,
+ unique_img_idx = unique_img_idx
+ )
+ logger.info(f"Camera {camera.cam_name} loaded.")
+ self.camera_data[cam_id] = camera
+
+ def load_objects(self):
+ """
+ get ground truth bounding boxes of the dynamic objects
+
+ instances_info = {
+ "0": # simplified instance id
+ {
+ "id": str,
+ "class_name": str,
+ "frame_annotations": {
+ "frame_idx": List,
+ "obj_to_world": List,
+ "box_size": List,
+ },
+ ...
+ }
+ frame_instances = {
+ "0": # frame idx
+ List[int] # list of simplified instance ids
+ ...
+ }
+ """
+ instances_info_path = os.path.join(self.data_path, "instances", "instances_info.json")
+ frame_instances_path = os.path.join(self.data_path, "instances", "frame_instances.json")
+ with open(instances_info_path, "r") as f:
+ instances_info = json.load(f)
+ with open(frame_instances_path, "r") as f:
+ frame_instances = json.load(f)
+ # get pose of each instance at each frame
+ # shape (num_frames, num_instances, 4, 4)
+ num_instances = len(instances_info)
+ num_full_frames = len(frame_instances)
+ instances_pose = np.zeros((num_full_frames, num_instances, 4, 4))
+ instances_size = np.zeros((num_full_frames, num_instances, 3))
+ instances_true_id = np.arange(num_instances)
+ instances_model_types = np.ones(num_instances) * -1
+
+ ego_to_world_start = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{self.start_timestep:03d}.txt")
+ )
+ for k, v in instances_info.items():
+ instances_model_types[int(k)] = OBJECT_CLASS_NODE_MAPPING[v["class_name"]]
+ for frame_idx, obj_to_world, box_size in zip(v["frame_annotations"]["frame_idx"], v["frame_annotations"]["obj_to_world"], v["frame_annotations"]["box_size"]):
+ # the first ego pose as the origin of the world coordinate system.
+ obj_to_world = np.array(obj_to_world).reshape(4, 4)
+ obj_to_world = np.linalg.inv(ego_to_world_start) @ obj_to_world
+ instances_pose[frame_idx, int(k)] = np.array(obj_to_world)
+ instances_size[frame_idx, int(k)] = np.array(box_size)
+
+ # get frame valid instances
+ # shape (num_frames, num_instances)
+ per_frame_instance_mask = np.zeros((num_full_frames, num_instances))
+ for frame_idx, valid_instances in frame_instances.items():
+ per_frame_instance_mask[int(frame_idx), valid_instances] = 1
+
+ # select the frames that are in the range of start_timestep and end_timestep
+ instances_pose = torch.from_numpy(instances_pose[self.start_timestep:self.end_timestep]).float()
+ instances_size = torch.from_numpy(instances_size[self.start_timestep:self.end_timestep]).float()
+ instances_true_id = torch.from_numpy(instances_true_id).long()
+ instances_model_types = torch.from_numpy(instances_model_types).long()
+ per_frame_instance_mask = torch.from_numpy(per_frame_instance_mask[self.start_timestep:self.end_timestep]).bool()
+
+ # filter out the instances that are not visible in selected frames
+ ins_frame_cnt = per_frame_instance_mask.sum(dim=0)
+ instances_pose = instances_pose[:, ins_frame_cnt > 0]
+ instances_size = instances_size[:, ins_frame_cnt > 0]
+ instances_true_id = instances_true_id[ins_frame_cnt > 0]
+ instances_model_types = instances_model_types[ins_frame_cnt > 0]
+ per_frame_instance_mask = per_frame_instance_mask[:, ins_frame_cnt > 0]
+
+ # assign to the class
+ # (num_frames, num_instances, 4, 4)
+ self.instances_pose = instances_pose
+ # (num_instances, 3)
+ self.instances_size = instances_size.sum(0) / per_frame_instance_mask.sum(0).unsqueeze(-1)
+ # (num_frames, num_instances)
+ self.per_frame_instance_mask = per_frame_instance_mask
+ # (num_instances)
+ self.instances_true_id = instances_true_id
+ # (num_instances)
+ self.instances_model_types = instances_model_types
+
+ if self.data_cfg.load_smpl:
+ # Collect camera-to-world matrices for all available cameras
+ cam_to_worlds = {}
+ for cam_id in AVAILABLE_CAM_LIST:
+ cam_to_worlds[cam_id] = KITTICameraData.get_camera2worlds(
+ self.data_path,
+ str(cam_id),
+ self.start_timestep,
+ self.end_timestep
+ )
+
+ # load SMPL parameters
+ smpl_dict = joblib.load(os.path.join(self.data_path, "humanpose", "smpl.pkl"))
+ frame_num = self.end_timestep - self.start_timestep
+
+ smpl_human_all = {}
+ for fi in tqdm(range(self.start_timestep, self.end_timestep), desc="Loading SMPL"):
+ for instance_id, ins_smpl in smpl_dict.items():
+ if instance_id not in smpl_human_all:
+ smpl_human_all[instance_id] = {
+ "smpl_quats": torch.zeros((frame_num, 24, 4), dtype=torch.float32),
+ "smpl_trans": torch.zeros((frame_num, 3), dtype=torch.float32),
+ "smpl_betas": torch.zeros((frame_num, 10), dtype=torch.float32),
+ "frame_valid": torch.zeros((frame_num), dtype=torch.bool)
+ }
+ smpl_human_all[instance_id]["smpl_quats"][:, :, 0] = 1.0
+ if ins_smpl["valid_mask"][fi]:
+ betas = ins_smpl["smpl"]["betas"][fi]
+ smpl_human_all[instance_id]["smpl_betas"][fi - self.start_timestep] = betas
+
+ body_pose = ins_smpl["smpl"]["body_pose"][fi]
+ smpl_orient = ins_smpl["smpl"]["global_orient"][fi]
+ cam_depend = ins_smpl["selected_cam_idx"][fi].item()
+
+ c2w = cam_to_worlds[cam_depend][fi - self.start_timestep]
+ world_orient = c2w[:3, :3].to(smpl_orient.device) @ smpl_orient.squeeze()
+ smpl_quats = matrix_to_quaternion(
+ torch.cat([world_orient[None, ...], body_pose], dim=0)
+ )
+
+ ii = instances_info[str(instance_id)]['frame_annotations']["frame_idx"].index(fi)
+ o2w = np.array(
+ instances_info[str(instance_id)]['frame_annotations']["obj_to_world"][ii]
+ )
+ o2w = torch.from_numpy(
+ np.linalg.inv(ego_to_world_start) @ o2w
+ )
+ # box_size = instances_info[str(instance_id)]['frame_annotations']["box_size"][ii]
+
+ smpl_human_all[instance_id]["smpl_quats"][fi - self.start_timestep] = smpl_quats
+ smpl_human_all[instance_id]["smpl_trans"][fi - self.start_timestep] = o2w[:3, 3]
+ smpl_human_all[instance_id]["frame_valid"][fi - self.start_timestep] = True
+
+ self.smpl_human_all = smpl_human_all
+
+class KITTILiDARSource(SceneLidarSource):
+ def __init__(
+ self,
+ lidar_data_config: OmegaConf,
+ data_path: str,
+ start_timestep: int,
+ end_timestep: int,
+ device: torch.device = torch.device("cpu"),
+ ):
+ super().__init__(lidar_data_config, device=device)
+ self.data_path = data_path
+ self.start_timestep = start_timestep
+ self.end_timestep = end_timestep
+ self.create_all_filelist()
+ self.load_data()
+
+ def create_all_filelist(self):
+ """
+ Create a list of all the files in the dataset.
+ e.g., a list of all the lidar scans in the dataset.
+ """
+ lidar_filepaths = []
+ for t in range(self.start_timestep, self.end_timestep):
+ lidar_filepaths.append(
+ os.path.join(self.data_path, "lidar", f"{t:03d}.bin")
+ )
+ self.lidar_filepaths = np.array(lidar_filepaths)
+
+ def load_calibrations(self):
+ """
+ Load the calibration files of the dataset.
+ e.g., lidar to world transformation matrices.
+ """
+ # Note that in the Waymo Open Dataset, the lidar coordinate system is the same
+ # as the vehicle coordinate system
+ lidar_to_worlds = []
+
+ # we tranform the poses w.r.t. the first timestep to make the origin of the
+ # first ego pose as the origin of the world coordinate system.
+ ego_to_world_start = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{self.start_timestep:03d}.txt")
+ )
+ for t in range(self.start_timestep, self.end_timestep):
+ ego_to_world_current = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{t:03d}.txt")
+ )
+ # compute ego_to_world transformation
+ lidar_to_world = np.linalg.inv(ego_to_world_start) @ ego_to_world_current
+ lidar_to_worlds.append(lidar_to_world)
+
+ self.lidar_to_worlds = torch.from_numpy(
+ np.stack(lidar_to_worlds, axis=0)
+ ).float()
+
+ def load_lidar(self):
+ origins, directions, ranges = [], [], []
+ timesteps = []
+
+ accumulated_num_original_rays = 0
+ accumulated_num_rays = 0
+ for t in trange(
+ 0, len(self.lidar_filepaths), desc="Loading lidar", dynamic_ncols=True
+ ):
+ lidar_info = np.fromfile(self.lidar_filepaths[t], dtype=np.float32).reshape(-1, 4)
+ original_length = len(lidar_info)
+ accumulated_num_original_rays += original_length
+
+ lidar_points = torch.from_numpy(lidar_info[:, :3]).float()
+ # convert the lidar points to opencv coordinate system
+ # since lidar points already include the information of origins
+ # NOTE: here we use the pesudo origins, TODO: consider a more accurate way
+ lidar_origins = torch.zeros_like(lidar_points)
+
+ lidar_origins = (
+ self.lidar_to_worlds[t][:3, :3] @ lidar_origins.T
+ + self.lidar_to_worlds[t][:3, 3:4]
+ ).T
+ lidar_points = (
+ self.lidar_to_worlds[t][:3, :3] @ lidar_points.T
+ + self.lidar_to_worlds[t][:3, 3:4]
+ ).T
+ lidar_directions = lidar_points - lidar_origins
+ lidar_ranges = torch.norm(lidar_directions, dim=-1, keepdim=True)
+ lidar_directions = lidar_directions / lidar_ranges
+ lidar_timestamp = torch.ones_like(lidar_ranges).squeeze(-1) * t
+ accumulated_num_rays += len(lidar_ranges)
+
+ origins.append(lidar_origins)
+ directions.append(lidar_directions)
+ ranges.append(lidar_ranges)
+ timesteps.append(lidar_timestamp)
+
+ logger.info(
+ f"Number of lidar rays: {accumulated_num_rays} "
+ f"({accumulated_num_rays / accumulated_num_original_rays * 100:.2f}% of "
+ f"{accumulated_num_original_rays} original rays)"
+ )
+
+ self.origins = torch.cat(origins, dim=0)
+ self.directions = torch.cat(directions, dim=0)
+ self.ranges = torch.cat(ranges, dim=0)
+ self.visible_masks = torch.zeros_like(self.ranges).squeeze().bool()
+ self.colors = torch.ones_like(self.directions)
+
+ self._timesteps = torch.cat(timesteps, dim=0)
+ self.register_normalized_timestamps()
+
+ def get_lidar_rays(self, time_idx: int) -> Dict[str, Tensor]:
+ origins = self.origins[self.timesteps == time_idx]
+ directions = self.directions[self.timesteps == time_idx]
+ ranges = self.ranges[self.timesteps == time_idx]
+ normalized_time = self.normalized_time[self.timesteps == time_idx]
+ return {
+ "lidar_origins": origins,
+ "lidar_viewdirs": directions,
+ "lidar_ranges": ranges,
+ "lidar_normed_time": normalized_time,
+ "lidar_mask": self.timesteps == time_idx,
+ }
+
+ def delete_invisible_pts(self) -> None:
+ if self.visible_masks is not None:
+ num_bf = self.origins.shape[0]
+ self.origins = self.origins[self.visible_masks]
+ self.directions = self.directions[self.visible_masks]
+ self.ranges = self.ranges[self.visible_masks]
+ self._timesteps = self._timesteps[self.visible_masks]
+ self._normalized_time = self._normalized_time[self.visible_masks]
+ self.colors = self.colors[self.visible_masks]
+ logger.info(
+ f"[Lidar] {num_bf - self.visible_masks.sum()} out of {num_bf} points are cleared. {self.visible_masks.sum()} points left."
+ )
+ self.visible_masks = None
+ else:
+ logger.info("[Lidar] No unvisible points to clear.")
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/kitti/trackletparser.py b/worldbench/third_party/drivestudio/datasets/kitti/trackletparser.py
new file mode 100644
index 0000000..a59e6eb
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/kitti/trackletparser.py
@@ -0,0 +1,255 @@
+#!/usr/bin/env python
+"""
+parse XML files containing tracklet info for kitti data base (raw data section)
+(http://cvlibs.net/datasets/kitti/raw_data.php)
+
+No guarantees that this code is correct, usage is at your own risk!
+
+created by Christian Herdtweck, Max Planck Institute for Biological Cybernetics
+ (christian.herdtweck@tuebingen.mpg.de)
+
+requires numpy!
+"""
+
+from xml.etree.ElementTree import ElementTree
+import numpy as np
+import itertools
+from warnings import warn
+
+STATE_UNSET = 0
+STATE_INTERP = 1
+STATE_LABELED = 2
+stateFromText = {'0': STATE_UNSET, '1': STATE_INTERP, '2': STATE_LABELED}
+
+OCC_UNSET = 255 # -1 as uint8
+OCC_VISIBLE = 0
+OCC_PARTLY = 1
+OCC_FULLY = 2
+occFromText = {'-1': OCC_UNSET, '0': OCC_VISIBLE, '1': OCC_PARTLY, '2': OCC_FULLY}
+
+TRUNC_UNSET = 255 # -1 as uint8, but in xml files the value '99' is used!
+TRUNC_IN_IMAGE = 0
+TRUNC_TRUNCATED = 1
+TRUNC_OUT_IMAGE = 2
+TRUNC_BEHIND_IMAGE = 3
+truncFromText = {'99': TRUNC_UNSET, '0': TRUNC_IN_IMAGE, '1': TRUNC_TRUNCATED, \
+ '2': TRUNC_OUT_IMAGE, '3': TRUNC_BEHIND_IMAGE}
+xrange = range
+
+class Tracklet(object):
+ r""" representation an annotated object track
+
+ Tracklets are created in function parseXML and can most conveniently used as follows:
+
+ for trackletObj in parseXML(trackletFile):
+ for translation, rotation, state, occlusion, truncation, amtOcclusion, amtBorders, absoluteFrameNumber in trackletObj:
+ ... your code here ...
+ #end: for all frames
+ #end: for all tracklets
+
+ absoluteFrameNumber is in range [firstFrame, firstFrame+nFrames[
+ amtOcclusion and amtBorders could be None
+
+ You can of course also directly access the fields objType (string), size (len-3 ndarray), firstFrame/nFrames (int),
+ trans/rots (nFrames x 3 float ndarrays), states/truncs (len-nFrames uint8 ndarrays), occs (nFrames x 2 uint8 ndarray),
+ and for some tracklets amtOccs (nFrames x 2 float ndarray) and amtBorders (nFrames x 3 float ndarray). The last two
+ can be None if the xml file did not include these fields in poses
+ """
+
+ objectType = None
+ size = None # len-3 float array: (height, width, length)
+ firstFrame = None
+ trans = None # n x 3 float array (x,y,z)
+ rots = None # n x 3 float array (x,y,z)
+ states = None # len-n uint8 array of states
+ occs = None # n x 2 uint8 array (occlusion, occlusion_kf)
+ truncs = None # len-n uint8 array of truncation
+ amtOccs = None # None or (n x 2) float array (amt_occlusion, amt_occlusion_kf)
+ amtBorders = None # None (n x 3) float array (amt_border_l / _r / _kf)
+ nFrames = None
+
+ def __init__(self):
+ r""" create Tracklet with no info set """
+ self.size = np.nan * np.ones(3, dtype=float)
+
+ def __str__(self):
+ r""" return human-readable string representation of tracklet object
+
+ called implicitly in
+ print trackletObj
+ or in
+ text = str(trackletObj)
+ """
+ return '[Tracklet over {0} frames for {1}]'.format(self.nFrames, self.objectType)
+
+ def __iter__(self):
+ r""" returns an iterator that yields tuple of all the available data for each frame
+
+ called whenever code iterates over a tracklet object, e.g. in
+ for translation, rotation, state, occlusion, truncation, amtOcclusion, amtBorders, absoluteFrameNumber in trackletObj:
+ ...do something ...
+ or
+ trackDataIter = iter(trackletObj)
+ """
+ if self.amtOccs is None:
+ return itertools.izip(self.trans, self.rots, self.states, self.occs, self.truncs, \
+ itertools.repeat(None), itertools.repeat(None),
+ xrange(self.firstFrame, self.firstFrame + self.nFrames))
+ else:
+ return itertools.izip(self.trans, self.rots, self.states, self.occs, self.truncs, \
+ self.amtOccs, self.amtBorders,
+ xrange(self.firstFrame, self.firstFrame + self.nFrames))
+
+def parseXML(trackletFile):
+ r""" parse tracklet xml file and convert results to list of Tracklet objects
+
+ :param trackletFile: name of a tracklet xml file
+ :returns: list of Tracklet objects read from xml file
+ """
+
+ # convert tracklet XML data to a tree structure
+ eTree = ElementTree()
+ print
+ 'parsing tracklet file', trackletFile
+ with open(trackletFile) as f:
+ eTree.parse(f)
+
+ # now convert output to list of Tracklet objects
+ trackletsElem = eTree.find('tracklets')
+ tracklets = []
+ trackletIdx = 0
+ nTracklets = None
+ for trackletElem in trackletsElem:
+ # print 'track:', trackletElem.tag
+ if trackletElem.tag == 'count':
+ nTracklets = int(trackletElem.text)
+ print
+ 'file contains', nTracklets, 'tracklets'
+ elif trackletElem.tag == 'item_version':
+ pass
+ elif trackletElem.tag == 'item':
+ # print 'tracklet {0} of {1}'.format(trackletIdx, nTracklets)
+ # a tracklet
+ newTrack = Tracklet()
+ isFinished = False
+ hasAmt = False
+ frameIdx = None
+ for info in trackletElem:
+ # print 'trackInfo:', info.tag
+ if isFinished:
+ raise ValueError('more info on element after finished!')
+ if info.tag == 'objectType':
+ newTrack.objectType = info.text
+ elif info.tag == 'h':
+ newTrack.size[0] = float(info.text)
+ elif info.tag == 'w':
+ newTrack.size[1] = float(info.text)
+ elif info.tag == 'l':
+ newTrack.size[2] = float(info.text)
+ elif info.tag == 'first_frame':
+ newTrack.firstFrame = int(info.text)
+ elif info.tag == 'poses':
+ # this info is the possibly long list of poses
+ for pose in info:
+ # print 'trackInfoPose:', pose.tag
+ if pose.tag == 'count': # this should come before the others
+ if newTrack.nFrames is not None:
+ raise ValueError('there are several pose lists for a single track!')
+ elif frameIdx is not None:
+ raise ValueError('?!')
+ newTrack.nFrames = int(pose.text)
+ newTrack.trans = np.nan * np.ones((newTrack.nFrames, 3), dtype=float)
+ newTrack.rots = np.nan * np.ones((newTrack.nFrames, 3), dtype=float)
+ newTrack.states = np.nan * np.ones(newTrack.nFrames, dtype='uint8')
+ newTrack.occs = np.nan * np.ones((newTrack.nFrames, 2), dtype='uint8')
+ newTrack.truncs = np.nan * np.ones(newTrack.nFrames, dtype='uint8')
+ newTrack.amtOccs = np.nan * np.ones((newTrack.nFrames, 2), dtype=float)
+ newTrack.amtBorders = np.nan * np.ones((newTrack.nFrames, 3), dtype=float)
+ frameIdx = 0
+ elif pose.tag == 'item_version':
+ pass
+ elif pose.tag == 'item':
+ # pose in one frame
+ if frameIdx is None:
+ raise ValueError('pose item came before number of poses!')
+ for poseInfo in pose:
+ # print 'trackInfoPoseInfo:', poseInfo.tag
+ if poseInfo.tag == 'tx':
+ newTrack.trans[frameIdx, 0] = float(poseInfo.text)
+ elif poseInfo.tag == 'ty':
+ newTrack.trans[frameIdx, 1] = float(poseInfo.text)
+ elif poseInfo.tag == 'tz':
+ newTrack.trans[frameIdx, 2] = float(poseInfo.text)
+ elif poseInfo.tag == 'rx':
+ newTrack.rots[frameIdx, 0] = float(poseInfo.text)
+ elif poseInfo.tag == 'ry':
+ newTrack.rots[frameIdx, 1] = float(poseInfo.text)
+ elif poseInfo.tag == 'rz':
+ newTrack.rots[frameIdx, 2] = float(poseInfo.text)
+ elif poseInfo.tag == 'state':
+ newTrack.states[frameIdx] = stateFromText[poseInfo.text]
+ elif poseInfo.tag == 'occlusion':
+ newTrack.occs[frameIdx, 0] = occFromText[poseInfo.text]
+ elif poseInfo.tag == 'occlusion_kf':
+ newTrack.occs[frameIdx, 1] = occFromText[poseInfo.text]
+ elif poseInfo.tag == 'truncation':
+ newTrack.truncs[frameIdx] = truncFromText[poseInfo.text]
+ elif poseInfo.tag == 'amt_occlusion':
+ newTrack.amtOccs[frameIdx, 0] = float(poseInfo.text)
+ hasAmt = True
+ elif poseInfo.tag == 'amt_occlusion_kf':
+ newTrack.amtOccs[frameIdx, 1] = float(poseInfo.text)
+ hasAmt = True
+ elif poseInfo.tag == 'amt_border_l':
+ newTrack.amtBorders[frameIdx, 0] = float(poseInfo.text)
+ hasAmt = True
+ elif poseInfo.tag == 'amt_border_r':
+ newTrack.amtBorders[frameIdx, 1] = float(poseInfo.text)
+ hasAmt = True
+ elif poseInfo.tag == 'amt_border_kf':
+ newTrack.amtBorders[frameIdx, 2] = float(poseInfo.text)
+ hasAmt = True
+ else:
+ raise ValueError('unexpected tag in poses item: {0}!'.format(poseInfo.tag))
+ frameIdx += 1
+ else:
+ raise ValueError('unexpected pose info: {0}!'.format(pose.tag))
+ elif info.tag == 'finished':
+ isFinished = True
+ else:
+ raise ValueError('unexpected tag in tracklets: {0}!'.format(info.tag))
+ # end: for all fields in current tracklet
+
+ # some final consistency checks on new tracklet
+ if not isFinished:
+ warn('tracklet {0} was not finished!'.format(trackletIdx))
+ if newTrack.nFrames is None:
+ warn('tracklet {0} contains no information!'.format(trackletIdx))
+ elif frameIdx != newTrack.nFrames:
+ warn('tracklet {0} is supposed to have {1} frames, but perser found {1}!'.format( \
+ trackletIdx, newTrack.nFrames, frameIdx))
+ if np.abs(newTrack.rots[:, :2]).sum() > 1e-16:
+ warn('track contains rotation other than yaw!')
+
+ # if amtOccs / amtBorders are not set, set them to None
+ if not hasAmt:
+ newTrack.amtOccs = None
+ newTrack.amtBorders = None
+
+ # add new tracklet to list
+ tracklets.append(newTrack)
+ trackletIdx += 1
+
+ else:
+ raise ValueError('unexpected tracklet info')
+ # end: for tracklet list items
+
+ print
+ 'loaded', trackletIdx, 'tracklets'
+
+ # final consistency check
+ if trackletIdx != nTracklets:
+ warn('according to xml information the file has {0} tracklets, but parser found {1}!'.format(nTracklets,
+ trackletIdx))
+
+ return tracklets
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/nuplan/nuplan_human_utils.py b/worldbench/third_party/drivestudio/datasets/nuplan/nuplan_human_utils.py
new file mode 100644
index 0000000..e2419d0
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/nuplan/nuplan_human_utils.py
@@ -0,0 +1,218 @@
+from typing import List
+import os
+import cv2
+import json
+import logging
+import numpy as np
+from tqdm import tqdm
+
+from utils.geometry import (
+ get_corners,
+ project_camera_points_to_image
+)
+from .nuplan_sourceloader import (
+ SMPLNODE_CLASSES,
+ OPENCV2DATASET,
+ AVAILABLE_CAM_LIST,
+)
+
+logger = logging.getLogger()
+
+CAMERA_LIST = AVAILABLE_CAM_LIST
+
+def project_human_boxes(
+ scene_dir: str,
+ camera_list: List[int],
+ save_temp=True,
+ verbose=False,
+ narrow_width_ratio=0.2,
+ fps=12
+):
+ """Project human boxes to 2D image space and save the results to pkl files
+
+ Args:
+ scene_dir: str, path to the scene directory
+ camera_list: List[int], a list of camera ids to be processed
+ save_temp: bool, whether to save the intermediate results
+ verbose: bool, whether to visualize the projected boxes
+ narrow_width_ratio: sometimes the projected boxes are too wide
+ we can narrow them with this ratio to get more accurate results
+ fps: int, FPS for the visualization video
+
+ Returns:
+ collector_all: dict, a dictionary containing the projected boxes for each camera
+ """
+ # check if the necessary directories exist
+ images_dir = f'{scene_dir}/images'
+ poses_dir = f'{scene_dir}/ego_pose'
+ extrinsics_dir = f'{scene_dir}/extrinsics'
+ intrinsics_dir = f'{scene_dir}/intrinsics'
+ instances_dir = f'{scene_dir}/instances'
+ valid_paths = [images_dir, poses_dir, extrinsics_dir, intrinsics_dir, instances_dir]
+ for path in valid_paths:
+ assert os.path.exists(path), \
+ f"Path {path} does not exist, you need to run waymo preprocess to generate the necessary files"
+
+ # create directories for saving the results
+ save_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks'
+ if not os.path.exists(save_dir):
+ os.makedirs(save_dir)
+
+ if verbose:
+ # create directories for saving the intermediate visualization results
+ video_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks/vis'
+ if not os.path.exists(video_dir):
+ os.makedirs(video_dir)
+ per_human_img_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks/vis/images'
+ if not os.path.exists(per_human_img_dir):
+ os.makedirs(per_human_img_dir)
+
+ # load instances and frame infos
+ frame_infos = json.load(
+ open(f'{instances_dir}/frame_instances.json')
+ )
+ instances_meta = json.load(
+ open(f'{instances_dir}/instances_info.json')
+ )
+
+ collector_all = {}
+ # iterate over each camera
+ for cam_id in camera_list:
+ # check if already processed
+ pkl_path = os.path.join(save_dir, f"{cam_id}.pkl")
+ if os.path.exists(pkl_path):
+ collector_all[cam_id] = json.load(open(pkl_path))
+ logger.info(f"Results for camera {cam_id} already exists at {pkl_path}")
+ continue
+
+ if verbose:
+ per_cam_vis_dir = os.path.join(per_human_img_dir, f"{cam_id}")
+ if not os.path.exists(per_cam_vis_dir):
+ os.makedirs(per_cam_vis_dir)
+
+ collector = {}
+ frames = []
+ for frame_id, frame_ins_list in frame_infos.items():
+ frame_id = int(frame_id)
+
+ # define empty instance collector for each frame
+ frame_collector = {
+ "gt_bbox": [],
+ "extra_data": {
+ "gt_track_id": [],
+ "gt_class": [],
+ }
+ }
+
+ # load extrinsic
+ cam_to_ego = np.loadtxt(os.path.join(extrinsics_dir, f"{cam_id}.txt"))
+ cam_to_ego = cam_to_ego @ OPENCV2DATASET
+ ego_to_world = np.loadtxt(os.path.join(poses_dir, f"{str(frame_id).zfill(3)}.txt"))
+ cam2world = ego_to_world @ cam_to_ego
+
+ # load intrinsic
+ Ks = np.loadtxt(os.path.join(intrinsics_dir, f"{cam_id}.txt"))
+ fx, fy, cx, cy = Ks[0], Ks[1], Ks[2], Ks[3]
+ # k1, k2, p1, p2, k3 = Ks[4], Ks[5], Ks[6], Ks[7], Ks[8]
+ intrinsic = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float32)
+
+ # load image
+ ori_image = cv2.imread(
+ os.path.join(images_dir, f"{str(frame_id).zfill(3)}_{cam_id}.jpg")
+ )
+ image = ori_image.copy()
+ H, W = image.shape[:2]
+
+ # iterate over each instance in the frame
+ if len(frame_ins_list) > 0:
+ # if there are pedestrians in this frame, project the boxes to the image
+ image_plotted = image.copy()
+ for instance_id in frame_ins_list:
+ ins = instances_meta[str(instance_id)]
+
+ if ins["class_name"] not in SMPLNODE_CLASSES:
+ continue
+
+ ins_anno = ins["frame_annotations"]
+ index = ins_anno['frame_idx'].index(frame_id)
+ obj_to_world = np.array(ins_anno['obj_to_world'][index])
+ l, w, h = ins_anno['box_size'][index]
+
+ # get box corners in object space
+ corners = get_corners(l, w, h)
+ # transform box corners to world space
+ corners_world = obj_to_world[:3, :3] @ corners + obj_to_world[:3, 3:4]
+ # transform box corners to image space
+ world2cam = np.linalg.inv(cam2world)
+ corners_cam = world2cam[:3, :3] @ corners_world + world2cam[:3, 3:4]
+ cam_points, depth = project_camera_points_to_image(corners_cam.T, intrinsic)
+
+ x_min, y_min = np.min(cam_points, axis=0)
+ x_max, y_max = np.max(cam_points, axis=0)
+ # clip left and right with this ratio
+ if narrow_width_ratio > 0.:
+ length = x_max - x_min
+ x_min += length * narrow_width_ratio
+ x_max -= length * narrow_width_ratio
+
+ # clip the box to the image
+ original_area = (x_max - x_min) * (y_max - y_min)
+ x_min, x_max = np.clip(x_min, 0, W), np.clip(x_max, 0, W)
+ y_min, y_max = np.clip(y_min, 0, H), np.clip(y_max, 0, H)
+ new_area = (x_max - x_min) * (y_max - y_min)
+
+ # filter out boxes that are too small or too large
+ behind = depth.max() < 0
+ too_small = new_area < W * H * (0.03)**2
+ too_large = new_area > W * H / 1.1
+ too_far = np.linalg.norm(obj_to_world[:3, 3] - cam2world[:3, 3]) > 40
+ clip_large = new_area / original_area < 1/3
+ if too_small or too_large or clip_large or behind or too_far:
+ continue
+
+ # gt box on image
+ gt_box = [x_min, y_min, x_max - x_min, y_max - y_min]
+
+ # save the projected box to the collector
+ frame_collector["gt_bbox"].append(gt_box)
+ frame_collector["extra_data"]["gt_track_id"].append(instance_id)
+ frame_collector["extra_data"]["gt_class"].append([0])
+
+ if verbose:
+ # visualize the projected boxes of ONE instance
+ raw_image = cv2.rectangle(
+ ori_image.copy(), (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 2
+ )
+ raw_image_path = os.path.join(per_cam_vis_dir, f"{frame_id}_{instance_id}.jpg")
+ cv2.imwrite(raw_image_path, raw_image)
+
+ # add this instance to the image
+ image_plotted = cv2.rectangle(
+ image, (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 2
+ )
+
+ if verbose:
+ frames.append(image_plotted)
+ else:
+ # if no instance in this frame, just save the original image
+ if verbose:
+ frames.append(ori_image)
+
+ collector[frame_id] = frame_collector
+
+ if verbose:
+ height, width = frames[0].shape[:2]
+ output_path = os.path.join(video_dir, f"cam_{cam_id}.mp4")
+ fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+ out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+ for frame in tqdm(frames, desc=f"Writing video for camera {cam_id}"):
+ out.write(frame)
+ out.release()
+
+ if save_temp:
+ # save collector to pkl
+ json.dump(collector, open(pkl_path, "w"))
+
+ collector_all[cam_id] = collector
+
+ return collector_all
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/nuplan/nuplan_preprocess.py b/worldbench/third_party/drivestudio/datasets/nuplan/nuplan_preprocess.py
new file mode 100644
index 0000000..efbc019
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/nuplan/nuplan_preprocess.py
@@ -0,0 +1,527 @@
+import json
+import os
+from typing import Dict, List
+
+import numpy as np
+from PIL import Image
+from pyquaternion import Quaternion
+
+from nuplan.database.nuplan_db.nuplan_scenario_queries import get_images_from_lidar_tokens
+from nuplan.database.nuplan_db_orm.nuplandb import NuPlanDB
+from nuplan.database.nuplan_db_orm.nuplandb_wrapper import NuPlanDBWrapper
+from nuplan.database.utils.pointclouds.lidar import LidarPointCloud
+
+from datasets.tools.multiprocess_utils import track_parallel_progress
+from utils.geometry import get_corners, project_camera_points_to_image
+from utils.visualization import color_mapper, dump_3d_bbox_on_image
+from .nuplan_utils import get_egopose3d_for_lidarpc_token_from_db, get_tracked_objects_for_lidarpc_token_from_db
+
+NUPLAN_LABELS = [
+ 'vehicle', 'pedestrian', 'bicycle'
+]
+NUPLAN_NONRIGID_DYNAMIC_CLASSES = [
+ 'pedestrian', 'bicycle'
+]
+NUPLAN_RIGID_DYNAMIC_CLASSES = [
+ 'vehicle'
+]
+NUPLAN_DYNAMIC_CLASSES = NUPLAN_NONRIGID_DYNAMIC_CLASSES + NUPLAN_RIGID_DYNAMIC_CLASSES
+
+class NuPlanProcessor(object):
+ """Process NUPLAN Dataset
+
+ NuPlan Datasets provides 8 cameras and Merged Lidar data.
+ Cameras works in 10Hz and Lidar works in 20Hz. Thus we process at 10Hz.
+ The duration of each scene is around 8 mins, resulting in ~5000 frames.
+ We only process the first max_frame_limit(default=300) frames.
+
+ Args:
+ load_dir (str): Directory to load data.
+ save_dir (str): Directory to save data in processed format.
+ prefix (str): Prefix of filename.
+ workers (int, optional): Number of workers for the parallel process.
+ process_keys (list, optional): List of keys to process. Default: ["images", "lidar", "calib", "dynamic_masks", "objects"]
+ process_log_list (list, optional): List of scene indices to process. Default: None
+ """
+
+ def __init__(
+ self,
+ load_dir='data/nuplan/raw',
+ save_dir='data/nuplan/processed',
+ prefix='mini',
+ # We skip the first start_frame_idx frames to avoid ego static frames
+ start_frame_idx=1000,
+ # We only process the max_frame_limit frames
+ max_frame_limit=300,
+ process_keys=[
+ "images",
+ "pose",
+ "calib",
+ "lidar",
+ "dynamic_masks",
+ "objects"
+ ],
+ process_id_list=None,
+ workers=64,
+ ):
+ self.HW = (1080, 1920)
+ print("Raw Image Resolution: ", self.HW)
+ self.process_keys = process_keys
+ print("will process keys: ", self.process_keys)
+ self.start_frame_idx = start_frame_idx
+ print("We will skip the first {} frames".format(self.start_frame_idx))
+ self.max_frame_limit = max_frame_limit
+ print("We will process the first {} frames each scene".format(self.max_frame_limit))
+ # the lidar data is collected at 20Hz, we need to downsample to 10Hz to match the camera data
+ self.lidar_idxs = range(self.start_frame_idx, self.start_frame_idx + self.max_frame_limit * 2, 2)
+
+ # NUPLAN Provides 8 cameras
+ self.cam_list = [ # {frame_idx}_{cam_id}.jpg
+ "CAM_F0", # "xxx_0.jpg"
+ "CAM_L0", # "xxx_1.jpg"
+ "CAM_R0", # "xxx_2.jpg"
+ "CAM_L1", # "xxx_3.jpg"
+ "CAM_R1", # "xxx_4.jpg"
+ "CAM_L2", # "xxx_5.jpg"
+ "CAM_R2", # "xxx_6.jpg"
+ "CAM_B0" # "xxx_7.jpg"
+ ]
+
+ self.sensor_blobs_dir = os.path.join(load_dir, 'nuplan-v1.1', 'sensor_blobs')
+ self.split_dir = os.path.join(load_dir, 'nuplan-v1.1', 'splits', prefix)
+ self.nuplandb_wrapper = NuPlanDBWrapper(
+ data_root=os.path.join(load_dir, 'nuplan-v1.1'),
+ map_root=os.path.join(load_dir, 'maps'),
+ db_files=self.split_dir,
+ map_version='nuplan-maps-v1.0',
+ )
+
+ process_log_list = []
+ for idx in process_id_list:
+ process_log_list.append(self.nuplandb_wrapper.log_names[idx])
+ self.process_log_list = process_log_list
+
+ self.save_dir = os.path.join(save_dir, prefix)
+ self.workers = int(workers)
+ self.create_folder()
+
+ def convert(self):
+ """Convert action."""
+ print("Start converting ...")
+ if self.process_log_list is None:
+ id_list = range(len(self))
+ else:
+ id_list = self.process_log_list
+ track_parallel_progress(self.convert_one, id_list, self.workers)
+ print("\nFinished ...")
+
+ def convert_one(self, scene_log_name):
+ """Convert action for single file."""
+ # get log db
+ log_db = self.nuplandb_wrapper.get_log_db(scene_log_name)
+
+ # since lidar and images are captured at different frequency
+ # we find the best start frame that lidar and images matches the best
+ # lidar_idx:[0] 1 [2] 3 [4] 5 [6]
+ # timestamp: 0 0.05 0.1 0.15 0.2 0.25 0.3
+ # lidar_pc: | | | | | | |
+ # Images: | | | |
+ # NOTE: the best match should be the frame with the closest timestamp to the lidar_pc (e.g. [0] [2] [4] [6])
+ # calulate time shift of original start frame
+ lidar_pc = log_db.lidar_pc[self.start_frame_idx]
+ images = get_images_from_lidar_tokens(
+ log_file=os.path.join(self.split_dir, log_db.log_name + '.db'),
+ tokens=[lidar_pc.token],
+ channels=self.cam_list,
+ )
+ images_timestamps = [image.timestamp for image in images]
+ lidar_timestamp = lidar_pc.timestamp
+ no_shift_time_diff = [abs(lidar_timestamp - timestamp) for timestamp in images_timestamps]
+ # calulate time shift of original start frame + 1
+ lidar_pc = log_db.lidar_pc[self.start_frame_idx + 1]
+ images = get_images_from_lidar_tokens(
+ log_file=os.path.join(self.split_dir, log_db.log_name + '.db'),
+ tokens=[lidar_pc.token],
+ channels=self.cam_list,
+ )
+ images_timestamps = [image.timestamp for image in images]
+ lidar_timestamp = lidar_pc.timestamp
+ shift_time_diff = [abs(lidar_timestamp - timestamp) for timestamp in images_timestamps]
+
+ if sum(no_shift_time_diff) > sum(shift_time_diff):
+ lidar_idxs = [idx + 1 for idx in self.lidar_idxs]
+ else:
+ lidar_idxs = self.lidar_idxs
+
+ if "images" in self.process_keys:
+ self.save_image(log_db, lidar_idxs)
+ print(f"Processed images for {scene_log_name}")
+ if "calib" in self.process_keys:
+ self.save_calib(log_db, lidar_idxs)
+ print(f"Processed calib for {scene_log_name}")
+ if "lidar" in self.process_keys:
+ self.save_lidar(log_db, lidar_idxs)
+ print(f"Processed lidar for {scene_log_name}")
+ if "pose" in self.process_keys:
+ self.save_pose(log_db, lidar_idxs)
+ print(f"Processed pose for {scene_log_name}")
+ if "dynamic_masks" in self.process_keys:
+ self.save_dynamic_mask(log_db, lidar_idxs, valid_classes=NUPLAN_DYNAMIC_CLASSES, dir_name='all')
+ self.save_dynamic_mask(log_db, lidar_idxs, valid_classes=NUPLAN_NONRIGID_DYNAMIC_CLASSES, dir_name='human')
+ self.save_dynamic_mask(log_db, lidar_idxs, valid_classes=NUPLAN_RIGID_DYNAMIC_CLASSES, dir_name='vehicle')
+ print(f"Processed dynamic masks for {scene_log_name}")
+
+ # process annotated objects
+ if "objects" in self.process_keys:
+ instances_info, frame_instances = self.save_objects(log_db, lidar_idxs)
+ print(f"Processed instances info for {scene_log_name}")
+
+ # Save instances info and frame instances
+ object_info_dir = f"{self.save_dir}/{scene_log_name}/instances"
+ with open(f"{object_info_dir}/instances_info.json", "w") as fp:
+ json.dump(instances_info, fp, indent=4)
+ with open(f"{object_info_dir}/frame_instances.json", "w") as fp:
+ json.dump(frame_instances, fp, indent=4)
+
+ if "objects_vis" in self.process_keys:
+ self.visualize_dynamic_objects(
+ log_db, lidar_idxs,
+ output_path=f"{object_info_dir}/debug_vis",
+ instances_info=instances_info,
+ frame_instances=frame_instances
+ )
+ print(f"Processed objects visualization for {scene_log_name}")
+
+ def __len__(self):
+ """Length of the filename list."""
+ return len(self.process_log_list)
+
+ def save_image(self, log_db: NuPlanDB, lidar_idxs: List[int]):
+ """Parse and save the images in jpg format."""
+ lidar_pcs = log_db.lidar_pc
+ for frame_idx, lidar_idx in enumerate(lidar_idxs):
+ if frame_idx >= self.max_frame_limit:
+ break
+ lidar_pc = lidar_pcs[lidar_idx]
+
+ images = get_images_from_lidar_tokens(
+ log_file=os.path.join(self.split_dir, log_db.log_name + '.db'),
+ tokens=[lidar_pc.token],
+ channels=self.cam_list,
+ )
+
+ image_cnt = 0
+ for cam_id, image in enumerate(images):
+ raw_image_path = os.path.join(self.sensor_blobs_dir, image.filename_jpg)
+ image_save_path = f"{self.save_dir}/{log_db.log_name}/images/{str(frame_idx).zfill(3)}_{cam_id}.jpg"
+
+ os.system(f"cp {raw_image_path} {image_save_path}")
+ image_cnt+=1
+
+ assert image_cnt == len(self.cam_list), \
+ f"Image number, camera number mismatch: {image_cnt} != {len(self.cam_list)}"
+
+ def get_cameras_calib(self, log_db: NuPlanDB):
+ """Get the camera calibration."""
+ cameras = log_db.camera
+ extrinsics, intrinsics, distortions = {}, {}, {}
+ for cam in cameras:
+ channel = cam.channel
+
+ extrinsic = Quaternion(cam.rotation).transformation_matrix
+ extrinsic[:3, 3] = np.array(cam.translation)
+ extrinsics[channel] = extrinsic
+
+ intrinsic = np.array(cam.intrinsic)
+ intrinsics[channel] = intrinsic
+
+ distortions[channel] = np.array(cam.distortion)
+
+ return extrinsics, intrinsics, distortions
+
+ def save_calib(self, log_db: NuPlanDB, lidar_idxs: List[int]):
+ """Parse and save the calibration data."""
+ extrinsics, intrinsics, distortions = self.get_cameras_calib(log_db)
+ for channel in self.cam_list:
+ cam_id = self.cam_list.index(channel)
+
+ intrinsic = intrinsics[channel]
+ fx, fy, cx, cy = intrinsic[0, 0], intrinsic[1, 1], intrinsic[0, 2], intrinsic[1, 2]
+ k1, k2, p1, p2, k3 = distortions[channel]
+ Ks = np.array([fx, fy, cx, cy, k1, k2, p1, p2, k3])
+
+ np.savetxt(
+ f"{self.save_dir}/{log_db.log_name}/extrinsics/"
+ + f"{str(cam_id)}.txt",
+ extrinsics[channel]
+ )
+ np.savetxt(
+ f"{self.save_dir}/{log_db.log_name}/intrinsics/"
+ + f"{str(cam_id)}.txt",
+ Ks
+ )
+
+ def save_lidar(self, log_db: NuPlanDB, lidar_idxs: List[int]):
+ """Parse and save the lidar data in psd format."""
+ # NOTE: lidar points is already in ego pose frame
+ lidar_pcs = log_db.lidar_pc
+ for frame_idx, lidar_idx in enumerate(lidar_idxs):
+ if frame_idx >= self.max_frame_limit:
+ break
+ lidar_pc = lidar_pcs[lidar_idx]
+
+ lidar_data: LidarPointCloud = lidar_pc.load(log_db, lidar_idxs)
+ # version 1: a numpy array with 5 cols (x, y, z, intensity, ring).
+ # version 2: a numpy array with 6 cols (x, y, z, intensity, ring, lidar_id).
+ lidar_save_path = f"{self.save_dir}/{log_db.log_name}/lidar/{str(frame_idx).zfill(3)}.bin"
+ lidar_data.points.T.astype(np.float32).tofile(lidar_save_path)
+
+ def save_pose(self, log_db: NuPlanDB, lidar_idxs: List[int]):
+ """Parse and save the pose data."""
+ lidar_pcs = log_db.lidar_pc
+ for frame_idx, lidar_idx in enumerate(lidar_idxs):
+ if frame_idx >= self.max_frame_limit:
+ break
+ lidar_pc = lidar_pcs[lidar_idx]
+
+ ego_pose = get_egopose3d_for_lidarpc_token_from_db(
+ log_file=os.path.join(self.split_dir, log_db.log_name + '.db'),
+ token=lidar_pc.token
+ )
+
+ np.savetxt(
+ f"{self.save_dir}/{log_db.log_name}/ego_pose/"
+ + f"{str(frame_idx).zfill(3)}.txt",
+ ego_pose
+ )
+
+ def save_dynamic_mask(
+ self, log_db: NuPlanDB, lidar_idxs: List[int],
+ valid_classes: List[str], dir_name: str
+ ):
+ """Parse and save the segmentation data."""
+ extrinsics, intrinsics, _ = self.get_cameras_calib(log_db)
+ lidar_pcs = log_db.lidar_pc
+ for frame_idx, lidar_idx in enumerate(lidar_idxs):
+ if frame_idx >= self.max_frame_limit:
+ break
+ lidar_pc = lidar_pcs[lidar_idx]
+
+ ego_pose = get_egopose3d_for_lidarpc_token_from_db(
+ log_file=os.path.join(self.split_dir, log_db.log_name + '.db'),
+ token=lidar_pc.token
+ )
+
+ objects_generator = get_tracked_objects_for_lidarpc_token_from_db(
+ log_file=os.path.join(self.split_dir, log_db.log_name + '.db'),
+ token=lidar_pc.token
+ )
+ objects = [obj for obj in objects_generator if obj.category in valid_classes]
+
+ for channel in self.cam_list:
+ cam_id = self.cam_list.index(channel)
+ dynamic_mask = np.zeros(self.HW, dtype=np.float32)
+
+ # compute dynamic mask according to the instances' bbox projection
+ for obj in objects:
+ obj_to_world = obj.pose
+ l, w, h = obj.box_size
+ corners = get_corners(l, w, h)
+ corners_world = obj_to_world[:3, :3] @ corners + obj_to_world[:3, 3:4]
+
+ world_to_ego = np.linalg.inv(ego_pose)
+ corners_ego = world_to_ego[:3, :3] @ corners_world + world_to_ego[:3, 3:4]
+
+ ego_to_cam = np.linalg.inv(extrinsics[channel])
+ corners_cam = ego_to_cam[:3, :3] @ corners_ego + ego_to_cam[:3, 3:4]
+
+ corners_2d, _ = project_camera_points_to_image(corners_cam.T, intrinsics[channel])
+
+ # Check if the object is in front of the camera and all corners are in the image
+ # NOTE: we use strict visibility check here, requiring all corners to be visible
+ in_front = np.all(corners_cam[2, :] > 0.1)
+ in_image = np.all(corners_2d[0, :] >= 0) & np.all(corners_2d[0, :] < self.HW[1]) & \
+ np.all(corners_2d[1, :] >= 0) & np.all(corners_2d[1, :] < self.HW[0])
+ if not (in_front and in_image):
+ continue
+
+ # Fill the mask
+ u, v = corners_2d[0, :].astype(np.int32), corners_2d[1, :].astype(np.int32)
+ u = np.clip(u, 0, self.HW[1] - 1)
+ v = np.clip(v, 0, self.HW[0] - 1)
+
+ if u.max() - u.min() == 0 or v.max() - v.min() == 0:
+ continue
+
+ xy = (u.min(), v.min())
+ width = u.max() - u.min()
+ height = v.max() - v.min()
+
+ dynamic_mask[
+ int(xy[1]): int(xy[1] + height),
+ int(xy[0]): int(xy[0] + width)
+ ] = np.maximum(
+ dynamic_mask[
+ int(xy[1]): int(xy[1] + height),
+ int(xy[0]): int(xy[0] + width)
+ ],
+ 1
+ )
+
+ dynamic_mask = np.clip((dynamic_mask > 0.) * 255, 0, 255).astype(np.uint8)
+ dynamic_mask = Image.fromarray(dynamic_mask, "L")
+ dynamic_mask_path = os.path.join(
+ self.save_dir, log_db.log_name, "dynamic_masks", dir_name,
+ f"{str(frame_idx).zfill(3)}_{cam_id}.png"
+ )
+ dynamic_mask.save(dynamic_mask_path)
+
+ def save_objects(self, log_db: NuPlanDB, lidar_idxs: List[int]):
+ """Parse and save the objects annotation data."""
+ instances_info, frame_instances = {}, {}
+
+ lidar_pcs = log_db.lidar_pc
+ for frame_idx, lidar_idx in enumerate(lidar_idxs):
+ if frame_idx >= self.max_frame_limit:
+ break
+ lidar_pc = lidar_pcs[lidar_idx]
+
+ objects_generator = get_tracked_objects_for_lidarpc_token_from_db(
+ log_file=os.path.join(self.split_dir, log_db.log_name + '.db'),
+ token=lidar_pc.token
+ )
+ objects = [obj for obj in objects_generator if obj.category in NUPLAN_DYNAMIC_CLASSES]
+
+ frame_instances[frame_idx] = []
+ for obj in objects:
+ obj_id = obj.track_token
+ if obj_id not in instances_info:
+ instances_info[obj_id] = {
+ "id": obj_id,
+ "class_name": obj.category,
+ "frame_annotations": {
+ "frame_idx": [],
+ "obj_to_world": [],
+ "box_size": [],
+ }
+ }
+
+ obj_to_world = obj.pose
+ l, w, h = obj.box_size
+
+ instances_info[obj_id]['frame_annotations']['frame_idx'].append(frame_idx)
+ instances_info[obj_id]['frame_annotations']['obj_to_world'].append(obj_to_world.tolist())
+ instances_info[obj_id]['frame_annotations']['box_size'].append([l, w, h])
+
+ frame_instances[frame_idx].append(obj_id)
+
+ # Correct ID mapping
+ id_map = {}
+ for i, (k, v) in enumerate(instances_info.items()):
+ id_map[v["id"]] = i
+
+ # Update keys in instances_info
+ new_instances_info = {}
+ for k, v in instances_info.items():
+ new_instances_info[id_map[v["id"]]] = v
+
+ # Update keys in frame_instances
+ new_frame_instances = {}
+ for k, v in frame_instances.items():
+ new_frame_instances[k] = [id_map[i] for i in v]
+
+ return new_instances_info, new_frame_instances
+
+ def visualize_dynamic_objects(
+ self, log_db: NuPlanDB,
+ lidar_idxs: List[int], output_path: str,
+ instances_info: Dict, frame_instances: Dict
+ ):
+ """Visualize the dynamic objects' box with different colors on the image."""
+ extrinsics, intrinsics, _ = self.get_cameras_calib(log_db)
+ lidar_pcs = log_db.lidar_pc
+
+ for frame_idx, lidar_idx in enumerate(lidar_idxs):
+ if frame_idx >= self.max_frame_limit:
+ break
+ lidar_pc = lidar_pcs[lidar_idx]
+
+ ego_pose = get_egopose3d_for_lidarpc_token_from_db(
+ log_file=os.path.join(self.split_dir, log_db.log_name + '.db'),
+ token=lidar_pc.token
+ )
+
+ images = get_images_from_lidar_tokens(
+ log_file=os.path.join(self.split_dir, log_db.log_name + '.db'),
+ tokens=[lidar_pc.token],
+ channels=self.cam_list,
+ )
+
+ for cam_id, image in enumerate(images):
+ raw_image_path = os.path.join(self.sensor_blobs_dir, image.filename_jpg)
+ canvas = np.array(Image.open(raw_image_path))
+
+ channel = self.cam_list[cam_id]
+ objects = frame_instances[frame_idx]
+ if len(objects) == 0:
+ img_plotted = canvas
+ else:
+ lstProj2d = []
+ color_list = []
+
+ for obj_id in objects:
+ idx_in_obj = instances_info[obj_id]['frame_annotations']['frame_idx'].index(frame_idx)
+ obj_to_world = np.array(instances_info[obj_id]['frame_annotations']['obj_to_world'][idx_in_obj])
+ l, w, h = instances_info[obj_id]['frame_annotations']['box_size'][idx_in_obj]
+
+ corners = get_corners(l, w, h)
+ corners_world = obj_to_world[:3, :3] @ corners + obj_to_world[:3, 3:4]
+
+ world_to_ego = np.linalg.inv(ego_pose)
+ corners_ego = world_to_ego[:3, :3] @ corners_world + world_to_ego[:3, 3:4]
+
+ ego_to_cam = np.linalg.inv(extrinsics[channel])
+ corners_cam = ego_to_cam[:3, :3] @ corners_ego + ego_to_cam[:3, 3:4]
+
+ corners_2d, _ = project_camera_points_to_image(corners_cam.T, intrinsics[channel])
+
+ # Check if the object is in front of the camera and at least one corner is in the image
+ in_front = np.all(corners_cam[2, :] > 0.1)
+ intersected = np.any(
+ (corners_2d[0, :] >= 0) & (corners_2d[0, :] < self.HW[1]) & \
+ (corners_2d[1, :] >= 0) & (corners_2d[1, :] < self.HW[0])
+ )
+ if in_front and intersected:
+ lstProj2d.append(corners_2d.T.tolist())
+ color_list.append(color_mapper(str(obj_id)))
+
+ img_plotted = dump_3d_bbox_on_image(coords=np.array(lstProj2d), img=canvas, color=color_list)
+
+ img_path = os.path.join(output_path, f"{str(frame_idx).zfill(3)}_{cam_id}.jpg")
+ Image.fromarray(img_plotted).save(img_path)
+
+ def create_folder(self):
+ """Create folder for data preprocessing."""
+ if self.process_log_list is None:
+ id_list = range(len(self))
+ else:
+ id_list = self.process_log_list
+ for scene_log_name in id_list:
+ if "images" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{scene_log_name}/images", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{scene_log_name}/sky_masks", exist_ok=True)
+ if "pose" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{scene_log_name}/ego_pose", exist_ok=True)
+ if "calib" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{scene_log_name}/extrinsics", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{scene_log_name}/intrinsics", exist_ok=True)
+ if "lidar" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{scene_log_name}/lidar", exist_ok=True)
+ if "dynamic_masks" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{scene_log_name}/dynamic_masks/all", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{scene_log_name}/dynamic_masks/human", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{scene_log_name}/dynamic_masks/vehicle", exist_ok=True)
+ if "objects" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{scene_log_name}/instances", exist_ok=True)
+ if "objects_vis" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{scene_log_name}/instances/debug_vis", exist_ok=True)
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/nuplan/nuplan_preview.py b/worldbench/third_party/drivestudio/datasets/nuplan/nuplan_preview.py
new file mode 100644
index 0000000..2640ac0
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/nuplan/nuplan_preview.py
@@ -0,0 +1,202 @@
+import os
+from typing import List
+
+import cv2
+import numpy as np
+import numpy.typing as npt
+from moviepy.editor import VideoFileClip, clips_array
+from tqdm import tqdm
+
+NDArrayByte = npt.NDArray[np.uint8]
+
+def layout_nuplan(
+ imgs: List[np.ndarray], cam_names: List[str]
+) -> np.ndarray:
+ """Combine cameras into a tiled image for NuPlan dataset."""
+ channel = imgs[0].shape[-1]
+ height, width = imgs[0].shape[:2]
+
+ tiled_height = height * 3
+ tiled_width = width * 3
+ tiled_img = np.zeros((tiled_height, tiled_width, channel), dtype=np.uint8)
+ filled_mask = np.zeros((tiled_height, tiled_width), dtype=np.uint8)
+
+ for idx, cam_name in enumerate(cam_names):
+ img = imgs[idx]
+ if cam_name == "CAM_F0":
+ tiled_img[:height, width:2*width] = img
+ filled_mask[:height, width:2*width] = 1
+ elif cam_name == "CAM_L0":
+ tiled_img[:height, :width] = img
+ filled_mask[:height, :width] = 1
+ elif cam_name == "CAM_R0":
+ tiled_img[:height, 2*width:] = img
+ filled_mask[:height, 2*width:] = 1
+ elif cam_name == "CAM_L1":
+ tiled_img[height:2*height, :width] = img
+ filled_mask[height:2*height, :width] = 1
+ elif cam_name == "CAM_R1":
+ tiled_img[height:2*height, 2*width:] = img
+ filled_mask[height:2*height, 2*width:] = 1
+ elif cam_name == "CAM_L2":
+ tiled_img[2*height:, :width] = img
+ filled_mask[2*height:, :width] = 1
+ elif cam_name == "CAM_R2":
+ tiled_img[2*height:, 2*width:] = img
+ filled_mask[2*height:, 2*width:] = 1
+ elif cam_name == "CAM_B0":
+ tiled_img[2*height:, width:2*width] = img
+ filled_mask[2*height:, width:2*width] = 1
+
+ # Crop the image according to the largest filled area
+ min_y, max_y = np.where(filled_mask)[0].min(), np.where(filled_mask)[0].max()
+ min_x, max_x = np.where(filled_mask)[1].min(), np.where(filled_mask)[1].max()
+ tiled_img = tiled_img[min_y:max_y+1, min_x:max_x+1]
+
+ return tiled_img
+
+def preview_instances(
+ input_root: str = 'data/nuplan/processed/mini',
+ output_root: str = 'data/nuplan/preview/instances',
+ frame_width: int = 1920,
+ frame_height: int = 1080,
+ fps: int = 24,
+):
+ if not os.path.exists(output_root):
+ os.makedirs(output_root)
+
+ # Iterate over each scene directory in the input root directory
+ for scene_id in tqdm(sorted(os.listdir(input_root))):
+ scene_folder = os.path.join(input_root, scene_id, "instances", "debug_vis")
+
+ # Define the output video file path
+ output_video_path = os.path.join(output_root, f'{scene_id}.mp4')
+
+ # Get paths to each camera image folder
+ images_all = sorted([f for f in sorted(os.listdir(scene_folder)) if f.lower().endswith('.jpg')])
+ num_images = len(images_all)
+ num_frames = num_images // 8 # NuPlan has 8 cameras
+
+ out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
+ for frame_id in range(num_frames):
+ imgs = []
+ cam_names = []
+ for cam_id in range(8):
+ img = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 8 + cam_id]))
+ imgs.append(img)
+ cam_names.append(f"CAM_{['F0', 'L0', 'R0', 'L1', 'R1', 'L2', 'R2', 'B0'][cam_id]}")
+
+ combined_img = layout_nuplan(imgs, cam_names)
+ combined_img = cv2.resize(combined_img, (frame_width, frame_height))
+
+ out.write(combined_img)
+ out.release()
+
+def preview_camera_captures(
+ input_root: str = 'data/nuplan/processed/mini',
+ output_root: str = 'data/nuplan/preview/camera_captures',
+ frame_width: int = 1920,
+ frame_height: int = 1080,
+ fps: int = 24
+):
+ if not os.path.exists(output_root):
+ os.makedirs(output_root)
+
+ # Iterate over each scene directory in the input root directory
+ for scene_id in tqdm(sorted(os.listdir(input_root))):
+ scene_folder = os.path.join(input_root, scene_id, "images")
+
+ # Define the output video file path
+ output_video_path = os.path.join(output_root, f'{scene_id}.mp4')
+
+ # Get paths to each camera image folder
+ images_all = sorted([f for f in sorted(os.listdir(scene_folder)) if f.lower().endswith('.jpg')])
+ num_images = len(images_all)
+ num_frames = num_images // 8 # NuPlan has 8 cameras
+
+ out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
+ for frame_id in range(num_frames):
+ imgs = []
+ cam_names = []
+ for cam_id in range(8):
+ img = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 8 + cam_id]))
+ imgs.append(img)
+ cam_names.append(f"CAM_{['F0', 'L0', 'R0', 'L1', 'R1', 'L2', 'R2', 'B0'][cam_id]}")
+
+ combined_img = layout_nuplan(imgs, cam_names)
+ combined_img = cv2.resize(combined_img, (frame_width, frame_height))
+
+ out.write(combined_img)
+ out.release()
+
+def preview_dynamic_masks(
+ input_root: str = 'data/nuplan/processed/mini',
+ output_root: str = 'data/nuplan/preview/dynamic_masks',
+ frame_width: int = 1920,
+ frame_height: int = 1080,
+ fps: int = 24,
+ dynamic_class: str = 'all' # 'all' or 'vehicle' or 'human'
+):
+ if not os.path.exists(output_root):
+ os.makedirs(output_root)
+
+ # Iterate over each scene directory in the input root directory
+ for scene_id in tqdm(sorted(os.listdir(input_root))):
+ scene_folder = os.path.join(input_root, scene_id, "dynamic_masks", dynamic_class)
+
+ # Define the output video file path
+ output_video_path = os.path.join(output_root, f'{scene_id}.mp4')
+
+ # Get paths to each camera image folder
+ images_all = sorted([f for f in sorted(os.listdir(scene_folder)) if f.lower().endswith('.png')])
+ num_images = len(images_all)
+ num_frames = num_images // 8 # NuPlan has 8 cameras
+
+ out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
+ for frame_id in range(num_frames):
+ imgs = []
+ cam_names = []
+ for cam_id in range(8):
+ img = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 8 + cam_id]))
+ imgs.append(img)
+ cam_names.append(f"CAM_{['F0', 'L0', 'R0', 'L1', 'R1', 'L2', 'R2', 'B0'][cam_id]}")
+
+ combined_img = layout_nuplan(imgs, cam_names)
+ combined_img = cv2.resize(combined_img, (frame_width, frame_height))
+
+ out.write(combined_img)
+ out.release()
+
+def cat_all_videos(
+ keys: list,
+ output_root: str = 'data/nuplan/preview/catted_videos',
+ scene_id: str = 'all',
+ fps: int = 24
+):
+ if not os.path.exists(output_root):
+ os.makedirs(output_root)
+
+ if scene_id == 'all':
+ scene_ids = sorted(os.listdir('data/nuplan/processed/mini'))
+ else:
+ scene_ids = [scene_id]
+
+ for scene_id in scene_ids:
+ video_paths = []
+ for key in keys:
+ video_paths.append(f'data/nuplan/preview/{key}/{scene_id}.mp4')
+
+ output_video_path = os.path.join(output_root, f'{scene_id}.mp4')
+
+ # use moviepy to concatenate videos
+ clips = [VideoFileClip(path) for path in video_paths]
+ # cat videos from top to bottom, vertically
+ final_clip = clips_array([[clip] for clip in clips])
+ final_clip.write_videofile(output_video_path, fps=fps)
+ final_clip.close()
+
+if __name__ == '__main__':
+ # preview_camera_captures(fps=24)
+ # preview_dynamic_masks(fps=24)
+ preview_instances(fps=24)
+ # cat_all_videos(keys=['instances'], scene_id='all', fps=24)
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/nuplan/nuplan_sourceloader.py b/worldbench/third_party/drivestudio/datasets/nuplan/nuplan_sourceloader.py
new file mode 100644
index 0000000..bdd4947
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/nuplan/nuplan_sourceloader.py
@@ -0,0 +1,492 @@
+import json
+import logging
+import os
+from typing import Dict
+
+import joblib
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from pytorch3d.transforms import matrix_to_quaternion
+from torch import Tensor
+from tqdm import tqdm, trange
+
+from datasets.base.scene_dataset import ModelType
+from datasets.base.lidar_source import SceneLidarSource
+from datasets.base.pixel_source import CameraData, ScenePixelSource
+
+logger = logging.getLogger()
+
+# define each class's node type
+OBJECT_CLASS_NODE_MAPPING = {
+ "vehicle": ModelType.RigidNodes,
+ "pedestrian": ModelType.SMPLNodes,
+ "bicycle": ModelType.DeformableNodes
+}
+SMPLNODE_CLASSES = ["pedestrian"]
+
+# OpenCV to Dataset coordinate transformation
+# opencv coordinate system: x right, y down, z front
+OPENCV2DATASET = np.array(
+ [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]
+)
+
+# NuPlan Camera List:
+# 0: CAM_F0
+# 1: CAM_L0
+# 2: CAM_R0
+# 3: CAM_L1
+# 4: CAM_R1
+# 5: CAM_L2
+# 6: CAM_R2
+# 7: CAM_B0
+AVAILABLE_CAM_LIST = [0, 1, 2, 3, 4, 5, 6, 7]
+
+class NuPlanCameraData(CameraData):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ def load_calibrations(self):
+ """
+ Load the camera intrinsics, extrinsics, timestamps, etc.
+ Compute the camera-to-world matrices, ego-to-world matrices, etc.
+ """
+ # load camera intrinsics
+ # 1d Array of [f_u, f_v, c_u, c_v, k{1, 2}, p{1, 2}, k{3}].
+ # ====!! we did not use distortion parameters for simplicity !!====
+ # to be improved!!
+ intrinsic = np.loadtxt(
+ os.path.join(self.data_path, "intrinsics", f"{self.cam_id}.txt")
+ )
+ fx, fy, cx, cy = intrinsic[0], intrinsic[1], intrinsic[2], intrinsic[3]
+ k1, k2, p1, p2, k3 = intrinsic[4], intrinsic[5], intrinsic[6], intrinsic[7], intrinsic[8]
+ # scale intrinsics w.r.t. load size
+ fx, fy = (
+ fx * self.load_size[1] / self.original_size[1],
+ fy * self.load_size[0] / self.original_size[0],
+ )
+ cx, cy = (
+ cx * self.load_size[1] / self.original_size[1],
+ cy * self.load_size[0] / self.original_size[0],
+ )
+ _intrinsics = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
+ _distortions = np.array([k1, k2, p1, p2, k3])
+
+ # load camera extrinsics
+ cam_to_ego = np.loadtxt(
+ os.path.join(self.data_path, "extrinsics", f"{self.cam_id}.txt")
+ )
+ cam_to_ego = cam_to_ego @ OPENCV2DATASET
+
+ # compute per-image poses and intrinsics
+ cam_to_worlds, ego_to_worlds = [], []
+ intrinsics, distortions = [], []
+
+ # we tranform the camera poses w.r.t. the first timestep to make the translation vector of
+ # the first ego pose as the origin of the world coordinate system.
+ ego_to_world_start = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{self.start_timestep:03d}.txt")
+ )
+ for t in range(self.start_timestep, self.end_timestep):
+ ego_to_world_current = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{t:03d}.txt")
+ )
+ # compute ego_to_world transformation
+ ego_to_world = np.linalg.inv(ego_to_world_start) @ ego_to_world_current
+ ego_to_worlds.append(ego_to_world)
+ cam2world = ego_to_world @ cam_to_ego
+ cam_to_worlds.append(cam2world)
+ intrinsics.append(_intrinsics)
+ distortions.append(_distortions)
+
+ self.intrinsics = torch.from_numpy(np.stack(intrinsics, axis=0)).float()
+ self.distortions = torch.from_numpy(np.stack(distortions, axis=0)).float()
+ self.cam_to_worlds = torch.from_numpy(np.stack(cam_to_worlds, axis=0)).float()
+
+ @classmethod
+ def get_camera2worlds(cls, data_path: str, cam_id: str, start_timestep: int, end_timestep: int) -> torch.Tensor:
+ """
+ Returns camera-to-world matrices for the specified camera and time range.
+
+ Args:
+ data_path (str): Path to the dataset.
+ cam_id (str): Camera ID.
+ start_timestep (int): Start timestep.
+ end_timestep (int): End timestep.
+
+ Returns:
+ torch.Tensor: Camera-to-world matrices of shape (num_frames, 4, 4).
+ """
+ # Load the camera-to-ego transformation (this is constant for NuPlan)
+ cam_to_ego = np.loadtxt(
+ os.path.join(data_path, "extrinsics", f"{cam_id}.txt")
+ )
+ cam_to_ego = cam_to_ego @ OPENCV2DATASET
+
+ # Load the initial ego pose to align the world coordinate system
+ ego_to_world_start = np.loadtxt(
+ os.path.join(data_path, "ego_pose", f"{start_timestep:03d}.txt")
+ )
+
+ cam_to_worlds = []
+
+ for t in range(start_timestep, end_timestep):
+ # Load ego pose for each timestep
+ ego_to_world_current = np.loadtxt(
+ os.path.join(data_path, "ego_pose", f"{t:03d}.txt")
+ )
+
+ # Compute ego_to_world transformation relative to the start
+ ego_to_world = np.linalg.inv(ego_to_world_start) @ ego_to_world_current
+
+ # Compute camera-to-world transformation
+ cam2world = ego_to_world @ cam_to_ego
+
+ cam_to_worlds.append(cam2world)
+
+ return torch.from_numpy(np.stack(cam_to_worlds, axis=0)).float()
+
+class NuPlanPixelSource(ScenePixelSource):
+ def __init__(
+ self,
+ dataset_name: str,
+ pixel_data_config: OmegaConf,
+ data_path: str,
+ start_timestep: int,
+ end_timestep: int,
+ device: torch.device = torch.device("cpu"),
+ ):
+ super().__init__(dataset_name, pixel_data_config, device=device)
+ self.data_path = data_path
+ self.start_timestep = start_timestep
+ self.end_timestep = end_timestep
+ self.load_data()
+
+ def load_cameras(self):
+ self._timesteps = torch.arange(self.start_timestep, self.end_timestep)
+ self.register_normalized_timestamps()
+
+ for idx, cam_id in enumerate(self.camera_list):
+ logger.info(f"Loading camera {cam_id}")
+ camera = NuPlanCameraData(
+ dataset_name=self.dataset_name,
+ data_path=self.data_path,
+ cam_id=cam_id,
+ start_timestep=self.start_timestep,
+ end_timestep=self.end_timestep,
+ load_dynamic_mask=self.data_cfg.load_dynamic_mask,
+ load_sky_mask=self.data_cfg.load_sky_mask,
+ downscale_when_loading=self.data_cfg.downscale_when_loading[idx],
+ undistort=self.data_cfg.undistort,
+ buffer_downscale=self.buffer_downscale,
+ device=self.device,
+ )
+ camera.load_time(self.normalized_time)
+ unique_img_idx = torch.arange(len(camera), device=self.device) * len(self.camera_list) + idx
+ camera.set_unique_ids(
+ unique_cam_idx = idx,
+ unique_img_idx = unique_img_idx
+ )
+ logger.info(f"Camera {camera.cam_name} loaded.")
+ self.camera_data[cam_id] = camera
+
+ def load_objects(self):
+ """
+ get ground truth bounding boxes of the dynamic objects
+
+ instances_info = {
+ "0": # simplified instance id
+ {
+ "id": str,
+ "class_name": str,
+ "frame_annotations": {
+ "frame_idx": List,
+ "obj_to_world": List,
+ "box_size": List,
+ },
+ ...
+ }
+ frame_instances = {
+ "0": # frame idx
+ List[int] # list of simplified instance ids
+ ...
+ }
+ """
+ instances_info_path = os.path.join(self.data_path, "instances", "instances_info.json")
+ frame_instances_path = os.path.join(self.data_path, "instances", "frame_instances.json")
+ with open(instances_info_path, "r") as f:
+ instances_info = json.load(f)
+ with open(frame_instances_path, "r") as f:
+ frame_instances = json.load(f)
+ # get pose of each instance at each frame
+ # shape (num_frames, num_instances, 4, 4)
+ num_instances = len(instances_info)
+ num_full_frames = len(frame_instances)
+ instances_pose = np.zeros((num_full_frames, num_instances, 4, 4))
+ instances_size = np.zeros((num_full_frames, num_instances, 3))
+ instances_true_id = np.arange(num_instances)
+ instances_model_types = np.ones(num_instances) * -1
+
+ ego_to_world_start = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{self.start_timestep:03d}.txt")
+ )
+ for k, v in instances_info.items():
+ instances_model_types[int(k)] = OBJECT_CLASS_NODE_MAPPING[v["class_name"]]
+ for frame_idx, obj_to_world, box_size in zip(v["frame_annotations"]["frame_idx"], v["frame_annotations"]["obj_to_world"], v["frame_annotations"]["box_size"]):
+ # the first ego pose as the origin of the world coordinate system.
+ obj_to_world = np.array(obj_to_world).reshape(4, 4)
+ obj_to_world = np.linalg.inv(ego_to_world_start) @ obj_to_world
+ instances_pose[frame_idx, int(k)] = np.array(obj_to_world)
+ instances_size[frame_idx, int(k)] = np.array(box_size)
+
+ # get frame valid instances
+ # shape (num_frames, num_instances)
+ per_frame_instance_mask = np.zeros((num_full_frames, num_instances))
+ for frame_idx, valid_instances in frame_instances.items():
+ per_frame_instance_mask[int(frame_idx), valid_instances] = 1
+
+ # select the frames that are in the range of start_timestep and end_timestep
+ instances_pose = torch.from_numpy(instances_pose[self.start_timestep:self.end_timestep]).float()
+ instances_size = torch.from_numpy(instances_size[self.start_timestep:self.end_timestep]).float()
+ instances_true_id = torch.from_numpy(instances_true_id).long()
+ instances_model_types = torch.from_numpy(instances_model_types).long()
+ per_frame_instance_mask = torch.from_numpy(per_frame_instance_mask[self.start_timestep:self.end_timestep]).bool()
+
+ # filter out the instances that are not visible in selected frames
+ ins_frame_cnt = per_frame_instance_mask.sum(dim=0)
+ instances_pose = instances_pose[:, ins_frame_cnt > 0]
+ instances_size = instances_size[:, ins_frame_cnt > 0]
+ instances_true_id = instances_true_id[ins_frame_cnt > 0]
+ instances_model_types = instances_model_types[ins_frame_cnt > 0]
+ per_frame_instance_mask = per_frame_instance_mask[:, ins_frame_cnt > 0]
+
+ # assign to the class
+ # (num_frames, num_instances, 4, 4)
+ self.instances_pose = instances_pose
+ # (num_instances, 3)
+ self.instances_size = instances_size.sum(0) / per_frame_instance_mask.sum(0).unsqueeze(-1)
+ # (num_frames, num_instances)
+ self.per_frame_instance_mask = per_frame_instance_mask
+ # (num_instances)
+ self.instances_true_id = instances_true_id
+ # (num_instances)
+ self.instances_model_types = instances_model_types
+
+ if self.data_cfg.load_smpl:
+ # Collect camera-to-world matrices for all available cameras
+ cam_to_worlds = {}
+ for cam_id in AVAILABLE_CAM_LIST:
+ cam_to_worlds[cam_id] = NuPlanCameraData.get_camera2worlds(
+ self.data_path,
+ str(cam_id),
+ self.start_timestep,
+ self.end_timestep
+ )
+
+ # load SMPL parameters
+ smpl_dict = joblib.load(os.path.join(self.data_path, "humanpose", "smpl.pkl"))
+ frame_num = self.end_timestep - self.start_timestep
+
+ smpl_human_all = {}
+ for fi in tqdm(range(self.start_timestep, self.end_timestep), desc="Loading SMPL"):
+ for instance_id, ins_smpl in smpl_dict.items():
+ if instance_id not in smpl_human_all:
+ smpl_human_all[instance_id] = {
+ "smpl_quats": torch.zeros((frame_num, 24, 4), dtype=torch.float32),
+ "smpl_trans": torch.zeros((frame_num, 3), dtype=torch.float32),
+ "smpl_betas": torch.zeros((frame_num, 10), dtype=torch.float32),
+ "frame_valid": torch.zeros((frame_num), dtype=torch.bool)
+ }
+ smpl_human_all[instance_id]["smpl_quats"][:, :, 0] = 1.0
+ if ins_smpl["valid_mask"][fi]:
+ betas = ins_smpl["smpl"]["betas"][fi]
+ smpl_human_all[instance_id]["smpl_betas"][fi - self.start_timestep] = betas
+
+ body_pose = ins_smpl["smpl"]["body_pose"][fi]
+ smpl_orient = ins_smpl["smpl"]["global_orient"][fi]
+ cam_depend = ins_smpl["selected_cam_idx"][fi].item()
+
+ c2w = cam_to_worlds[cam_depend][fi - self.start_timestep]
+ world_orient = c2w[:3, :3].to(smpl_orient.device) @ smpl_orient.squeeze()
+ smpl_quats = matrix_to_quaternion(
+ torch.cat([world_orient[None, ...], body_pose], dim=0)
+ )
+
+ ii = instances_info[str(instance_id)]['frame_annotations']["frame_idx"].index(fi)
+ o2w = np.array(
+ instances_info[str(instance_id)]['frame_annotations']["obj_to_world"][ii]
+ )
+ o2w = torch.from_numpy(
+ np.linalg.inv(ego_to_world_start) @ o2w
+ )
+ # box_size = instances_info[str(instance_id)]['frame_annotations']["box_size"][ii]
+
+ smpl_human_all[instance_id]["smpl_quats"][fi - self.start_timestep] = smpl_quats
+ smpl_human_all[instance_id]["smpl_trans"][fi - self.start_timestep] = o2w[:3, 3]
+ smpl_human_all[instance_id]["frame_valid"][fi - self.start_timestep] = True
+
+ self.smpl_human_all = smpl_human_all
+
+class NuPlanLiDARSource(SceneLidarSource):
+ def __init__(
+ self,
+ lidar_data_config: OmegaConf,
+ data_path: str,
+ start_timestep: int,
+ end_timestep: int,
+ device: torch.device = torch.device("cpu"),
+ ):
+ super().__init__(lidar_data_config, device=device)
+ self.data_path = data_path
+ self.start_timestep = start_timestep
+ self.end_timestep = end_timestep
+ self.create_all_filelist()
+ self.load_data()
+
+ def create_all_filelist(self):
+ """
+ Create a list of all the files in the dataset.
+ e.g., a list of all the lidar scans in the dataset.
+ """
+ lidar_filepaths = []
+ for t in range(self.start_timestep, self.end_timestep):
+ lidar_filepaths.append(
+ os.path.join(self.data_path, "lidar", f"{t:03d}.bin")
+ )
+ self.lidar_filepaths = np.array(lidar_filepaths)
+
+ def load_calibrations(self):
+ """
+ Load the calibration files of the dataset.
+ e.g., lidar to world transformation matrices.
+ """
+ # Note that in the Waymo Open Dataset, the lidar coordinate system is the same
+ # as the vehicle coordinate system
+ lidar_to_worlds = []
+
+ # we tranform the poses w.r.t. the first timestep to make the origin of the
+ # first ego pose as the origin of the world coordinate system.
+ ego_to_world_start = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{self.start_timestep:03d}.txt")
+ )
+ for t in range(self.start_timestep, self.end_timestep):
+ ego_to_world_current = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{t:03d}.txt")
+ )
+ # compute ego_to_world transformation
+ lidar_to_world = np.linalg.inv(ego_to_world_start) @ ego_to_world_current
+ lidar_to_worlds.append(lidar_to_world)
+
+ self.lidar_to_worlds = torch.from_numpy(
+ np.stack(lidar_to_worlds, axis=0)
+ ).float()
+
+ def load_lidar(self):
+ """
+ Load the lidar data of the dataset from the filelist.
+ """
+ origins, directions, ranges = [], [], []
+ # we simplify timestamps as the time indices
+ timesteps = []
+
+ accumulated_num_original_rays = 0
+ accumulated_num_rays = 0
+ for t in trange(
+ 0, len(self.lidar_filepaths), desc="Loading lidar", dynamic_ncols=True
+ ):
+ # each lidar_info contains an Nx4 array
+ # from left to right:
+ # points: 3d, intensity: 1d, ring: 1d, lidar_id: 1d
+ lidar_info = np.memmap(
+ self.lidar_filepaths[t],
+ dtype=np.float32,
+ mode="r",
+ ).reshape(-1, 6)
+ original_length = len(lidar_info)
+ accumulated_num_original_rays += original_length
+
+ lidar_points = torch.from_numpy(lidar_info[:, :3]).float()
+ # convert the lidar points to opencv coordinate system
+ # since lidar points already include the information of origins
+ # NOTE: here we use the pesudo origins, TODO: consider a more accurate way
+ lidar_origins = torch.zeros_like(lidar_points)
+ intensity = torch.from_numpy(lidar_info[:, 3]).float()
+ lidar_ids = torch.from_numpy(lidar_info[:, 5]).float()
+
+ # transform lidar points from lidar coordinate system to world coordinate system
+ lidar_origins = (
+ self.lidar_to_worlds[t][:3, :3] @ lidar_origins.T
+ + self.lidar_to_worlds[t][:3, 3:4]
+ ).T
+ lidar_points = (
+ self.lidar_to_worlds[t][:3, :3] @ lidar_points.T
+ + self.lidar_to_worlds[t][:3, 3:4]
+ ).T
+ # compute lidar directions
+ lidar_directions = lidar_points - lidar_origins
+ lidar_ranges = torch.norm(lidar_directions, dim=-1, keepdim=True)
+ lidar_directions = lidar_directions / lidar_ranges
+ # we use time indices as the timestamp
+ lidar_timestamp = torch.ones_like(lidar_ranges).squeeze(-1) * t
+ accumulated_num_rays += len(lidar_ranges)
+
+ origins.append(lidar_origins)
+ directions.append(lidar_directions)
+ ranges.append(lidar_ranges)
+ # we use time indices as the timestamp
+ timesteps.append(lidar_timestamp)
+
+ logger.info(
+ f"Number of lidar rays: {accumulated_num_rays} "
+ f"({accumulated_num_rays / accumulated_num_original_rays * 100:.2f}% of "
+ f"{accumulated_num_original_rays} original rays)"
+ )
+
+ self.origins = torch.cat(origins, dim=0)
+ self.directions = torch.cat(directions, dim=0)
+ self.ranges = torch.cat(ranges, dim=0)
+ self.visible_masks = torch.zeros_like(self.ranges).squeeze().bool()
+ self.colors = torch.ones_like(self.directions)
+
+ # the underscore here is important.
+ self._timesteps = torch.cat(timesteps, dim=0)
+ self.register_normalized_timestamps()
+
+ def get_lidar_rays(self, time_idx: int) -> Dict[str, Tensor]:
+ """
+ Get the of rays for rendering at the given timestep.
+ Args:
+ time_idx: the index of the lidar scan to render.
+ Returns:
+ a dict of the sampled rays.
+ """
+ origins = self.origins[self.timesteps == time_idx]
+ directions = self.directions[self.timesteps == time_idx]
+ ranges = self.ranges[self.timesteps == time_idx]
+ normalized_time = self.normalized_time[self.timesteps == time_idx]
+ return {
+ "lidar_origins": origins,
+ "lidar_viewdirs": directions,
+ "lidar_ranges": ranges,
+ "lidar_normed_time": normalized_time,
+ "lidar_mask": self.timesteps == time_idx,
+ }
+
+ def delete_invisible_pts(self) -> None:
+ """
+ Clear the unvisible points.
+ """
+ if self.visible_masks is not None:
+ num_bf = self.origins.shape[0]
+ self.origins = self.origins[self.visible_masks]
+ self.directions = self.directions[self.visible_masks]
+ self.ranges = self.ranges[self.visible_masks]
+ self._timesteps = self._timesteps[self.visible_masks]
+ self._normalized_time = self._normalized_time[self.visible_masks]
+ self.colors = self.colors[self.visible_masks]
+ logger.info(
+ f"[Lidar] {num_bf - self.visible_masks.sum()} out of {num_bf} points are cleared. {self.visible_masks.sum()} points left."
+ )
+ self.visible_masks = None
+ else:
+ logger.info("[Lidar] No unvisible points to clear.")
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/nuplan/nuplan_utils.py b/worldbench/third_party/drivestudio/datasets/nuplan/nuplan_utils.py
new file mode 100644
index 0000000..924756e
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/nuplan/nuplan_utils.py
@@ -0,0 +1,157 @@
+import sqlite3
+from typing import Generator
+
+import numpy as np
+from dataclasses import dataclass
+from pyquaternion import Quaternion
+
+from nuplan.common.actor_state.agent import Agent
+from nuplan.common.actor_state.oriented_box import OrientedBox
+from nuplan.common.actor_state.scene_object import SceneObjectMetadata
+from nuplan.common.actor_state.state_representation import StateSE2, StateVector2D
+from nuplan.common.actor_state.static_object import StaticObject
+from nuplan.common.actor_state.tracked_objects import TrackedObject
+from nuplan.common.actor_state.tracked_objects_types import AGENT_TYPES, TrackedObjectType
+from nuplan.common.utils.helpers import get_unique_incremental_track_id
+from nuplan.database.nuplan_db.query_session import execute_many, execute_one
+from nuplan.database.utils.label.utils import local2agent_type, raw_mapping
+
+@dataclass
+class SimplifiedTrackedObject:
+ track_token: str
+ category: str
+ pose: np.array
+ box_size: np.array
+
+def _parse_tracked_object_row(row: sqlite3.Row) -> SimplifiedTrackedObject:
+ """
+ A convenience method to parse a TrackedObject from a sqlite3 row.
+ :param row: The row from the DB query.
+ :return: The parsed TrackedObject.
+ """
+ category_name = row["category_name"]
+ pose = StateSE2(row["x"], row["y"], row["yaw"])
+ oriented_box = OrientedBox(pose, width=row["width"], length=row["length"], height=row["height"])
+
+ # These next two are globals
+ label_local = raw_mapping["global2local"][category_name]
+ tracked_object_type = TrackedObjectType[local2agent_type[label_local]]
+
+ if tracked_object_type in AGENT_TYPES:
+ obj = Agent(
+ tracked_object_type=tracked_object_type,
+ oriented_box=oriented_box,
+ velocity=StateVector2D(row["vx"], row["vy"]),
+ predictions=[], # to be filled in later
+ angular_velocity=np.nan,
+ metadata=SceneObjectMetadata(
+ token=row["token"].hex(),
+ track_token=row["track_token"].hex(),
+ track_id=get_unique_incremental_track_id(str(row["track_token"].hex())),
+ timestamp_us=row["timestamp"],
+ category_name=category_name,
+ ),
+ )
+ else:
+ obj = StaticObject(
+ tracked_object_type=tracked_object_type,
+ oriented_box=oriented_box,
+ metadata=SceneObjectMetadata(
+ token=row["token"].hex(),
+ track_token=row["track_token"].hex(),
+ track_id=get_unique_incremental_track_id(str(row["track_token"].hex())),
+ timestamp_us=row["timestamp"],
+ category_name=category_name,
+ ),
+ )
+
+ # object pose in 3D space
+ obj_pose3d = obj.box.center.as_matrix_3d() # StateSE2 to np.array
+ obj_pose3d[2, 3] = row["z"] # add z to the pose to get 3D pose
+
+ # box_size is [l, w, h]
+ l, w, h = obj.box.length, obj.box.width, obj.box.height
+ box_size = np.array([l, w, h])
+
+ return SimplifiedTrackedObject(
+ category=tracked_object_type.name.lower(),
+ pose=obj_pose3d,
+ box_size=box_size,
+ track_token=row["track_token"].hex()
+ )
+
+def get_egopose3d_for_lidarpc_token_from_db(log_file: str, token: str) -> np.array:
+ """
+ Get the ego pose as a StateSE2 from the db for a given lidar_pc token.
+ :param log_file: The db file to query.
+ :param token: The token for which to query the current state.
+ :return: The current ego state, as a StateSE2 object.
+ """
+ query = """
+ SELECT ep.x,
+ ep.y,
+ ep.z,
+ ep.qw,
+ ep.qx,
+ ep.qy,
+ ep.qz
+ FROM ego_pose AS ep
+ INNER JOIN lidar_pc AS lp
+ ON lp.ego_pose_token = ep.token
+ WHERE lp.token = ?
+ """
+
+ row = execute_one(query, (bytearray.fromhex(token),), log_file)
+ if row is None:
+ return None
+
+ q = Quaternion(row["qw"], row["qx"], row["qy"], row["qz"])
+ rotation = q.rotation_matrix
+ translation = np.array([row["x"], row["y"], row["z"]])
+
+ ego_pose = np.eye(4)
+ ego_pose[:3, :3] = rotation
+ ego_pose[:3, 3] = translation
+
+ return ego_pose
+
+def get_tracked_objects_for_lidarpc_token_from_db(log_file: str, token: str) -> Generator[TrackedObject, None, None]:
+ """
+ Get all tracked objects for a given lidar_pc.
+ This includes both agents and static objects.
+ The values are returned in random order.
+
+ For agents, this query will not obtain the future waypoints.
+ For that, call `get_future_waypoints_for_agents_from_db()`
+ with the tokens of the agents of interest.
+
+ :param log_file: The log file to query.
+ :param token: The lidar_pc token for which to obtain the objects.
+ :return: The tracked objects associated with the token.
+ """
+ query = """
+ SELECT c.name AS category_name,
+ lb.x,
+ lb.y,
+ lb.z,
+ lb.yaw,
+ lb.width,
+ lb.length,
+ lb.height,
+ lb.vx,
+ lb.vy,
+ lb.token,
+ lb.track_token,
+ lp.timestamp
+ FROM lidar_box AS lb
+ INNER JOIN track AS t
+ ON t.token = lb.track_token
+ INNER JOIN category AS c
+ ON c.token = t.category_token
+ INNER JOIN lidar_pc AS lp
+ ON lp.token = lb.lidar_pc_token
+ WHERE lp.token = ?
+ """
+
+ for row in execute_many(query, (bytearray.fromhex(token),), log_file):
+ yield _parse_tracked_object_row(row)
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/nuscenes/nuscenes_human_utils.py b/worldbench/third_party/drivestudio/datasets/nuscenes/nuscenes_human_utils.py
new file mode 100644
index 0000000..31a9085
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/nuscenes/nuscenes_human_utils.py
@@ -0,0 +1,216 @@
+from typing import List
+import os
+import cv2
+import json
+import logging
+import numpy as np
+from tqdm import tqdm
+
+from utils.geometry import (
+ get_corners,
+ project_camera_points_to_image
+)
+from .nuscenes_sourceloader import (
+ SMPLNODE_CLASSES,
+ OPENCV2DATASET,
+ AVAILABLE_CAM_LIST,
+)
+
+logger = logging.getLogger()
+
+CAMERA_LIST = AVAILABLE_CAM_LIST
+
+def project_human_boxes(
+ scene_dir: str,
+ camera_list: List[int],
+ save_temp=True,
+ verbose=False,
+ narrow_width_ratio=0.2,
+ fps=12
+):
+ """Project human boxes to 2D image space and save the results to pkl files
+
+ Args:
+ scene_dir: str, path to the scene directory
+ camera_list: List[int], a list of camera ids to be processed
+ save_temp: bool, whether to save the intermediate results
+ verbose: bool, whether to visualize the projected boxes
+ narrow_width_ratio: sometimes the projected boxes are too wide
+ we can narrow them with this ratio to get more accurate results
+ fps: int, FPS for the visualization video
+
+ Returns:
+ collector_all: dict, a dictionary containing the projected boxes for each camera
+ """
+ # check if the necessary directories exist
+ images_dir = f'{scene_dir}/images'
+ extrinsics_dir = f'{scene_dir}/extrinsics'
+ intrinsics_dir = f'{scene_dir}/intrinsics'
+ instances_dir = f'{scene_dir}/instances'
+ valid_paths = [images_dir, extrinsics_dir, intrinsics_dir, instances_dir]
+ for path in valid_paths:
+ assert os.path.exists(path), \
+ f"Path {path} does not exist, you need to run waymo preprocess to generate the necessary files"
+
+ # create directories for saving the results
+ save_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks'
+ if not os.path.exists(save_dir):
+ os.makedirs(save_dir)
+
+ if verbose:
+ # create directories for saving the intermediate visualization results
+ video_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks/vis'
+ if not os.path.exists(video_dir):
+ os.makedirs(video_dir)
+ per_human_img_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks/vis/images'
+ if not os.path.exists(per_human_img_dir):
+ os.makedirs(per_human_img_dir)
+
+ # load instances and frame infos
+ frame_infos = json.load(
+ open(f'{instances_dir}/frame_instances.json')
+ )
+ instances_meta = json.load(
+ open(f'{instances_dir}/instances_info.json')
+ )
+
+ collector_all = {}
+ # iterate over each camera
+ for cam_id in camera_list:
+ # check if already processed
+ pkl_path = os.path.join(save_dir, f"{cam_id}.pkl")
+ if os.path.exists(pkl_path):
+ collector_all[cam_id] = json.load(open(pkl_path))
+ logger.info(f"Results for camera {cam_id} already exists at {pkl_path}")
+ continue
+
+ if verbose:
+ per_cam_vis_dir = os.path.join(per_human_img_dir, f"{cam_id}")
+ if not os.path.exists(per_cam_vis_dir):
+ os.makedirs(per_cam_vis_dir)
+
+ collector = {}
+ frames = []
+ for frame_id, frame_ins_list in frame_infos.items():
+ frame_id = int(frame_id)
+
+ # define empty instance collector for each frame
+ frame_collector = {
+ "gt_bbox": [],
+ "extra_data": {
+ "gt_track_id": [],
+ "gt_class": [],
+ }
+ }
+
+ # Load extrinsics for this frame and camera
+ cam2world = np.loadtxt(os.path.join(extrinsics_dir, f"{frame_id:03d}_{cam_id}.txt"))
+ # Convert to OpenCV coordinate system
+ cam2world = cam2world @ OPENCV2DATASET
+
+ # load intrinsic
+ Ks = np.loadtxt(os.path.join(intrinsics_dir, f"{cam_id}.txt"))
+ fx, fy, cx, cy = Ks[0], Ks[1], Ks[2], Ks[3]
+ # k1, k2, p1, p2, k3 = Ks[4], Ks[5], Ks[6], Ks[7], Ks[8]
+ intrinsic = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float32)
+
+ # load image
+ ori_image = cv2.imread(
+ os.path.join(images_dir, f"{str(frame_id).zfill(3)}_{cam_id}.jpg")
+ )
+ image = ori_image.copy()
+ H, W = image.shape[:2]
+
+ # iterate over each instance in the frame
+ if len(frame_ins_list) > 0:
+ # if there are pedestrians in this frame, project the boxes to the image
+ image_plotted = image.copy()
+ for instance_id in frame_ins_list:
+ ins = instances_meta[str(instance_id)]
+
+ if ins["class_name"] not in SMPLNODE_CLASSES:
+ continue
+
+ ins_anno = ins["frame_annotations"]
+ index = ins_anno['frame_idx'].index(frame_id)
+ obj_to_world = np.array(ins_anno['obj_to_world'][index])
+ l, w, h = ins_anno['box_size'][index]
+
+ # get box corners in object space
+ corners = get_corners(l, w, h)
+ # transform box corners to world space
+ corners_world = obj_to_world[:3, :3] @ corners + obj_to_world[:3, 3:4]
+ # transform box corners to image space
+ world2cam = np.linalg.inv(cam2world)
+ corners_cam = world2cam[:3, :3] @ corners_world + world2cam[:3, 3:4]
+ cam_points, depth = project_camera_points_to_image(corners_cam.T, intrinsic)
+
+ x_min, y_min = np.min(cam_points, axis=0)
+ x_max, y_max = np.max(cam_points, axis=0)
+ # clip left and right with this ratio
+ if narrow_width_ratio > 0.:
+ length = x_max - x_min
+ x_min += length * narrow_width_ratio
+ x_max -= length * narrow_width_ratio
+
+ # clip the box to the image
+ original_area = (x_max - x_min) * (y_max - y_min)
+ x_min, x_max = np.clip(x_min, 0, W), np.clip(x_max, 0, W)
+ y_min, y_max = np.clip(y_min, 0, H), np.clip(y_max, 0, H)
+ new_area = (x_max - x_min) * (y_max - y_min)
+
+ # filter out boxes that are too small or too large
+ behind = depth.max() < 0
+ too_small = new_area < W * H * (0.03)**2
+ too_large = new_area > W * H / 1.1
+ too_far = np.linalg.norm(obj_to_world[:3, 3] - cam2world[:3, 3]) > 40
+ clip_large = new_area / original_area < 1/3
+ if too_small or too_large or clip_large or behind or too_far:
+ continue
+
+ # gt box on image
+ gt_box = [x_min, y_min, x_max - x_min, y_max - y_min]
+
+ # save the projected box to the collector
+ frame_collector["gt_bbox"].append(gt_box)
+ frame_collector["extra_data"]["gt_track_id"].append(instance_id)
+ frame_collector["extra_data"]["gt_class"].append([0])
+
+ if verbose:
+ # visualize the projected boxes of ONE instance
+ raw_image = cv2.rectangle(
+ ori_image.copy(), (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 2
+ )
+ raw_image_path = os.path.join(per_cam_vis_dir, f"{frame_id}_{instance_id}.jpg")
+ cv2.imwrite(raw_image_path, raw_image)
+
+ # add this instance to the image
+ image_plotted = cv2.rectangle(
+ image, (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 2
+ )
+
+ if verbose:
+ frames.append(image_plotted)
+ else:
+ # if no instance in this frame, just save the original image
+ if verbose:
+ frames.append(ori_image)
+
+ collector[frame_id] = frame_collector
+
+ if verbose:
+ height, width = frames[0].shape[:2]
+ output_path = os.path.join(video_dir, f"cam_{cam_id}.mp4")
+ fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+ out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+ for frame in tqdm(frames, desc=f"Writing video for camera {cam_id}"):
+ out.write(frame)
+ out.release()
+
+ if save_temp:
+ # save collector to pkl
+ json.dump(collector, open(pkl_path, "w"))
+
+ collector_all[cam_id] = collector
+
+ return collector_all
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/nuscenes/nuscenes_preprocess.py b/worldbench/third_party/drivestudio/datasets/nuscenes/nuscenes_preprocess.py
new file mode 100644
index 0000000..bb2ca18
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/nuscenes/nuscenes_preprocess.py
@@ -0,0 +1,1085 @@
+import json
+import os
+from collections import Counter
+from typing import List
+
+import cv2
+import numpy as np
+from PIL import Image
+from pyquaternion import Quaternion
+from tqdm import tqdm
+
+from nuscenes.nuscenes import LidarPointCloud, NuScenes
+from nuscenes.utils.data_classes import Box
+from nuscenes.utils.geometry_utils import view_points
+
+from datasets.tools.multiprocess_utils import track_parallel_progress
+from utils.geometry import get_corners
+from utils.visualization import color_mapper, dump_3d_bbox_on_image
+
+NUSCENES_LABELS = [
+ 'animal',
+ 'human.pedestrian.adult',
+ 'human.pedestrian.child',
+ 'human.pedestrian.construction_worker',
+ 'human.pedestrian.personal_mobility',
+ 'human.pedestrian.police_officer',
+ 'human.pedestrian.stroller',
+ 'human.pedestrian.wheelchair',
+ 'movable_object.barrier',
+ 'movable_object.debris',
+ 'movable_object.pushable_pullable',
+ 'movable_object.trafficcone',
+ 'static_object.bicycle_rack',
+ 'vehicle.bicycle',
+ 'vehicle.bus.bendy',
+ 'vehicle.bus.rigid',
+ 'vehicle.car',
+ 'vehicle.construction',
+ 'vehicle.emergency.ambulance',
+ 'vehicle.emergency.police',
+ 'vehicle.motorcycle',
+ 'vehicle.trailer',
+ 'vehicle.truck'
+]
+
+NUSCENES_NONRIGID_DYNAMIC_CLASSES = [
+ 'animal',
+ 'human.pedestrian.adult',
+ 'human.pedestrian.child',
+ 'human.pedestrian.construction_worker',
+ 'human.pedestrian.personal_mobility',
+ 'human.pedestrian.police_officer',
+ 'human.pedestrian.stroller',
+ 'human.pedestrian.wheelchair',
+ 'vehicle.bicycle',
+ 'vehicle.motorcycle',
+]
+
+NUSCENES_RIGID_DYNAMIC_CLASSES = [
+ 'vehicle.bus.bendy',
+ 'vehicle.bus.rigid',
+ 'vehicle.car',
+ 'vehicle.construction',
+ 'vehicle.emergency.ambulance',
+ 'vehicle.emergency.police',
+ 'vehicle.trailer',
+ 'vehicle.truck'
+]
+
+NUSCENES_DYNAMIC_CLASSES = NUSCENES_NONRIGID_DYNAMIC_CLASSES + NUSCENES_RIGID_DYNAMIC_CLASSES
+
+class NuScenesProcessor(object):
+ """Process NuScenes dataset.
+
+ NuScenes Data Frequencies:
+ - CAMERA: 12 Hz
+ - LIDAR: 20 Hz
+ - Annotated Keyframes: 2 Hz
+
+ This function processes the NuScenes dataset, offering two processing modes:
+
+ 1. Original Mode (2 Hz):
+ - Uses only the annotated keyframes.
+ - Sparse but contains original annotations.
+ - Default mode when 'interpolate_N' is set to 0.
+
+ 2. Interpolated Mode:
+ - Interpolate N frames between annotated keyframes, where N <= 4 is recommended.
+ - Provides denser data by interpolating between annotated frames.
+ - Activated by setting 'interpolate_N' to a positive integer.
+ - Note: Using N > 4 may lead to insufficient image frames due to potential frame drops,
+ despite the camera's 12 Hz capture rate.
+
+ Interpolated mode may enhance training due to higher temporal resolution,
+ but uses estimated annotations for non-keyframes. Choose the mode that
+ best fits your use case.
+
+ Args:
+ load_dir (str): Source directory of NuScenes data.
+ save_dir (str): Target directory for processed data.
+ split (str): Dataset split to process. Default: 'v1.0-mini'.
+ interpolate_N (int): Number of frames to interpolate between keyframes.
+ If 0, use original 2 Hz data.
+ workers (int): Number of parallel processing workers.
+ process_keys (list): Data types to process.
+ Default: ["images", "lidar", "calib", "dynamic_masks", "objects"]
+ process_id_list (list): Specific scene IDs to process.
+ """
+
+ def __init__(
+ self,
+ load_dir,
+ save_dir,
+ split='v1.0-mini',
+ interpolate_N=0,
+ process_keys=[
+ "images",
+ "lidar",
+ "calib",
+ "dynamic_masks",
+ "objects"
+ ],
+ process_id_list=None,
+ workers=64,
+ ):
+ self.process_id_list = process_id_list
+ self.process_keys = process_keys
+ print("will process keys: ", self.process_keys)
+ self.interpolate_N = interpolate_N
+ assert self.interpolate_N <= 4, \
+ "Interpolation frames should be less than 4. \or there will be frame drop issue."
+ if self.interpolate_N:
+ print(f"We will interpolate {self.interpolate_N} frames between keyframes.")
+
+ # NuScenes Provides 6 cameras
+ self.cam_list = [ # {frame_idx}_{cam_id}.jpg
+ "CAM_FRONT", # "xxx_0.jpg"
+ "CAM_FRONT_LEFT", # "xxx_1.jpg"
+ "CAM_FRONT_RIGHT", # "xxx_2.jpg"
+ "CAM_BACK_LEFT", # "xxx_3.jpg"
+ "CAM_BACK_RIGHT", # "xxx_4.jpg"
+ "CAM_BACK" # "xxx_5.jpg"
+ ]
+ # For each keyframe, the total number of points from LIDAR_TOP is ~35000
+ # the total number of points from ALL RADAR is ~300
+ # Since Radar points are too sparse, we consider only LIDAR_TOP
+ self.lidar_list = ['LIDAR_TOP']
+
+ self.load_dir = load_dir
+ post_fix = f"_{(interpolate_N+1)*2}Hz" if interpolate_N > 0 else ""
+ save_dir = save_dir.replace("processed", "processed"+post_fix)
+ self.save_dir = os.path.join(save_dir, split.split('-')[-1])
+ self.workers = int(workers)
+ self.nusc = NuScenes(
+ version=split, dataroot=load_dir, verbose=True
+ )
+ self.create_folder()
+
+ def convert(self):
+ """Convert action."""
+ print("Start converting ...")
+ if self.process_id_list is None:
+ id_list = range(len(self))
+ else:
+ id_list = self.process_id_list
+ if self.interpolate_N > 0:
+ convert_fn = self.convert_one_interpolated
+ else:
+ convert_fn = self.convert_one
+ track_parallel_progress(convert_fn, id_list, self.workers)
+ print("\nFinished ...")
+
+ def convert_one(self, scene_idx):
+ """Convert action for single file."""
+ scene = self.nusc.scene[scene_idx]
+ scene_data = self.nusc.get('scene', scene['token'])
+ if "images" in self.process_keys:
+ self.save_image(scene_data, scene_idx)
+ print(f"Processed images for scene {str(scene_idx).zfill(3)}")
+ if "calib" in self.process_keys:
+ self.save_calib(scene_data, scene_idx)
+ print(f"Processed calib for scene {str(scene_idx).zfill(3)}")
+ if "lidar" in self.process_keys:
+ self.save_lidar(scene_data, scene_idx)
+ print(f"Processed lidar for scene {str(scene_idx).zfill(3)}")
+ if "dynamic_masks" in self.process_keys:
+ self.save_dynamic_mask(scene_data, scene_idx, class_valid='all')
+ self.save_dynamic_mask(scene_data, scene_idx, class_valid='human')
+ self.save_dynamic_mask(scene_data, scene_idx, class_valid='vehicle')
+ print(f"Processed dynamic masks for scene {str(scene_idx).zfill(3)}")
+
+ # process annotated objects
+ if "objects" in self.process_keys:
+ instances_info, frame_instances = self.save_objects(scene_data)
+ print(f"Processed objects for scene {str(scene_idx).zfill(3)}")
+
+ # Save instances info and frame instances
+ instances_info_save_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/instances"
+ with open(f"{instances_info_save_path}/instances_info.json", "w") as fp:
+ json.dump(instances_info, fp, indent=4)
+ with open(f"{instances_info_save_path}/frame_instances.json", "w") as fp:
+ json.dump(frame_instances, fp, indent=4)
+
+ if "objects_vis" in self.process_keys:
+ self.visualize_dynamic_objects(scene_data, scene_idx, instances_info, frame_instances)
+ print(f"Processed objects visualization for scene {str(scene_idx).zfill(3)}")
+
+ def convert_one_interpolated(self, scene_idx):
+ """Convert action for single file."""
+ scene = self.nusc.scene[scene_idx]
+ scene_data = self.nusc.get('scene', scene['token'])
+ keyframe_timestamps = self.get_keyframe_timestamps(scene_data)
+ interpolated_timestamps = self.get_interpolated_timestamps(keyframe_timestamps, self.interpolate_N)
+ interpolated_timestamps = np.array(interpolated_timestamps, dtype=np.int64)
+ if "images" in self.process_keys:
+ self.save_image_interpolated(scene_data, scene_idx, interpolated_timestamps)
+ print(f"Processed images for scene {str(scene_idx).zfill(3)}")
+ if "calib" in self.process_keys:
+ self.save_calib_interpolated(scene_data, scene_idx, interpolated_timestamps)
+ print(f"Processed calib for scene {str(scene_idx).zfill(3)}")
+ if "lidar" in self.process_keys:
+ self.save_lidar_interpolated(scene_data, scene_idx, interpolated_timestamps)
+ print(f"Processed lidar for scene {str(scene_idx).zfill(3)}")
+
+ # process annotated objects
+ if "objects" in self.process_keys:
+ instances_info, _ = self.save_objects(scene_data)
+ print(f"Processed objects for scene {str(scene_idx).zfill(3)}")
+
+ # interpolate the instances info
+ instances_info, frame_instances = self.interpolate_boxes(instances_info)
+ print(f"Interpolated objects for scene {str(scene_idx).zfill(3)}")
+
+ # Save instances info and frame instances
+ instances_info_save_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/instances"
+ with open(f"{instances_info_save_path}/instances_info.json", "w") as fp:
+ json.dump(instances_info, fp, indent=4)
+ with open(f"{instances_info_save_path}/frame_instances.json", "w") as fp:
+ json.dump(frame_instances, fp, indent=4)
+
+ if "objects_vis" in self.process_keys:
+ self.visualize_dynamic_objects_interpolated(scene_data, scene_idx, interpolated_timestamps, instances_info, frame_instances)
+ print(f"Processed objects visualization for scene {str(scene_idx).zfill(3)}")
+
+ if "dynamic_masks" in self.process_keys:
+ self.save_dynamic_mask_interpolated(scene_data, scene_idx, interpolated_timestamps, instances_info, class_valid='all')
+ self.save_dynamic_mask_interpolated(scene_data, scene_idx, interpolated_timestamps, instances_info, class_valid='human')
+ self.save_dynamic_mask_interpolated(scene_data, scene_idx, interpolated_timestamps, instances_info, class_valid='vehicle')
+ print(f"Processed dynamic masks for scene {str(scene_idx).zfill(3)}")
+
+ def __len__(self):
+ """Length of the filename list."""
+ return len(self.process_id_list)
+
+ def get_percam_meta(self, scene_data):
+ """Get the timestamp for each image."""
+ first_sample_token, last_sample_token = scene_data['first_sample_token'], scene_data['last_sample_token']
+ first_sample_record = self.nusc.get('sample', first_sample_token)
+ cams_mata = {
+ cam_name: {
+ "timestamps": [],
+ "is_key_frame": [],
+ "token": []
+ } for cam_name in self.cam_list
+ }
+
+ for cam_name in self.cam_list:
+ cur_img_data = self.nusc.get('sample_data', first_sample_record['data'][cam_name])
+ while True:
+ cams_mata[cam_name]["timestamps"].append(cur_img_data['timestamp'])
+ cams_mata[cam_name]["is_key_frame"].append(cur_img_data['is_key_frame'])
+ cams_mata[cam_name]["token"].append(cur_img_data['token'])
+
+ if cur_img_data['sample_token'] == last_sample_token:
+ break
+ cur_img_data = self.nusc.get('sample_data', cur_img_data['next'])
+
+ return cams_mata
+
+ def get_keyframe_timestamps(self, scene_data):
+ """Get timestamps for keyframes in the scene."""
+ first_sample_token = scene_data['first_sample_token']
+ last_sample_token = scene_data['last_sample_token']
+ curr_sample_record = self.nusc.get('sample', first_sample_token)
+
+ keyframe_timestamps = []
+
+ while True:
+ # Add the timestamp of the current keyframe
+ keyframe_timestamps.append(curr_sample_record['timestamp'])
+
+ if curr_sample_record['token'] == last_sample_token:
+ break
+
+ # Move to the next keyframe
+ curr_sample_record = self.nusc.get('sample', curr_sample_record['next'])
+
+ return keyframe_timestamps
+
+ def get_interpolated_timestamps(self, keyframe_timestamps: List[int], N):
+ """Interpolate timestamps between keyframes."""
+ interpolated_timestamps = []
+
+ for i in range(len(keyframe_timestamps) - 1):
+ start_time = keyframe_timestamps[i]
+ end_time = keyframe_timestamps[i + 1]
+
+ # Calculate the time step for interpolation
+ time_step = (end_time - start_time) / (N + 1)
+
+ # Add the start timestamp
+ interpolated_timestamps.append(start_time)
+
+ # Add N interpolated timestamps
+ for j in range(1, N + 1):
+ interpolated_time = start_time + j * time_step
+ interpolated_timestamps.append(int(interpolated_time))
+
+ # Add the last keyframe timestamp
+ interpolated_timestamps.append(keyframe_timestamps[-1])
+
+ return interpolated_timestamps
+
+ def find_cloest_lidar_tokens(self, scene_data, timestamps: List[int]):
+ """Find the closest LiDAR tokens for given timestamps."""
+ first_sample_token = scene_data['first_sample_token']
+ first_sample_record = self.nusc.get('sample', first_sample_token)
+ lidar_token = first_sample_record['data'][self.lidar_list[0]]
+ lidar_data = self.nusc.get('sample_data', lidar_token)
+
+ # Collect all LiDAR timestamps and tokens
+ lidar_timestamps = []
+ lidar_tokens = []
+ current_lidar = lidar_data
+ while True:
+ lidar_timestamps.append(current_lidar['timestamp'])
+ lidar_tokens.append(current_lidar['token'])
+ if current_lidar['next'] == '':
+ break
+ current_lidar = self.nusc.get('sample_data', current_lidar['next'])
+
+ lidar_timestamps = np.array(lidar_timestamps, dtype=np.int64)
+
+ # Find closest LiDAR tokens for each timestamp
+ closest_tokens = []
+ for timestamp in timestamps:
+ idx = np.argmin(np.abs(lidar_timestamps - timestamp))
+ closest_tokens.append(lidar_tokens[idx])
+
+ # DEBUG USAGE: find is there any duplicated tokens
+ # if len(closest_tokens) != len(set(closest_tokens)):
+ # duplicates = [token for token, count in Counter(closest_tokens).items() if count > 1]
+ # print(f"\nWARNING: {len(duplicates)} duplicated tokens in lidar")
+
+ return closest_tokens
+
+ def find_closest_img_tokens(self, scene_data, timestamps: List[int], cam_name):
+ """Find the closest image tokens for given timestamps for a specific camera."""
+ first_sample_token = scene_data['first_sample_token']
+ first_sample_record = self.nusc.get('sample', first_sample_token)
+ img_token = first_sample_record['data'][cam_name]
+ img_data = self.nusc.get('sample_data', img_token)
+
+ # Collect all image timestamps and tokens for the specified camera
+ img_timestamps = []
+ img_tokens = []
+ current_img = img_data
+ while True:
+ img_timestamps.append(current_img['timestamp'])
+ img_tokens.append(current_img['token'])
+ if current_img['next'] == '':
+ break
+ current_img = self.nusc.get('sample_data', current_img['next'])
+
+ img_timestamps = np.array(img_timestamps, dtype=np.int64)
+
+ # Find closest image tokens for each timestamp
+ closest_tokens = []
+ for timestamp in timestamps:
+ idx = np.argmin(np.abs(img_timestamps - timestamp))
+ closest_tokens.append(img_tokens[idx])
+
+ # DEBUG USAGE: find is there any duplicated tokens
+ # if len(closest_tokens) != len(set(closest_tokens)):
+ # duplicates = [token for token, count in Counter(closest_tokens).items() if count > 1]
+ # print(f"\nWARNING: {len(duplicates)} duplicated tokens in {cam_name}")
+
+ return closest_tokens
+
+ def save_image(self, scene_data, scene_idx):
+ """Parse and save the images in jpg format."""
+ first_sample_token, last_sample_token = scene_data['first_sample_token'], scene_data['last_sample_token']
+ curr_sample_record = self.nusc.get('sample', first_sample_token)
+ key_frame_idx = 0
+
+ while True:
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ cam_data = self.nusc.get('sample_data', curr_sample_record['data'][cam_name])
+ source_path = os.path.join(self.nusc.dataroot, cam_data['filename'])
+ img_path = (
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/images/"
+ + f"{str(key_frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+ )
+ # cp the image
+ os.system(f"cp {source_path} {img_path}")
+
+ if curr_sample_record['next'] == '' or curr_sample_record['token'] == last_sample_token:
+ break
+ key_frame_idx += 1
+ curr_sample_record = self.nusc.get('sample', curr_sample_record['next'])
+
+ def save_image_interpolated(self, scene_data, scene_idx, timestamps: np.array):
+ """Parse and save the interpolated images in jpg format."""
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ # Find the closest image tokens for each timestamp
+ closest_tokens = self.find_closest_img_tokens(scene_data, timestamps, cam_name)
+
+ for frame_idx, token in enumerate(closest_tokens):
+ cam_data = self.nusc.get('sample_data', token)
+ source_path = os.path.join(self.nusc.dataroot, cam_data['filename'])
+ img_path = (
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/images/"
+ + f"{str(frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+ )
+ # Copy the image
+ os.system(f"cp {source_path} {img_path}")
+
+ def save_calib(self, scene_data, scene_idx):
+ """Parse and save the calibration data."""
+ first_sample_token, last_sample_token = scene_data['first_sample_token'], scene_data['last_sample_token']
+ curr_sample_record = self.nusc.get('sample', first_sample_token)
+ key_frame_idx = 0
+
+ while True:
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ cam_data = self.nusc.get('sample_data', curr_sample_record['data'][cam_name])
+ calib_data = self.nusc.get('calibrated_sensor', cam_data['calibrated_sensor_token'])
+
+ # Extrinsics (camera to ego)
+ extrinsics_cam_to_ego = np.eye(4)
+ extrinsics_cam_to_ego[:3, :3] = Quaternion(calib_data['rotation']).rotation_matrix
+ extrinsics_cam_to_ego[:3, 3] = np.array(calib_data['translation'])
+
+ # Get ego pose (ego to world)
+ ego_pose_data = self.nusc.get('ego_pose', cam_data['ego_pose_token'])
+ ego_to_world = np.eye(4)
+ ego_to_world[:3, :3] = Quaternion(ego_pose_data['rotation']).rotation_matrix
+ ego_to_world[:3, 3] = np.array(ego_pose_data['translation'])
+
+ # Transform camera extrinsics to world coordinates
+ extrinsics_cam_to_world = ego_to_world @ extrinsics_cam_to_ego
+
+ np.savetxt(
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/extrinsics/"
+ f"{str(key_frame_idx).zfill(3)}_{str(cam_idx)}.txt",
+ extrinsics_cam_to_world
+ )
+
+ # Intrinsics
+ intrinsics = np.array(calib_data['camera_intrinsic'])
+ # fx fy cx cy
+ Ks = [intrinsics[0, 0], intrinsics[1, 1], intrinsics[0, 2], intrinsics[1, 2], 0., 0., 0., 0., 0.]
+ np.savetxt(
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/intrinsics/"
+ f"{str(cam_idx)}.txt",
+ Ks
+ )
+
+ if curr_sample_record['next'] == '' or curr_sample_record['token'] == last_sample_token:
+ break
+ key_frame_idx += 1
+ curr_sample_record = self.nusc.get('sample', curr_sample_record['next'])
+
+ def save_calib_interpolated(self, scene_data, scene_idx, timestamps: np.array):
+ """Parse and save the interpolated calibration data."""
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ # Find the closest image tokens for each timestamp
+ closest_tokens = self.find_closest_img_tokens(scene_data, timestamps, cam_name)
+
+ for frame_idx, token in enumerate(closest_tokens):
+ cam_data = self.nusc.get('sample_data', token)
+ calib_data = self.nusc.get('calibrated_sensor', cam_data['calibrated_sensor_token'])
+
+ # Extrinsics (camera to ego)
+ extrinsics_cam_to_ego = np.eye(4)
+ extrinsics_cam_to_ego[:3, :3] = Quaternion(calib_data['rotation']).rotation_matrix
+ extrinsics_cam_to_ego[:3, 3] = np.array(calib_data['translation'])
+
+ # Get ego pose (ego to world)
+ ego_pose_data = self.nusc.get('ego_pose', cam_data['ego_pose_token'])
+ ego_to_world = np.eye(4)
+ ego_to_world[:3, :3] = Quaternion(ego_pose_data['rotation']).rotation_matrix
+ ego_to_world[:3, 3] = np.array(ego_pose_data['translation'])
+
+ # Transform camera extrinsics to world coordinates
+ extrinsics_cam_to_world = ego_to_world @ extrinsics_cam_to_ego
+
+ np.savetxt(
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/extrinsics/"
+ f"{str(frame_idx).zfill(3)}_{str(cam_idx)}.txt",
+ extrinsics_cam_to_world
+ )
+
+ # Intrinsics
+ intrinsics = np.array(calib_data['camera_intrinsic'])
+ # fx fy cx cy
+ Ks = [intrinsics[0, 0], intrinsics[1, 1], intrinsics[0, 2], intrinsics[1, 2], 0., 0., 0., 0., 0.]
+ np.savetxt(
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/intrinsics/"
+ f"{str(cam_idx)}.txt",
+ Ks
+ )
+
+ def save_lidar(self, scene_data, scene_idx):
+ """Parse and save the lidar data in bin format and lidar pose in world coordinates."""
+ first_sample_token, last_sample_token = scene_data['first_sample_token'], scene_data['last_sample_token']
+ curr_sample_record = self.nusc.get('sample', first_sample_token)
+ key_frame_idx = 0
+
+ while True:
+ lidar_token = curr_sample_record['data'][self.lidar_list[0]]
+ lidar_data = self.nusc.get('sample_data', lidar_token)
+ lidar_path = os.path.join(self.nusc.dataroot, lidar_data['filename'])
+
+ # Load point cloud
+ pc = LidarPointCloud.from_file(lidar_path)
+
+ # Get lidar extrinsics (lidar to ego)
+ calib_data = self.nusc.get('calibrated_sensor', lidar_data['calibrated_sensor_token'])
+ lidar_to_ego = np.eye(4)
+ lidar_to_ego[:3, :3] = Quaternion(calib_data['rotation']).rotation_matrix
+ lidar_to_ego[:3, 3] = np.array(calib_data['translation'])
+
+ # Save lidar points in ego frame
+ lidar_save_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/lidar/{str(key_frame_idx).zfill(3)}.bin"
+ pc.points.T.astype(np.float32).tofile(lidar_save_path)
+
+ # Get ego pose (ego to world)
+ ego_pose_data = self.nusc.get('ego_pose', lidar_data['ego_pose_token'])
+ ego_to_world = np.eye(4)
+ ego_to_world[:3, :3] = Quaternion(ego_pose_data['rotation']).rotation_matrix
+ ego_to_world[:3, 3] = np.array(ego_pose_data['translation'])
+
+ # Calculate lidar pose in world coordinates
+ lidar_to_world = ego_to_world @ lidar_to_ego
+
+ # Save lidar pose in world coordinates
+ np.savetxt(
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/lidar_pose/"
+ f"{str(key_frame_idx).zfill(3)}.txt",
+ lidar_to_world
+ )
+
+ if curr_sample_record['next'] == '' or curr_sample_record['token'] == last_sample_token:
+ break
+ key_frame_idx += 1
+ curr_sample_record = self.nusc.get('sample', curr_sample_record['next'])
+
+ def save_lidar_interpolated(self, scene_data, scene_idx, timestamps: np.array):
+ """Parse and save the interpolated lidar data in bin format and lidar pose in world coordinates."""
+ # Find the closest LiDAR tokens for each timestamp
+ closest_tokens = self.find_cloest_lidar_tokens(scene_data, timestamps)
+
+ for frame_idx, token in enumerate(closest_tokens):
+ lidar_data = self.nusc.get('sample_data', token)
+ lidar_path = os.path.join(self.nusc.dataroot, lidar_data['filename'])
+
+ # Load point cloud
+ pc = LidarPointCloud.from_file(lidar_path)
+
+ # Get lidar extrinsics (lidar to ego)
+ calib_data = self.nusc.get('calibrated_sensor', lidar_data['calibrated_sensor_token'])
+ lidar_to_ego = np.eye(4)
+ lidar_to_ego[:3, :3] = Quaternion(calib_data['rotation']).rotation_matrix
+ lidar_to_ego[:3, 3] = np.array(calib_data['translation'])
+
+ # Save lidar points in ego frame
+ lidar_save_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/lidar/{str(frame_idx).zfill(3)}.bin"
+ pc.points.T.astype(np.float32).tofile(lidar_save_path)
+
+ # Get ego pose (ego to world)
+ ego_pose_data = self.nusc.get('ego_pose', lidar_data['ego_pose_token'])
+ ego_to_world = np.eye(4)
+ ego_to_world[:3, :3] = Quaternion(ego_pose_data['rotation']).rotation_matrix
+ ego_to_world[:3, 3] = np.array(ego_pose_data['translation'])
+
+ # Calculate lidar pose in world coordinates
+ lidar_to_world = ego_to_world @ lidar_to_ego
+
+ # Save lidar pose in world coordinates
+ np.savetxt(
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/lidar_pose/"
+ f"{str(frame_idx).zfill(3)}.txt",
+ lidar_to_world
+ )
+
+ def save_dynamic_mask(self, scene_data, scene_idx, class_valid='all'):
+ """Parse and save the segmentation data."""
+ assert class_valid in ['all', 'human', 'vehicle'], "Invalid class valid"
+ if class_valid == 'all':
+ VALID_CLASSES = NUSCENES_DYNAMIC_CLASSES
+ elif class_valid == 'human':
+ VALID_CLASSES = NUSCENES_NONRIGID_DYNAMIC_CLASSES
+ elif class_valid == 'vehicle':
+ VALID_CLASSES = NUSCENES_RIGID_DYNAMIC_CLASSES
+ mask_dir = f"{self.save_dir}/{str(scene_idx).zfill(3)}/dynamic_masks/{class_valid}"
+ if not os.path.exists(mask_dir):
+ os.makedirs(mask_dir)
+
+ first_sample_token, last_sample_token = scene_data['first_sample_token'], scene_data['last_sample_token']
+ curr_sample_record = self.nusc.get('sample', first_sample_token)
+ key_frame_idx = 0
+
+ while True:
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ cam_data = self.nusc.get('sample_data', curr_sample_record['data'][cam_name])
+ img_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/images/{str(key_frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+
+ img = cv2.imread(img_path)
+ dynamic_mask = np.zeros(img.shape[:2], dtype=np.float32)
+
+ anns = [self.nusc.get('sample_annotation', token) for token in curr_sample_record['anns']]
+ valid_anns = [ann for ann in anns if ann['category_name'] in VALID_CLASSES]
+
+ # Get camera calibration data
+ cs_record = self.nusc.get('calibrated_sensor', cam_data['calibrated_sensor_token'])
+ camera_intrinsic = np.array(cs_record['camera_intrinsic'])
+
+ # Get ego pose
+ pose_record = self.nusc.get('ego_pose', cam_data['ego_pose_token'])
+
+ # Project 3D boxes to 2D and create mask
+ for ann in valid_anns:
+ # Create Box object
+ box = Box(ann['translation'], ann['size'], Quaternion(ann['rotation']))
+
+ # Move box to ego vehicle coordinate system
+ box.translate(-np.array(pose_record['translation']))
+ box.rotate(Quaternion(pose_record['rotation']).inverse)
+
+ # Move box to sensor coordinate system
+ box.translate(-np.array(cs_record['translation']))
+ box.rotate(Quaternion(cs_record['rotation']).inverse)
+
+ # Project 3D box to 2D
+ corners_3d = box.corners()
+ corners_2d = view_points(corners_3d, camera_intrinsic, normalize=True)
+
+ # Check if the object is in front of the camera and all corners are in the image
+ # NOTE: we use strict visibility check here, requiring all corners to be visible
+ in_front = np.all(corners_3d[2, :] > 0.1)
+ in_image = np.all(corners_2d[0, :] >= 0) & np.all(corners_2d[0, :] < img.shape[1]) & \
+ np.all(corners_2d[1, :] >= 0) & np.all(corners_2d[1, :] < img.shape[0])
+ if not (in_front and in_image):
+ continue
+
+ # If valid, extract x and y coordinates
+ corners_2d = corners_2d[:2, :]
+
+ # Fill the mask
+ u, v = corners_2d[0, :].astype(np.int32), corners_2d[1, :].astype(np.int32)
+ u = np.clip(u, 0, img.shape[1] - 1)
+ v = np.clip(v, 0, img.shape[0] - 1)
+
+ if u.max() - u.min() == 0 or v.max() - v.min() == 0:
+ continue
+
+ xy = (u.min(), v.min())
+ width = u.max() - u.min()
+ height = v.max() - v.min()
+
+ dynamic_mask[
+ int(xy[1]): int(xy[1] + height),
+ int(xy[0]): int(xy[0] + width)
+ ] = np.maximum(
+ dynamic_mask[
+ int(xy[1]): int(xy[1] + height),
+ int(xy[0]): int(xy[0] + width)
+ ],
+ 1
+ )
+
+ # Save dynamic mask
+ dynamic_mask = np.clip((dynamic_mask > 0.) * 255, 0, 255).astype(np.uint8)
+ dynamic_mask = Image.fromarray(dynamic_mask, "L")
+ dynamic_mask_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/dynamic_masks/{class_valid}/{str(key_frame_idx).zfill(3)}_{str(cam_idx)}.png"
+ dynamic_mask.save(dynamic_mask_path)
+
+ if curr_sample_record['next'] == '' or curr_sample_record['token'] == last_sample_token:
+ break
+ key_frame_idx += 1
+ curr_sample_record = self.nusc.get('sample', curr_sample_record['next'])
+
+ def save_dynamic_mask_interpolated(self, scene_data, scene_idx, timestamps, instances_info, class_valid='all'):
+ """Parse and save the interpolated segmentation data."""
+ assert class_valid in ['all', 'human', 'vehicle'], "Invalid class valid"
+ if class_valid == 'all':
+ VALID_CLASSES = NUSCENES_DYNAMIC_CLASSES
+ elif class_valid == 'human':
+ VALID_CLASSES = NUSCENES_NONRIGID_DYNAMIC_CLASSES
+ elif class_valid == 'vehicle':
+ VALID_CLASSES = NUSCENES_RIGID_DYNAMIC_CLASSES
+ mask_dir = f"{self.save_dir}/{str(scene_idx).zfill(3)}/dynamic_masks/{class_valid}"
+ if not os.path.exists(mask_dir):
+ os.makedirs(mask_dir)
+
+ for frame_idx, timestamp in enumerate(timestamps):
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ # Find the closest image token for this timestamp
+ closest_token = self.find_closest_img_tokens(scene_data, [timestamp], cam_name)[0]
+ cam_data = self.nusc.get('sample_data', closest_token)
+ img_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/images/{str(frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+
+ img = cv2.imread(img_path)
+ dynamic_mask = np.zeros(img.shape[:2], dtype=np.float32)
+
+ objects = [obj_id for obj_id, obj_info in instances_info.items()
+ if frame_idx in obj_info['frame_annotations']['frame_idx']
+ and obj_info['class_name'] in VALID_CLASSES]
+
+ # Get camera calibration data
+ cs_record = self.nusc.get('calibrated_sensor', cam_data['calibrated_sensor_token'])
+ camera_intrinsic = np.array(cs_record['camera_intrinsic'])
+
+ # Get ego pose
+ pose_record = self.nusc.get('ego_pose', cam_data['ego_pose_token'])
+
+ # Project 3D boxes to 2D and create mask
+ for obj_id in objects:
+ obj_info = instances_info[obj_id]
+ idx_in_obj = obj_info['frame_annotations']['frame_idx'].index(frame_idx)
+ o2w = np.array(obj_info['frame_annotations']['obj_to_world'][idx_in_obj])
+ length, width, height = obj_info['frame_annotations']['box_size'][idx_in_obj]
+
+ # Create Box object
+ box = Box(o2w[:3, 3], [width, length, height], Quaternion(matrix=o2w[:3, :3]))
+
+ # Move box to ego vehicle coordinate system
+ box.translate(-np.array(pose_record['translation']))
+ box.rotate(Quaternion(pose_record['rotation']).inverse)
+
+ # Move box to sensor coordinate system
+ box.translate(-np.array(cs_record['translation']))
+ box.rotate(Quaternion(cs_record['rotation']).inverse)
+
+ # Project 3D box to 2D
+ corners_3d = box.corners()
+ corners_2d = view_points(corners_3d, camera_intrinsic, normalize=True)
+
+ # Check if the object is in front of the camera and all corners are in the image
+ in_front = np.all(corners_3d[2, :] > 0.1)
+ in_image = np.all(corners_2d[0, :] >= 0) & np.all(corners_2d[0, :] < img.shape[1]) & \
+ np.all(corners_2d[1, :] >= 0) & np.all(corners_2d[1, :] < img.shape[0])
+ if not (in_front and in_image):
+ continue
+
+ # If valid, extract x and y coordinates
+ corners_2d = corners_2d[:2, :]
+
+ # Fill the mask
+ u, v = corners_2d[0, :].astype(np.int32), corners_2d[1, :].astype(np.int32)
+ u = np.clip(u, 0, img.shape[1] - 1)
+ v = np.clip(v, 0, img.shape[0] - 1)
+
+ if u.max() - u.min() == 0 or v.max() - v.min() == 0:
+ continue
+
+ xy = (u.min(), v.min())
+ width = u.max() - u.min()
+ height = v.max() - v.min()
+
+ dynamic_mask[
+ int(xy[1]): int(xy[1] + height),
+ int(xy[0]): int(xy[0] + width)
+ ] = np.maximum(
+ dynamic_mask[
+ int(xy[1]): int(xy[1] + height),
+ int(xy[0]): int(xy[0] + width)
+ ],
+ 1
+ )
+
+ # Save dynamic mask
+ dynamic_mask = np.clip((dynamic_mask > 0.) * 255, 0, 255).astype(np.uint8)
+ dynamic_mask = Image.fromarray(dynamic_mask, "L")
+ dynamic_mask_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/dynamic_masks/{class_valid}/{str(frame_idx).zfill(3)}_{str(cam_idx)}.png"
+ dynamic_mask.save(dynamic_mask_path)
+
+ def save_objects(self, scene_data):
+ """Parse and save the objects annotation data."""
+ first_sample_token, last_sample_token = scene_data['first_sample_token'], scene_data['last_sample_token']
+ curr_sample_record = self.nusc.get('sample', first_sample_token)
+ key_frame_idx = 0
+
+ instances_info, frame_instances = {}, {}
+ while True:
+ anns = [self.nusc.get('sample_annotation', token) for token in curr_sample_record['anns']]
+
+ for ann in anns:
+ if ann['category_name'] not in NUSCENES_DYNAMIC_CLASSES:
+ continue
+
+ instance_token = ann['instance_token']
+ if instance_token not in instances_info:
+ instances_info[instance_token] = {
+ 'id': instance_token,
+ 'class_name': ann['category_name'],
+ 'frame_annotations': {
+ 'frame_idx': [],
+ 'obj_to_world': [],
+ 'box_size': [],
+ }
+ }
+
+ # Object to world transformation
+ o2w = np.eye(4)
+ o2w[:3, :3] = Quaternion(ann['rotation']).rotation_matrix
+ o2w[:3, 3] = np.array(ann['translation'])
+
+ # Key frames are spaced (interpolate_N + 1) frames apart in the new sequence
+ obj_frame_idx = key_frame_idx * (self.interpolate_N + 1)
+ instances_info[instance_token]['frame_annotations']['frame_idx'].append(obj_frame_idx)
+ instances_info[instance_token]['frame_annotations']['obj_to_world'].append(o2w.tolist())
+ # convert wlh to lwh
+ lwh = [ann['size'][1], ann['size'][0], ann['size'][2]]
+ instances_info[instance_token]['frame_annotations']['box_size'].append(lwh)
+
+ if key_frame_idx not in frame_instances:
+ frame_instances[key_frame_idx] = []
+ frame_instances[key_frame_idx].extend([ann['instance_token'] for ann in anns if ann['category_name'] in NUSCENES_DYNAMIC_CLASSES])
+
+ if curr_sample_record['next'] == '' or curr_sample_record['token'] == last_sample_token:
+ break
+ key_frame_idx += 1
+ curr_sample_record = self.nusc.get('sample', curr_sample_record['next'])
+
+ # Correct ID mapping
+ id_map = {}
+ for i, (k, v) in enumerate(instances_info.items()):
+ id_map[v["id"]] = i
+
+ # Update keys in instances_info
+ new_instances_info = {}
+ for k, v in instances_info.items():
+ new_instances_info[id_map[v["id"]]] = v
+
+ # Update keys in frame_instances
+ new_frame_instances = {}
+ for k, v in frame_instances.items():
+ new_frame_instances[k] = [id_map[i] for i in v]
+
+ return new_instances_info, new_frame_instances
+
+ def interpolate_boxes(self, instances_info):
+ """Interpolate object positions and sizes between keyframes."""
+ new_instances_info = {}
+ new_frame_instances = {}
+
+ for obj_id, obj_info in instances_info.items():
+ frame_annotations = obj_info['frame_annotations']
+ keyframe_indices = frame_annotations['frame_idx']
+ obj_to_world_list = frame_annotations['obj_to_world']
+ box_size_list = frame_annotations['box_size']
+
+ new_frame_idx = []
+ new_obj_to_world = []
+ new_box_size = []
+
+ for i in range(len(keyframe_indices) - 1):
+ start_frame = keyframe_indices[i]
+ start_transform = np.array(obj_to_world_list[i])
+ end_transform = np.array(obj_to_world_list[i + 1])
+ start_quat = Quaternion(matrix=start_transform[:3, :3])
+ end_quat = Quaternion(matrix=end_transform[:3, :3])
+ start_size = np.array(box_size_list[i])
+ end_size = np.array(box_size_list[i + 1])
+
+ for j in range(self.interpolate_N + 1):
+ t = j / (self.interpolate_N + 1)
+ current_frame = start_frame + j
+
+ # Interpolate translation
+ translation = (1 - t) * start_transform[:3, 3] + t * end_transform[:3, 3]
+
+ # Interpolate rotation using Quaternions
+ current_quat = Quaternion.slerp(start_quat, end_quat, t)
+
+ # Construct interpolated transformation matrix
+ current_transform = np.eye(4)
+ current_transform[:3, :3] = current_quat.rotation_matrix
+ current_transform[:3, 3] = translation
+
+ # Interpolate box size
+ current_size = (1 - t) * start_size + t * end_size
+
+ new_frame_idx.append(current_frame)
+ new_obj_to_world.append(current_transform.tolist())
+ new_box_size.append(current_size.tolist())
+
+ # Add the last keyframe
+ new_frame_idx.append(keyframe_indices[-1])
+ new_obj_to_world.append(obj_to_world_list[-1])
+ new_box_size.append(box_size_list[-1])
+
+ # Update instance info
+ new_instances_info[obj_id] = {
+ 'id': obj_info['id'],
+ 'class_name': obj_info['class_name'],
+ 'frame_annotations': {
+ 'frame_idx': new_frame_idx,
+ 'obj_to_world': new_obj_to_world,
+ 'box_size': new_box_size,
+ }
+ }
+
+ # Update frame instances
+ for frame in new_frame_idx:
+ if frame not in new_frame_instances:
+ new_frame_instances[frame] = []
+ new_frame_instances[frame].append(obj_id)
+
+ return new_instances_info, new_frame_instances
+
+ def visualize_dynamic_objects(self, scene_data, scene_idx, instances_info, frame_instances):
+ """DEBUG: Visualize the dynamic objects' boxes with different colors on the image."""
+ output_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/instances/debug_vis"
+
+ # Get all sample tokens for this scene
+ first_sample_token, last_sample_token = scene_data['first_sample_token'], scene_data['last_sample_token']
+ curr_sample_record = self.nusc.get('sample', first_sample_token)
+ key_frame_idx = 0
+ while True:
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ cam_data = self.nusc.get('sample_data', curr_sample_record['data'][cam_name])
+ img_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/images/{str(key_frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+ canvas = np.array(Image.open(img_path))
+ objects = frame_instances[key_frame_idx]
+
+ if len(objects) > 0:
+ lstProj2d = []
+ color_list = []
+ for obj_id in objects:
+ idx_in_obj = instances_info[obj_id]['frame_annotations']['frame_idx'].index(key_frame_idx)
+ o2w = np.array(
+ instances_info[obj_id]['frame_annotations']['obj_to_world'][idx_in_obj]
+ )
+ length, width, height = instances_info[obj_id]['frame_annotations']['box_size'][idx_in_obj]
+ corners = get_corners(length, width, height)
+
+ # Transform corners to world coordinates
+ corners_world = (o2w[:3, :3] @ corners + o2w[:3, 3:4]).T
+
+ # Get camera calibration data
+ cs_record = self.nusc.get('calibrated_sensor', cam_data['calibrated_sensor_token'])
+ camera_intrinsic = np.array(cs_record['camera_intrinsic'])
+
+ # Get ego pose
+ pose_record = self.nusc.get('ego_pose', cam_data['ego_pose_token'])
+
+ # Transform from world to ego vehicle coordinates
+ corners_ego = corners_world - np.array(pose_record['translation'])
+ corners_ego = Quaternion(pose_record['rotation']).inverse.rotation_matrix @ corners_ego.T
+
+ # Transform from ego to camera coordinates
+ corners_cam = corners_ego - np.array(cs_record['translation']).reshape(3, 1)
+ corners_cam = Quaternion(cs_record['rotation']).inverse.rotation_matrix @ corners_cam
+
+ # Project to 2D
+ corners_2d = view_points(corners_cam, camera_intrinsic, normalize=True)
+
+ # Check if the object is in front of the camera and all corners are in the image
+ in_front = np.all(corners_cam[2, :] > 0.1)
+ in_image = np.all(corners_2d[0, :] >= 0) & np.all(corners_2d[0, :] < canvas.shape[1]) & \
+ np.all(corners_2d[1, :] >= 0) & np.all(corners_2d[1, :] < canvas.shape[0])
+ ok = in_front and in_image
+
+ if ok:
+ projected_points2d = corners_2d[:2, :].T
+ lstProj2d.append(projected_points2d)
+ color_list.append(color_mapper(str(obj_id)))
+
+ lstProj2d = np.asarray(lstProj2d)
+ img_plotted = dump_3d_bbox_on_image(coords=lstProj2d, img=canvas, color=color_list)
+ else:
+ img_plotted = canvas
+
+ img_path = (
+ f"{output_path}/"
+ + f"{str(key_frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+ )
+ Image.fromarray(img_plotted).save(img_path)
+
+ if curr_sample_record['next'] == '' or curr_sample_record['token'] == last_sample_token:
+ break
+ key_frame_idx += 1
+ curr_sample_record = self.nusc.get('sample', curr_sample_record['next'])
+
+ def visualize_dynamic_objects_interpolated(self, scene_data, scene_idx, timestamps, instances_info, frame_instances):
+ """DEBUG: Visualize the interpolated dynamic objects' boxes with different colors on the image."""
+ output_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/instances/debug_vis"
+
+ for frame_idx, timestamp in enumerate(timestamps):
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ # Find the closest image token for this timestamp
+ closest_token = self.find_closest_img_tokens(scene_data, [timestamp], cam_name)[0]
+ cam_data = self.nusc.get('sample_data', closest_token)
+
+ img_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/images/{str(frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+ canvas = np.array(Image.open(img_path))
+ objects = frame_instances.get(frame_idx, [])
+
+ if len(objects) > 0:
+ lstProj2d = []
+ color_list = []
+ for obj_id in objects:
+ obj_info = instances_info[obj_id]
+ idx_in_obj = obj_info['frame_annotations']['frame_idx'].index(frame_idx)
+ o2w = np.array(obj_info['frame_annotations']['obj_to_world'][idx_in_obj])
+ length, width, height = obj_info['frame_annotations']['box_size'][idx_in_obj]
+ corners = get_corners(length, width, height)
+
+ # Transform corners to world coordinates
+ corners_world = (o2w[:3, :3] @ corners + o2w[:3, 3:4]).T
+
+ # Get camera calibration data
+ cs_record = self.nusc.get('calibrated_sensor', cam_data['calibrated_sensor_token'])
+ camera_intrinsic = np.array(cs_record['camera_intrinsic'])
+
+ # Get ego pose
+ pose_record = self.nusc.get('ego_pose', cam_data['ego_pose_token'])
+
+ # Transform from world to ego vehicle coordinates
+ corners_ego = corners_world - np.array(pose_record['translation'])
+ corners_ego = Quaternion(pose_record['rotation']).inverse.rotation_matrix @ corners_ego.T
+
+ # Transform from ego to camera coordinates
+ corners_cam = corners_ego - np.array(cs_record['translation']).reshape(3, 1)
+ corners_cam = Quaternion(cs_record['rotation']).inverse.rotation_matrix @ corners_cam
+
+ # Project to 2D
+ corners_2d = view_points(corners_cam, camera_intrinsic, normalize=True)
+
+ # Check if the object is in front of the camera and all corners are in the image
+ in_front = np.all(corners_cam[2, :] > 0.1)
+ in_image = np.all(corners_2d[0, :] >= 0) & np.all(corners_2d[0, :] < canvas.shape[1]) & \
+ np.all(corners_2d[1, :] >= 0) & np.all(corners_2d[1, :] < canvas.shape[0])
+ ok = in_front and in_image
+
+ if ok:
+ projected_points2d = corners_2d[:2, :].T
+ lstProj2d.append(projected_points2d)
+ color_list.append(color_mapper(str(obj_id)))
+
+ lstProj2d = np.asarray(lstProj2d)
+ img_plotted = dump_3d_bbox_on_image(coords=lstProj2d, img=canvas, color=color_list)
+ else:
+ img_plotted = canvas
+
+ img_path = f"{output_path}/{str(frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+ Image.fromarray(img_plotted).save(img_path)
+
+ def create_folder(self):
+ """Create folder for data preprocessing."""
+ if self.process_id_list is None:
+ id_list = range(len(self))
+ else:
+ id_list = self.process_id_list
+ for i in id_list:
+ if "images" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/images", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/sky_masks", exist_ok=True)
+ if "calib" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/extrinsics", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/intrinsics", exist_ok=True)
+ if "lidar" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/lidar", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/lidar_pose", exist_ok=True)
+ if "dynamic_masks" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/dynamic_masks/all", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/dynamic_masks/human", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/dynamic_masks/vehicle", exist_ok=True)
+ if "objects" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/instances", exist_ok=True)
+ if "objects_vis" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/instances/debug_vis", exist_ok=True)
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/nuscenes/nuscenes_sourceloader.py b/worldbench/third_party/drivestudio/datasets/nuscenes/nuscenes_sourceloader.py
new file mode 100644
index 0000000..7be6721
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/nuscenes/nuscenes_sourceloader.py
@@ -0,0 +1,460 @@
+import json
+import logging
+import os
+from typing import Dict
+
+import joblib
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from pytorch3d.transforms import matrix_to_quaternion
+from torch import Tensor
+from tqdm import tqdm, trange
+
+from datasets.base.scene_dataset import ModelType
+from datasets.base.lidar_source import SceneLidarSource
+from datasets.base.pixel_source import CameraData, ScenePixelSource
+
+logger = logging.getLogger()
+
+# define each class's node type
+OBJECT_CLASS_NODE_MAPPING = {
+ # Rigid objects (vehicles)
+ "vehicle.bus.bendy": ModelType.RigidNodes,
+ "vehicle.bus.rigid": ModelType.RigidNodes,
+ "vehicle.car": ModelType.RigidNodes,
+ "vehicle.construction": ModelType.RigidNodes,
+ "vehicle.emergency.ambulance": ModelType.RigidNodes,
+ "vehicle.emergency.police": ModelType.RigidNodes,
+ "vehicle.motorcycle": ModelType.RigidNodes,
+ "vehicle.trailer": ModelType.RigidNodes,
+ "vehicle.truck": ModelType.RigidNodes,
+
+ # Humans (SMPL model)
+ "human.pedestrian.adult": ModelType.SMPLNodes,
+ "human.pedestrian.child": ModelType.SMPLNodes,
+ "human.pedestrian.construction_worker": ModelType.SMPLNodes,
+ "human.pedestrian.police_officer": ModelType.SMPLNodes,
+
+ # Potentially deformable objects
+ "human.pedestrian.personal_mobility": ModelType.DeformableNodes,
+ "human.pedestrian.stroller": ModelType.DeformableNodes,
+ "human.pedestrian.wheelchair": ModelType.DeformableNodes,
+ "animal": ModelType.DeformableNodes,
+ "vehicle.bicycle": ModelType.DeformableNodes
+}
+SMPLNODE_CLASSES = [
+ "human.pedestrian.adult",
+ "human.pedestrian.child",
+ "human.pedestrian.construction_worker",
+ "human.pedestrian.police_officer"
+]
+
+# OpenCV to Dataset coordinate transformation
+# opencv coordinate system: x right, y down, z front
+OPENCV2DATASET = np.array(
+ [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]
+)
+
+# NuScenes Camera List:
+# 0: CAM_FRONT
+# 1: CAM_FRONT_LEFT
+# 2: CAM_FRONT_RIGHT
+# 3: CAM_BACK_LEFT
+# 4: CAM_BACK_RIGHT
+# 5: CAM_BACK
+AVAILABLE_CAM_LIST = [0, 1, 2, 3, 4, 5]
+
+class NuScenesCameraData(CameraData):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ def load_calibrations(self):
+ cam_to_worlds, intrinsics, distortions = [], [], []
+
+ # Load the first camera (front) pose to align the world
+ camera_front_start = np.loadtxt(
+ os.path.join(self.data_path, "extrinsics", f"{self.start_timestep:03d}_0.txt")
+ )
+
+ for t in range(self.start_timestep, self.end_timestep):
+ # Load intrinsics
+ intrinsic = np.loadtxt(
+ os.path.join(self.data_path, "intrinsics", f"{self.cam_id}.txt")
+ )
+ fx, fy, cx, cy = intrinsic[0], intrinsic[1], intrinsic[2], intrinsic[3]
+ k1, k2, p1, p2, k3 = intrinsic[4], intrinsic[5], intrinsic[6], intrinsic[7], intrinsic[8]
+
+ # Scale intrinsics
+ fx, fy = (
+ fx * self.load_size[1] / self.original_size[1],
+ fy * self.load_size[0] / self.original_size[0],
+ )
+ cx, cy = (
+ cx * self.load_size[1] / self.original_size[1],
+ cy * self.load_size[0] / self.original_size[0],
+ )
+ intrinsic = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
+ intrinsics.append(intrinsic)
+ distortions.append(np.array([k1, k2, p1, p2, k3]))
+
+ # Load extrinsics
+ cam2world = np.loadtxt(
+ os.path.join(self.data_path, "extrinsics", f"{t:03d}_{self.cam_id}.txt")
+ )
+
+ # Align camera poses with the first camera pose
+ cam2world = np.linalg.inv(camera_front_start) @ cam2world
+
+ # Convert to OpenCV coordinate system
+ cam2world = cam2world @ OPENCV2DATASET
+ cam_to_worlds.append(cam2world)
+
+ self.intrinsics = torch.from_numpy(np.stack(intrinsics, axis=0)).float()
+ self.distortions = torch.from_numpy(np.stack(distortions, axis=0)).float()
+ self.cam_to_worlds = torch.from_numpy(np.stack(cam_to_worlds, axis=0)).float()
+
+ @classmethod
+ def get_camera2worlds(cls, data_path: str, cam_id: str, start_timestep: int, end_timestep: int) -> torch.Tensor:
+ """
+ Returns camera-to-world matrices for the specified camera and time range.
+
+ Args:
+ data_path (str): Path to the dataset.
+ cam_id (str): Camera ID.
+ start_timestep (int): Start timestep.
+ end_timestep (int): End timestep.
+
+ Returns:
+ torch.Tensor: Camera-to-world matrices of shape (num_frames, 4, 4).
+ """
+ # Load the first camera (front) pose to align the world
+ camera_front_start = np.loadtxt(
+ os.path.join(data_path, "extrinsics", f"{start_timestep:03d}_0.txt")
+ )
+
+ cam_to_worlds = []
+
+ for t in range(start_timestep, end_timestep):
+ # Load extrinsics for each timestep
+ cam2world = np.loadtxt(
+ os.path.join(data_path, "extrinsics", f"{t:03d}_{cam_id}.txt")
+ )
+
+ # Align camera poses with the first camera pose
+ cam2world = np.linalg.inv(camera_front_start) @ cam2world
+
+ # Convert to OpenCV coordinate system
+ cam2world = cam2world @ OPENCV2DATASET
+
+ cam_to_worlds.append(cam2world)
+
+ return torch.from_numpy(np.stack(cam_to_worlds, axis=0)).float()
+
+class NuScenesPixelSource(ScenePixelSource):
+ def __init__(
+ self,
+ dataset_name: str,
+ pixel_data_config: OmegaConf,
+ data_path: str,
+ start_timestep: int,
+ end_timestep: int,
+ device: torch.device = torch.device("cpu"),
+ ):
+ super().__init__(dataset_name, pixel_data_config, device=device)
+ self.data_path = data_path
+ self.start_timestep = start_timestep
+ self.end_timestep = end_timestep
+ self.load_data()
+
+ def load_cameras(self):
+ self._timesteps = torch.arange(self.start_timestep, self.end_timestep)
+ self.register_normalized_timestamps()
+
+ for idx, cam_id in enumerate(self.camera_list):
+ logger.info(f"Loading camera {cam_id}")
+ camera = NuScenesCameraData(
+ dataset_name=self.dataset_name,
+ data_path=self.data_path,
+ cam_id=cam_id,
+ start_timestep=self.start_timestep,
+ end_timestep=self.end_timestep,
+ load_dynamic_mask=self.data_cfg.load_dynamic_mask,
+ load_sky_mask=self.data_cfg.load_sky_mask,
+ downscale_when_loading=self.data_cfg.downscale_when_loading[idx],
+ undistort=self.data_cfg.undistort,
+ buffer_downscale=self.buffer_downscale,
+ device=self.device,
+ )
+ camera.load_time(self.normalized_time)
+ unique_img_idx = torch.arange(len(camera), device=self.device) * len(self.camera_list) + idx
+ camera.set_unique_ids(
+ unique_cam_idx = idx,
+ unique_img_idx = unique_img_idx
+ )
+ logger.info(f"Camera {camera.cam_name} loaded.")
+ self.camera_data[cam_id] = camera
+
+ def load_objects(self):
+ """
+ get ground truth bounding boxes of the dynamic objects
+
+ instances_info = {
+ "0": # simplified instance id
+ {
+ "id": str,
+ "class_name": str,
+ "frame_annotations": {
+ "frame_idx": List,
+ "obj_to_world": List,
+ "box_size": List,
+ },
+ ...
+ }
+ frame_instances = {
+ "0": # frame idx
+ List[int] # list of simplified instance ids
+ ...
+ }
+ """
+ instances_info_path = os.path.join(self.data_path, "instances", "instances_info.json")
+ frame_instances_path = os.path.join(self.data_path, "instances", "frame_instances.json")
+ with open(instances_info_path, "r") as f:
+ instances_info = json.load(f)
+ with open(frame_instances_path, "r") as f:
+ frame_instances = json.load(f)
+ # get pose of each instance at each frame
+ # shape (num_frames, num_instances, 4, 4)
+ num_instances = len(instances_info)
+ num_full_frames = len(frame_instances)
+ instances_pose = np.zeros((num_full_frames, num_instances, 4, 4))
+ instances_size = np.zeros((num_full_frames, num_instances, 3))
+ instances_true_id = np.arange(num_instances)
+ instances_model_types = np.ones(num_instances) * -1
+
+ # Load the first camera (front) pose to align the world
+ camera_front_start = np.loadtxt(
+ os.path.join(self.data_path, "extrinsics", f"{self.start_timestep:03d}_0.txt")
+ )
+ for k, v in instances_info.items():
+ instances_model_types[int(k)] = OBJECT_CLASS_NODE_MAPPING[v["class_name"]]
+ for frame_idx, obj_to_world, box_size in zip(v["frame_annotations"]["frame_idx"], v["frame_annotations"]["obj_to_world"], v["frame_annotations"]["box_size"]):
+ # the first ego pose as the origin of the world coordinate system.
+ obj_to_world = np.array(obj_to_world).reshape(4, 4)
+ obj_to_world = np.linalg.inv(camera_front_start) @ obj_to_world
+ instances_pose[frame_idx, int(k)] = np.array(obj_to_world)
+ instances_size[frame_idx, int(k)] = np.array(box_size)
+
+ # get frame valid instances
+ # shape (num_frames, num_instances)
+ per_frame_instance_mask = np.zeros((num_full_frames, num_instances))
+ for frame_idx, valid_instances in frame_instances.items():
+ per_frame_instance_mask[int(frame_idx), valid_instances] = 1
+
+ # select the frames that are in the range of start_timestep and end_timestep
+ instances_pose = torch.from_numpy(instances_pose[self.start_timestep:self.end_timestep]).float()
+ instances_size = torch.from_numpy(instances_size[self.start_timestep:self.end_timestep]).float()
+ instances_true_id = torch.from_numpy(instances_true_id).long()
+ instances_model_types = torch.from_numpy(instances_model_types).long()
+ per_frame_instance_mask = torch.from_numpy(per_frame_instance_mask[self.start_timestep:self.end_timestep]).bool()
+
+ # filter out the instances that are not visible in selected frames
+ ins_frame_cnt = per_frame_instance_mask.sum(dim=0)
+ instances_pose = instances_pose[:, ins_frame_cnt > 0]
+ instances_size = instances_size[:, ins_frame_cnt > 0]
+ instances_true_id = instances_true_id[ins_frame_cnt > 0]
+ instances_model_types = instances_model_types[ins_frame_cnt > 0]
+ per_frame_instance_mask = per_frame_instance_mask[:, ins_frame_cnt > 0]
+
+ # assign to the class
+ # (num_frames, num_instances, 4, 4)
+ self.instances_pose = instances_pose
+ # (num_instances, 3)
+ self.instances_size = instances_size.sum(0) / per_frame_instance_mask.sum(0).unsqueeze(-1)
+ # (num_frames, num_instances)
+ self.per_frame_instance_mask = per_frame_instance_mask
+ # (num_instances)
+ self.instances_true_id = instances_true_id
+ # (num_instances)
+ self.instances_model_types = instances_model_types
+
+ if self.data_cfg.load_smpl:
+ # Collect camera-to-world matrices for all available cameras
+ cam_to_worlds = {}
+ for cam_id in AVAILABLE_CAM_LIST:
+ cam_to_worlds[cam_id] = NuScenesCameraData.get_camera2worlds(
+ self.data_path,
+ str(cam_id),
+ self.start_timestep,
+ self.end_timestep
+ )
+
+ # load SMPL parameters
+ smpl_dict = joblib.load(os.path.join(self.data_path, "humanpose", "smpl.pkl"))
+ frame_num = self.end_timestep - self.start_timestep
+
+ smpl_human_all = {}
+ for fi in tqdm(range(self.start_timestep, self.end_timestep), desc="Loading SMPL"):
+ for instance_id, ins_smpl in smpl_dict.items():
+ if instance_id not in smpl_human_all:
+ smpl_human_all[instance_id] = {
+ "smpl_quats": torch.zeros((frame_num, 24, 4), dtype=torch.float32),
+ "smpl_trans": torch.zeros((frame_num, 3), dtype=torch.float32),
+ "smpl_betas": torch.zeros((frame_num, 10), dtype=torch.float32),
+ "frame_valid": torch.zeros((frame_num), dtype=torch.bool)
+ }
+ smpl_human_all[instance_id]["smpl_quats"][:, :, 0] = 1.0
+ if ins_smpl["valid_mask"][fi]:
+ betas = ins_smpl["smpl"]["betas"][fi]
+ smpl_human_all[instance_id]["smpl_betas"][fi - self.start_timestep] = betas
+
+ body_pose = ins_smpl["smpl"]["body_pose"][fi]
+ smpl_orient = ins_smpl["smpl"]["global_orient"][fi]
+ cam_depend = ins_smpl["selected_cam_idx"][fi].item()
+
+ c2w = cam_to_worlds[cam_depend][fi - self.start_timestep]
+ world_orient = c2w[:3, :3].to(smpl_orient.device) @ smpl_orient.squeeze()
+ smpl_quats = matrix_to_quaternion(
+ torch.cat([world_orient[None, ...], body_pose], dim=0)
+ )
+
+ ii = instances_info[str(instance_id)]['frame_annotations']["frame_idx"].index(fi)
+ o2w = np.array(
+ instances_info[str(instance_id)]['frame_annotations']["obj_to_world"][ii]
+ )
+ o2w = torch.from_numpy(
+ np.linalg.inv(camera_front_start) @ o2w
+ )
+ # box_size = instances_info[str(instance_id)]['frame_annotations']["box_size"][ii]
+
+ smpl_human_all[instance_id]["smpl_quats"][fi - self.start_timestep] = smpl_quats
+ smpl_human_all[instance_id]["smpl_trans"][fi - self.start_timestep] = o2w[:3, 3]
+ smpl_human_all[instance_id]["frame_valid"][fi - self.start_timestep] = True
+
+ self.smpl_human_all = smpl_human_all
+
+class NuScenesLiDARSource(SceneLidarSource):
+ def __init__(
+ self,
+ lidar_data_config: OmegaConf,
+ data_path: str,
+ start_timestep: int,
+ end_timestep: int,
+ device: torch.device = torch.device("cpu"),
+ ):
+ super().__init__(lidar_data_config, device=device)
+ self.data_path = data_path
+ self.start_timestep = start_timestep
+ self.end_timestep = end_timestep
+ self.create_all_filelist()
+ self.load_data()
+
+ def create_all_filelist(self):
+ lidar_filepaths = []
+ lidar_pose_filepaths = []
+ for t in range(self.start_timestep, self.end_timestep):
+ lidar_filepaths.append(
+ os.path.join(self.data_path, "lidar", f"{t:03d}.bin")
+ )
+ lidar_pose_filepaths.append(
+ os.path.join(self.data_path, "lidar_pose", f"{t:03d}.txt")
+ )
+ self.lidar_filepaths = np.array(lidar_filepaths)
+ self.lidar_pose_filepaths = np.array(lidar_pose_filepaths)
+
+ def load_calibrations(self):
+ lidar_to_worlds = []
+
+ # Load the first camera (front) pose to align the world
+ camera_front_start = np.loadtxt(
+ os.path.join(self.data_path, "extrinsics", f"{self.start_timestep:03d}_0.txt")
+ )
+
+ for pose_file in self.lidar_pose_filepaths:
+ lidar_to_world = np.loadtxt(pose_file)
+ # Align lidar poses with the first camera pose
+ lidar_to_world = np.linalg.inv(camera_front_start) @ lidar_to_world
+ lidar_to_worlds.append(lidar_to_world)
+
+ self.lidar_to_worlds = torch.from_numpy(
+ np.stack(lidar_to_worlds, axis=0)
+ ).float()
+
+ def load_lidar(self):
+ origins, directions, ranges = [], [], []
+ timesteps = []
+
+ accumulated_num_original_rays = 0
+ accumulated_num_rays = 0
+ for t in trange(
+ 0, len(self.lidar_filepaths), desc="Loading lidar", dynamic_ncols=True
+ ):
+ lidar_info = np.fromfile(self.lidar_filepaths[t], dtype=np.float32).reshape(-1, 4)
+ original_length = len(lidar_info)
+ accumulated_num_original_rays += original_length
+
+ lidar_points = torch.from_numpy(lidar_info[:, :3]).float()
+ lidar_origins = torch.zeros_like(lidar_points)
+
+ lidar_origins = (
+ self.lidar_to_worlds[t][:3, :3] @ lidar_origins.T
+ + self.lidar_to_worlds[t][:3, 3:4]
+ ).T
+ lidar_points = (
+ self.lidar_to_worlds[t][:3, :3] @ lidar_points.T
+ + self.lidar_to_worlds[t][:3, 3:4]
+ ).T
+ lidar_directions = lidar_points - lidar_origins
+ lidar_ranges = torch.norm(lidar_directions, dim=-1, keepdim=True)
+ lidar_directions = lidar_directions / lidar_ranges
+ lidar_timestamp = torch.ones_like(lidar_ranges).squeeze(-1) * t
+ accumulated_num_rays += len(lidar_ranges)
+
+ origins.append(lidar_origins)
+ directions.append(lidar_directions)
+ ranges.append(lidar_ranges)
+ timesteps.append(lidar_timestamp)
+
+ logger.info(
+ f"Number of lidar rays: {accumulated_num_rays} "
+ f"({accumulated_num_rays / accumulated_num_original_rays * 100:.2f}% of "
+ f"{accumulated_num_original_rays} original rays)"
+ )
+
+ self.origins = torch.cat(origins, dim=0)
+ self.directions = torch.cat(directions, dim=0)
+ self.ranges = torch.cat(ranges, dim=0)
+ self.visible_masks = torch.zeros_like(self.ranges).squeeze().bool()
+ self.colors = torch.ones_like(self.directions)
+
+ self._timesteps = torch.cat(timesteps, dim=0)
+ self.register_normalized_timestamps()
+
+ def get_lidar_rays(self, time_idx: int) -> Dict[str, Tensor]:
+ origins = self.origins[self.timesteps == time_idx]
+ directions = self.directions[self.timesteps == time_idx]
+ ranges = self.ranges[self.timesteps == time_idx]
+ normalized_time = self.normalized_time[self.timesteps == time_idx]
+ return {
+ "lidar_origins": origins,
+ "lidar_viewdirs": directions,
+ "lidar_ranges": ranges,
+ "lidar_normed_time": normalized_time,
+ "lidar_mask": self.timesteps == time_idx,
+ }
+
+ def delete_invisible_pts(self) -> None:
+ if self.visible_masks is not None:
+ num_bf = self.origins.shape[0]
+ self.origins = self.origins[self.visible_masks]
+ self.directions = self.directions[self.visible_masks]
+ self.ranges = self.ranges[self.visible_masks]
+ self._timesteps = self._timesteps[self.visible_masks]
+ self._normalized_time = self._normalized_time[self.visible_masks]
+ self.colors = self.colors[self.visible_masks]
+ logger.info(
+ f"[Lidar] {num_bf - self.visible_masks.sum()} out of {num_bf} points are cleared. {self.visible_masks.sum()} points left."
+ )
+ self.visible_masks = None
+ else:
+ logger.info("[Lidar] No unvisible points to clear.")
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/nuscenes_gen/nuscenes_human_utils.py b/worldbench/third_party/drivestudio/datasets/nuscenes_gen/nuscenes_human_utils.py
new file mode 100644
index 0000000..31a9085
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/nuscenes_gen/nuscenes_human_utils.py
@@ -0,0 +1,216 @@
+from typing import List
+import os
+import cv2
+import json
+import logging
+import numpy as np
+from tqdm import tqdm
+
+from utils.geometry import (
+ get_corners,
+ project_camera_points_to_image
+)
+from .nuscenes_sourceloader import (
+ SMPLNODE_CLASSES,
+ OPENCV2DATASET,
+ AVAILABLE_CAM_LIST,
+)
+
+logger = logging.getLogger()
+
+CAMERA_LIST = AVAILABLE_CAM_LIST
+
+def project_human_boxes(
+ scene_dir: str,
+ camera_list: List[int],
+ save_temp=True,
+ verbose=False,
+ narrow_width_ratio=0.2,
+ fps=12
+):
+ """Project human boxes to 2D image space and save the results to pkl files
+
+ Args:
+ scene_dir: str, path to the scene directory
+ camera_list: List[int], a list of camera ids to be processed
+ save_temp: bool, whether to save the intermediate results
+ verbose: bool, whether to visualize the projected boxes
+ narrow_width_ratio: sometimes the projected boxes are too wide
+ we can narrow them with this ratio to get more accurate results
+ fps: int, FPS for the visualization video
+
+ Returns:
+ collector_all: dict, a dictionary containing the projected boxes for each camera
+ """
+ # check if the necessary directories exist
+ images_dir = f'{scene_dir}/images'
+ extrinsics_dir = f'{scene_dir}/extrinsics'
+ intrinsics_dir = f'{scene_dir}/intrinsics'
+ instances_dir = f'{scene_dir}/instances'
+ valid_paths = [images_dir, extrinsics_dir, intrinsics_dir, instances_dir]
+ for path in valid_paths:
+ assert os.path.exists(path), \
+ f"Path {path} does not exist, you need to run waymo preprocess to generate the necessary files"
+
+ # create directories for saving the results
+ save_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks'
+ if not os.path.exists(save_dir):
+ os.makedirs(save_dir)
+
+ if verbose:
+ # create directories for saving the intermediate visualization results
+ video_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks/vis'
+ if not os.path.exists(video_dir):
+ os.makedirs(video_dir)
+ per_human_img_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks/vis/images'
+ if not os.path.exists(per_human_img_dir):
+ os.makedirs(per_human_img_dir)
+
+ # load instances and frame infos
+ frame_infos = json.load(
+ open(f'{instances_dir}/frame_instances.json')
+ )
+ instances_meta = json.load(
+ open(f'{instances_dir}/instances_info.json')
+ )
+
+ collector_all = {}
+ # iterate over each camera
+ for cam_id in camera_list:
+ # check if already processed
+ pkl_path = os.path.join(save_dir, f"{cam_id}.pkl")
+ if os.path.exists(pkl_path):
+ collector_all[cam_id] = json.load(open(pkl_path))
+ logger.info(f"Results for camera {cam_id} already exists at {pkl_path}")
+ continue
+
+ if verbose:
+ per_cam_vis_dir = os.path.join(per_human_img_dir, f"{cam_id}")
+ if not os.path.exists(per_cam_vis_dir):
+ os.makedirs(per_cam_vis_dir)
+
+ collector = {}
+ frames = []
+ for frame_id, frame_ins_list in frame_infos.items():
+ frame_id = int(frame_id)
+
+ # define empty instance collector for each frame
+ frame_collector = {
+ "gt_bbox": [],
+ "extra_data": {
+ "gt_track_id": [],
+ "gt_class": [],
+ }
+ }
+
+ # Load extrinsics for this frame and camera
+ cam2world = np.loadtxt(os.path.join(extrinsics_dir, f"{frame_id:03d}_{cam_id}.txt"))
+ # Convert to OpenCV coordinate system
+ cam2world = cam2world @ OPENCV2DATASET
+
+ # load intrinsic
+ Ks = np.loadtxt(os.path.join(intrinsics_dir, f"{cam_id}.txt"))
+ fx, fy, cx, cy = Ks[0], Ks[1], Ks[2], Ks[3]
+ # k1, k2, p1, p2, k3 = Ks[4], Ks[5], Ks[6], Ks[7], Ks[8]
+ intrinsic = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float32)
+
+ # load image
+ ori_image = cv2.imread(
+ os.path.join(images_dir, f"{str(frame_id).zfill(3)}_{cam_id}.jpg")
+ )
+ image = ori_image.copy()
+ H, W = image.shape[:2]
+
+ # iterate over each instance in the frame
+ if len(frame_ins_list) > 0:
+ # if there are pedestrians in this frame, project the boxes to the image
+ image_plotted = image.copy()
+ for instance_id in frame_ins_list:
+ ins = instances_meta[str(instance_id)]
+
+ if ins["class_name"] not in SMPLNODE_CLASSES:
+ continue
+
+ ins_anno = ins["frame_annotations"]
+ index = ins_anno['frame_idx'].index(frame_id)
+ obj_to_world = np.array(ins_anno['obj_to_world'][index])
+ l, w, h = ins_anno['box_size'][index]
+
+ # get box corners in object space
+ corners = get_corners(l, w, h)
+ # transform box corners to world space
+ corners_world = obj_to_world[:3, :3] @ corners + obj_to_world[:3, 3:4]
+ # transform box corners to image space
+ world2cam = np.linalg.inv(cam2world)
+ corners_cam = world2cam[:3, :3] @ corners_world + world2cam[:3, 3:4]
+ cam_points, depth = project_camera_points_to_image(corners_cam.T, intrinsic)
+
+ x_min, y_min = np.min(cam_points, axis=0)
+ x_max, y_max = np.max(cam_points, axis=0)
+ # clip left and right with this ratio
+ if narrow_width_ratio > 0.:
+ length = x_max - x_min
+ x_min += length * narrow_width_ratio
+ x_max -= length * narrow_width_ratio
+
+ # clip the box to the image
+ original_area = (x_max - x_min) * (y_max - y_min)
+ x_min, x_max = np.clip(x_min, 0, W), np.clip(x_max, 0, W)
+ y_min, y_max = np.clip(y_min, 0, H), np.clip(y_max, 0, H)
+ new_area = (x_max - x_min) * (y_max - y_min)
+
+ # filter out boxes that are too small or too large
+ behind = depth.max() < 0
+ too_small = new_area < W * H * (0.03)**2
+ too_large = new_area > W * H / 1.1
+ too_far = np.linalg.norm(obj_to_world[:3, 3] - cam2world[:3, 3]) > 40
+ clip_large = new_area / original_area < 1/3
+ if too_small or too_large or clip_large or behind or too_far:
+ continue
+
+ # gt box on image
+ gt_box = [x_min, y_min, x_max - x_min, y_max - y_min]
+
+ # save the projected box to the collector
+ frame_collector["gt_bbox"].append(gt_box)
+ frame_collector["extra_data"]["gt_track_id"].append(instance_id)
+ frame_collector["extra_data"]["gt_class"].append([0])
+
+ if verbose:
+ # visualize the projected boxes of ONE instance
+ raw_image = cv2.rectangle(
+ ori_image.copy(), (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 2
+ )
+ raw_image_path = os.path.join(per_cam_vis_dir, f"{frame_id}_{instance_id}.jpg")
+ cv2.imwrite(raw_image_path, raw_image)
+
+ # add this instance to the image
+ image_plotted = cv2.rectangle(
+ image, (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 2
+ )
+
+ if verbose:
+ frames.append(image_plotted)
+ else:
+ # if no instance in this frame, just save the original image
+ if verbose:
+ frames.append(ori_image)
+
+ collector[frame_id] = frame_collector
+
+ if verbose:
+ height, width = frames[0].shape[:2]
+ output_path = os.path.join(video_dir, f"cam_{cam_id}.mp4")
+ fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+ out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+ for frame in tqdm(frames, desc=f"Writing video for camera {cam_id}"):
+ out.write(frame)
+ out.release()
+
+ if save_temp:
+ # save collector to pkl
+ json.dump(collector, open(pkl_path, "w"))
+
+ collector_all[cam_id] = collector
+
+ return collector_all
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/nuscenes_gen/nuscenes_preprocess.py b/worldbench/third_party/drivestudio/datasets/nuscenes_gen/nuscenes_preprocess.py
new file mode 100644
index 0000000..441ce8b
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/nuscenes_gen/nuscenes_preprocess.py
@@ -0,0 +1,1147 @@
+import json
+import os
+from collections import Counter
+from typing import List
+
+import cv2
+import numpy as np
+from PIL import Image
+from pyquaternion import Quaternion
+from tqdm import tqdm
+
+from nuscenes.nuscenes import LidarPointCloud, NuScenes
+from nuscenes.utils.data_classes import Box
+from nuscenes.utils.geometry_utils import view_points
+from nuscenes.utils import splits
+
+from datasets.tools.multiprocess_utils import track_parallel_progress
+from utils.geometry import get_corners
+from utils.visualization import color_mapper, dump_3d_bbox_on_image
+
+NUSCENES_LABELS = [
+ 'animal',
+ 'human.pedestrian.adult',
+ 'human.pedestrian.child',
+ 'human.pedestrian.construction_worker',
+ 'human.pedestrian.personal_mobility',
+ 'human.pedestrian.police_officer',
+ 'human.pedestrian.stroller',
+ 'human.pedestrian.wheelchair',
+ 'movable_object.barrier',
+ 'movable_object.debris',
+ 'movable_object.pushable_pullable',
+ 'movable_object.trafficcone',
+ 'static_object.bicycle_rack',
+ 'vehicle.bicycle',
+ 'vehicle.bus.bendy',
+ 'vehicle.bus.rigid',
+ 'vehicle.car',
+ 'vehicle.construction',
+ 'vehicle.emergency.ambulance',
+ 'vehicle.emergency.police',
+ 'vehicle.motorcycle',
+ 'vehicle.trailer',
+ 'vehicle.truck'
+]
+
+NUSCENES_NONRIGID_DYNAMIC_CLASSES = [
+ 'animal',
+ 'human.pedestrian.adult',
+ 'human.pedestrian.child',
+ 'human.pedestrian.construction_worker',
+ 'human.pedestrian.personal_mobility',
+ 'human.pedestrian.police_officer',
+ 'human.pedestrian.stroller',
+ 'human.pedestrian.wheelchair',
+ 'vehicle.bicycle',
+ 'vehicle.motorcycle',
+]
+
+NUSCENES_RIGID_DYNAMIC_CLASSES = [
+ 'vehicle.bus.bendy',
+ 'vehicle.bus.rigid',
+ 'vehicle.car',
+ 'vehicle.construction',
+ 'vehicle.emergency.ambulance',
+ 'vehicle.emergency.police',
+ 'vehicle.trailer',
+ 'vehicle.truck'
+]
+
+NUSCENES_DYNAMIC_CLASSES = NUSCENES_NONRIGID_DYNAMIC_CLASSES + NUSCENES_RIGID_DYNAMIC_CLASSES
+
+def read_mp4_to_ndarray_list(video_path: str, target_size=None):
+ cap = cv2.VideoCapture(video_path)
+ frames = []
+
+ if not cap.isOpened():
+ raise IOError(f"无法打开视频文件: {video_path}")
+
+ while True:
+ ret, frame = cap.read()
+ if not ret: # 没有帧了
+ break
+ if target_size is not None:
+ # reshape frame
+ frame = cv2.resize(frame, target_size, interpolation=cv2.INTER_LINEAR)
+ frames.append(frame) # frame 默认是 np.ndarray，形状 (H, W, 3)
+
+ cap.release()
+ return frames
+
+class NuScenesProcessor(object):
+ """Process NuScenes dataset.
+
+ NuScenes Data Frequencies:
+ - CAMERA: 12 Hz
+ - LIDAR: 20 Hz
+ - Annotated Keyframes: 2 Hz
+
+ This function processes the NuScenes dataset, offering two processing modes:
+
+ 1. Original Mode (2 Hz):
+ - Uses only the annotated keyframes.
+ - Sparse but contains original annotations.
+ - Default mode when 'interpolate_N' is set to 0.
+
+ 2. Interpolated Mode:
+ - Interpolate N frames between annotated keyframes, where N <= 4 is recommended.
+ - Provides denser data by interpolating between annotated frames.
+ - Activated by setting 'interpolate_N' to a positive integer.
+ - Note: Using N > 4 may lead to insufficient image frames due to potential frame drops,
+ despite the camera's 12 Hz capture rate.
+
+ Interpolated mode may enhance training due to higher temporal resolution,
+ but uses estimated annotations for non-keyframes. Choose the mode that
+ best fits your use case.
+
+ Args:
+ load_dir (str): Source directory of NuScenes data.
+ save_dir (str): Target directory for processed data.
+ split (str): Dataset split to process. Default: 'v1.0-mini'.
+ interpolate_N (int): Number of frames to interpolate between keyframes.
+ If 0, use original 2 Hz data.
+ workers (int): Number of parallel processing workers.
+ process_keys (list): Data types to process.
+ Default: ["images", "lidar", "calib", "dynamic_masks", "objects"]
+ process_id_list (list): Specific scene IDs to process.
+ """
+
+ def __init__(
+ self,
+ load_dir,
+ save_dir,
+ split='v1.0-mini',
+ interpolate_N=0,
+ process_keys=[
+ "images",
+ "lidar",
+ "calib",
+ "dynamic_masks",
+ "objects"
+ ],
+ process_id_list=None,
+ workers=64,
+ num_scenes=150,
+ data_source="gt"
+ ):
+ self.data_source = data_source
+ self.process_id_list = process_id_list
+ self.process_keys = process_keys
+ print("will process keys: ", self.process_keys)
+ self.interpolate_N = interpolate_N
+ assert self.interpolate_N <= 4, \
+ "Interpolation frames should be less than 4. \or there will be frame drop issue."
+ if self.interpolate_N:
+ print(f"We will interpolate {self.interpolate_N} frames between keyframes.")
+
+ # NuScenes Provides 6 cameras
+ self.cam_list = [ # {frame_idx}_{cam_id}.jpg
+ "CAM_FRONT", # "xxx_0.jpg"
+ "CAM_FRONT_LEFT", # "xxx_1.jpg"
+ "CAM_FRONT_RIGHT", # "xxx_2.jpg"
+ "CAM_BACK_LEFT", # "xxx_3.jpg"
+ "CAM_BACK_RIGHT", # "xxx_4.jpg"
+ "CAM_BACK" # "xxx_5.jpg"
+ ]
+ # For each keyframe, the total number of points from LIDAR_TOP is ~35000
+ # the total number of points from ALL RADAR is ~300
+ # Since Radar points are too sparse, we consider only LIDAR_TOP
+ self.lidar_list = ['LIDAR_TOP']
+
+ self.load_dir = load_dir
+ post_fix = f"_{(interpolate_N+1)*2}Hz_{self.data_source}" if interpolate_N > 0 else f"_{self.data_source}"
+ save_dir = save_dir.replace("processed", "processed"+post_fix)
+ self.save_dir = os.path.join(save_dir, split.split('-')[-1])
+ self.workers = int(workers)
+ self.nusc = NuScenes(
+ version=split, dataroot=load_dir, verbose=True
+ )
+ _process_id_list = []
+ for scene_idx in range(len(self.nusc.scene)):
+ scene = self.nusc.scene[scene_idx]
+ if scene['name'] in splits.val:
+ _process_id_list.append(scene_idx)
+ self.process_id_list = _process_id_list[:num_scenes]
+ self.create_folder()
+
+ def convert(self):
+ """Convert action."""
+ print("Start converting ...")
+ if self.process_id_list is None:
+ id_list = range(len(self))
+ else:
+ id_list = self.process_id_list
+ if self.interpolate_N > 0:
+ convert_fn = self.convert_one_interpolated
+ else:
+ convert_fn = self.convert_one
+ track_parallel_progress(convert_fn, id_list, self.workers)
+ # convert_fn(id_list[0])
+ print("\nFinished ...")
+
+ def convert_one(self, scene_idx):
+ """Convert action for single file."""
+ scene = self.nusc.scene[scene_idx]
+ if scene['name'] not in splits.val:
+ return
+ scene_data = self.nusc.get('scene', scene['token'])
+
+ sample_tokens = []
+ sample = self.nusc.get('sample', scene['first_sample_token'])
+ sample_tokens.append(sample['token'])
+ while sample['next'] != '':
+ sample = self.nusc.get('sample', sample['next'])
+ sample_tokens.append(sample['token'])
+
+ # parse video
+ images = {}
+ for camera in self.cam_list:
+ if self.data_source == "drivedreamer2":
+ # TODO: note that the resolution of drivedreamer2 is [1600, 896], we assume directly reshape it to [1600, 900] is acceptable
+ images_single = read_mp4_to_ndarray_list(f'data/nuscenes_gen/{self.data_source}/video_submission/{scene_data["first_sample_token"]}_gen01/{scene_data["first_sample_token"]}_{camera}.mp4', target_size=(1600, 900))
+ elif self.data_source in ("dist4d", "opendwm", "xscene"):
+ images_single = read_mp4_to_ndarray_list(f'data/nuscenes_gen/{self.data_source}/video_submission/{scene_data["first_sample_token"]}_gen0/{scene_data["first_sample_token"]}_{camera}.mp4')
+ else:
+ images_single = read_mp4_to_ndarray_list(f'data/nuscenes_gen/{self.data_source}/video_submisson/{scene_data["first_sample_token"]}_gen0/{scene_data["first_sample_token"]}_{camera}.mp4')
+ images[camera] = dict(zip(sample_tokens[:len(images_single)], images_single))
+ max_frame_idx = len(images_single) - 1
+
+ selected_tokens = sample_tokens[:len(images_single)]
+ if "images" in self.process_keys:
+ for key_frame_idx in range(len(images_single)):
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ img_path = (
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/images/"
+ + f"{str(key_frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+ )
+ cv2.imwrite(img_path, images[cam_name][selected_tokens[key_frame_idx]])
+
+ # self.save_image(scene_data, scene_idx)
+ print(f"Processed images for scene {str(scene_idx).zfill(3)}")
+ if "calib" in self.process_keys:
+ self.save_calib(scene_data, scene_idx, max_frame_idx)
+ print(f"Processed calib for scene {str(scene_idx).zfill(3)}")
+ if "lidar" in self.process_keys:
+ self.save_lidar(scene_data, scene_idx, max_frame_idx)
+ print(f"Processed lidar for scene {str(scene_idx).zfill(3)}")
+ if "dynamic_masks" in self.process_keys:
+ self.save_dynamic_mask(scene_data, scene_idx, max_frame_idx, class_valid='all')
+ self.save_dynamic_mask(scene_data, scene_idx, max_frame_idx, class_valid='human')
+ self.save_dynamic_mask(scene_data, scene_idx, max_frame_idx, class_valid='vehicle')
+ print(f"Processed dynamic masks for scene {str(scene_idx).zfill(3)}")
+
+ # process annotated objects
+ if "objects" in self.process_keys:
+ instances_info, frame_instances = self.save_objects(scene_data, max_frame_idx)
+ print(f"Processed objects for scene {str(scene_idx).zfill(3)}")
+
+ # Save instances info and frame instances
+ instances_info_save_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/instances"
+ with open(f"{instances_info_save_path}/instances_info.json", "w") as fp:
+ json.dump(instances_info, fp, indent=4)
+ with open(f"{instances_info_save_path}/frame_instances.json", "w") as fp:
+ json.dump(frame_instances, fp, indent=4)
+
+ if "objects_vis" in self.process_keys:
+ self.visualize_dynamic_objects(scene_data, scene_idx, instances_info, frame_instances)
+ print(f"Processed objects visualization for scene {str(scene_idx).zfill(3)}")
+
+ def convert_one_interpolated(self, scene_idx):
+ """Convert action for single file."""
+ scene = self.nusc.scene[scene_idx]
+ scene_data = self.nusc.get('scene', scene['token'])
+ keyframe_timestamps = self.get_keyframe_timestamps(scene_data)
+ interpolated_timestamps = self.get_interpolated_timestamps(keyframe_timestamps, self.interpolate_N)
+ interpolated_timestamps = np.array(interpolated_timestamps, dtype=np.int64)
+ if "images" in self.process_keys:
+ self.save_image_interpolated(scene_data, scene_idx, interpolated_timestamps)
+ print(f"Processed images for scene {str(scene_idx).zfill(3)}")
+ if "calib" in self.process_keys:
+ self.save_calib_interpolated(scene_data, scene_idx, interpolated_timestamps)
+ print(f"Processed calib for scene {str(scene_idx).zfill(3)}")
+ if "lidar" in self.process_keys:
+ self.save_lidar_interpolated(scene_data, scene_idx, interpolated_timestamps)
+ print(f"Processed lidar for scene {str(scene_idx).zfill(3)}")
+
+ # process annotated objects
+ if "objects" in self.process_keys:
+ instances_info, _ = self.save_objects(scene_data)
+ print(f"Processed objects for scene {str(scene_idx).zfill(3)}")
+
+ # interpolate the instances info
+ instances_info, frame_instances = self.interpolate_boxes(instances_info)
+ print(f"Interpolated objects for scene {str(scene_idx).zfill(3)}")
+
+ # Save instances info and frame instances
+ instances_info_save_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/instances"
+ with open(f"{instances_info_save_path}/instances_info.json", "w") as fp:
+ json.dump(instances_info, fp, indent=4)
+ with open(f"{instances_info_save_path}/frame_instances.json", "w") as fp:
+ json.dump(frame_instances, fp, indent=4)
+
+ if "objects_vis" in self.process_keys:
+ self.visualize_dynamic_objects_interpolated(scene_data, scene_idx, interpolated_timestamps, instances_info, frame_instances)
+ print(f"Processed objects visualization for scene {str(scene_idx).zfill(3)}")
+
+ if "dynamic_masks" in self.process_keys:
+ self.save_dynamic_mask_interpolated(scene_data, scene_idx, interpolated_timestamps, instances_info, class_valid='all')
+ self.save_dynamic_mask_interpolated(scene_data, scene_idx, interpolated_timestamps, instances_info, class_valid='human')
+ self.save_dynamic_mask_interpolated(scene_data, scene_idx, interpolated_timestamps, instances_info, class_valid='vehicle')
+ print(f"Processed dynamic masks for scene {str(scene_idx).zfill(3)}")
+
+ def __len__(self):
+ """Length of the filename list."""
+ return len(self.process_id_list)
+
+ def get_percam_meta(self, scene_data):
+ """Get the timestamp for each image."""
+ first_sample_token, last_sample_token = scene_data['first_sample_token'], scene_data['last_sample_token']
+ first_sample_record = self.nusc.get('sample', first_sample_token)
+ cams_mata = {
+ cam_name: {
+ "timestamps": [],
+ "is_key_frame": [],
+ "token": []
+ } for cam_name in self.cam_list
+ }
+
+ for cam_name in self.cam_list:
+ cur_img_data = self.nusc.get('sample_data', first_sample_record['data'][cam_name])
+ while True:
+ cams_mata[cam_name]["timestamps"].append(cur_img_data['timestamp'])
+ cams_mata[cam_name]["is_key_frame"].append(cur_img_data['is_key_frame'])
+ cams_mata[cam_name]["token"].append(cur_img_data['token'])
+
+ if cur_img_data['sample_token'] == last_sample_token:
+ break
+ cur_img_data = self.nusc.get('sample_data', cur_img_data['next'])
+
+ return cams_mata
+
+ def get_keyframe_timestamps(self, scene_data):
+ """Get timestamps for keyframes in the scene."""
+ first_sample_token = scene_data['first_sample_token']
+ last_sample_token = scene_data['last_sample_token']
+ curr_sample_record = self.nusc.get('sample', first_sample_token)
+
+ keyframe_timestamps = []
+
+ while True:
+ # Add the timestamp of the current keyframe
+ keyframe_timestamps.append(curr_sample_record['timestamp'])
+
+ if curr_sample_record['token'] == last_sample_token:
+ break
+
+ # Move to the next keyframe
+ curr_sample_record = self.nusc.get('sample', curr_sample_record['next'])
+
+ return keyframe_timestamps
+
+ def get_interpolated_timestamps(self, keyframe_timestamps: List[int], N):
+ """Interpolate timestamps between keyframes."""
+ interpolated_timestamps = []
+
+ for i in range(len(keyframe_timestamps) - 1):
+ start_time = keyframe_timestamps[i]
+ end_time = keyframe_timestamps[i + 1]
+
+ # Calculate the time step for interpolation
+ time_step = (end_time - start_time) / (N + 1)
+
+ # Add the start timestamp
+ interpolated_timestamps.append(start_time)
+
+ # Add N interpolated timestamps
+ for j in range(1, N + 1):
+ interpolated_time = start_time + j * time_step
+ interpolated_timestamps.append(int(interpolated_time))
+
+ # Add the last keyframe timestamp
+ interpolated_timestamps.append(keyframe_timestamps[-1])
+
+ return interpolated_timestamps
+
+ def find_cloest_lidar_tokens(self, scene_data, timestamps: List[int]):
+ """Find the closest LiDAR tokens for given timestamps."""
+ first_sample_token = scene_data['first_sample_token']
+ first_sample_record = self.nusc.get('sample', first_sample_token)
+ lidar_token = first_sample_record['data'][self.lidar_list[0]]
+ lidar_data = self.nusc.get('sample_data', lidar_token)
+
+ # Collect all LiDAR timestamps and tokens
+ lidar_timestamps = []
+ lidar_tokens = []
+ current_lidar = lidar_data
+ while True:
+ lidar_timestamps.append(current_lidar['timestamp'])
+ lidar_tokens.append(current_lidar['token'])
+ if current_lidar['next'] == '':
+ break
+ current_lidar = self.nusc.get('sample_data', current_lidar['next'])
+
+ lidar_timestamps = np.array(lidar_timestamps, dtype=np.int64)
+
+ # Find closest LiDAR tokens for each timestamp
+ closest_tokens = []
+ for timestamp in timestamps:
+ idx = np.argmin(np.abs(lidar_timestamps - timestamp))
+ closest_tokens.append(lidar_tokens[idx])
+
+ # DEBUG USAGE: find is there any duplicated tokens
+ # if len(closest_tokens) != len(set(closest_tokens)):
+ # duplicates = [token for token, count in Counter(closest_tokens).items() if count > 1]
+ # print(f"\nWARNING: {len(duplicates)} duplicated tokens in lidar")
+
+ return closest_tokens
+
+ def find_closest_img_tokens(self, scene_data, timestamps: List[int], cam_name):
+ """Find the closest image tokens for given timestamps for a specific camera."""
+ first_sample_token = scene_data['first_sample_token']
+ first_sample_record = self.nusc.get('sample', first_sample_token)
+ img_token = first_sample_record['data'][cam_name]
+ img_data = self.nusc.get('sample_data', img_token)
+
+ # Collect all image timestamps and tokens for the specified camera
+ img_timestamps = []
+ img_tokens = []
+ current_img = img_data
+ while True:
+ img_timestamps.append(current_img['timestamp'])
+ img_tokens.append(current_img['token'])
+ if current_img['next'] == '':
+ break
+ current_img = self.nusc.get('sample_data', current_img['next'])
+
+ img_timestamps = np.array(img_timestamps, dtype=np.int64)
+
+ # Find closest image tokens for each timestamp
+ closest_tokens = []
+ for timestamp in timestamps:
+ idx = np.argmin(np.abs(img_timestamps - timestamp))
+ closest_tokens.append(img_tokens[idx])
+
+ # DEBUG USAGE: find is there any duplicated tokens
+ # if len(closest_tokens) != len(set(closest_tokens)):
+ # duplicates = [token for token, count in Counter(closest_tokens).items() if count > 1]
+ # print(f"\nWARNING: {len(duplicates)} duplicated tokens in {cam_name}")
+
+ return closest_tokens
+
+ def save_image(self, scene_data, scene_idx):
+ """Parse and save the images in jpg format."""
+ first_sample_token, last_sample_token = scene_data['first_sample_token'], scene_data['last_sample_token']
+ curr_sample_record = self.nusc.get('sample', first_sample_token)
+ key_frame_idx = 0
+
+ while True:
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ cam_data = self.nusc.get('sample_data', curr_sample_record['data'][cam_name])
+ source_path = os.path.join(self.nusc.dataroot, cam_data['filename'])
+ img_path = (
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/images/"
+ + f"{str(key_frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+ )
+ # cp the image
+ os.system(f"cp {source_path} {img_path}")
+
+ if curr_sample_record['next'] == '' or curr_sample_record['token'] == last_sample_token:
+ break
+ key_frame_idx += 1
+ curr_sample_record = self.nusc.get('sample', curr_sample_record['next'])
+
+ def save_image_interpolated(self, scene_data, scene_idx, timestamps: np.array):
+ """Parse and save the interpolated images in jpg format."""
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ # Find the closest image tokens for each timestamp
+ closest_tokens = self.find_closest_img_tokens(scene_data, timestamps, cam_name)
+
+ for frame_idx, token in enumerate(closest_tokens):
+ cam_data = self.nusc.get('sample_data', token)
+ source_path = os.path.join(self.nusc.dataroot, cam_data['filename'])
+ img_path = (
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/images/"
+ + f"{str(frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+ )
+ # Copy the image
+ os.system(f"cp {source_path} {img_path}")
+
+ def save_calib(self, scene_data, scene_idx, max_frame_idx):
+ """Parse and save the calibration data."""
+ first_sample_token, last_sample_token = scene_data['first_sample_token'], scene_data['last_sample_token']
+ curr_sample_record = self.nusc.get('sample', first_sample_token)
+ key_frame_idx = 0
+
+ while True:
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ cam_data = self.nusc.get('sample_data', curr_sample_record['data'][cam_name])
+ calib_data = self.nusc.get('calibrated_sensor', cam_data['calibrated_sensor_token'])
+
+ # Extrinsics (camera to ego)
+ extrinsics_cam_to_ego = np.eye(4)
+ extrinsics_cam_to_ego[:3, :3] = Quaternion(calib_data['rotation']).rotation_matrix
+ extrinsics_cam_to_ego[:3, 3] = np.array(calib_data['translation'])
+
+ # Get ego pose (ego to world)
+ ego_pose_data = self.nusc.get('ego_pose', cam_data['ego_pose_token'])
+ ego_to_world = np.eye(4)
+ ego_to_world[:3, :3] = Quaternion(ego_pose_data['rotation']).rotation_matrix
+ ego_to_world[:3, 3] = np.array(ego_pose_data['translation'])
+
+ # Transform camera extrinsics to world coordinates
+ extrinsics_cam_to_world = ego_to_world @ extrinsics_cam_to_ego
+
+ np.savetxt(
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/extrinsics/"
+ f"{str(key_frame_idx).zfill(3)}_{str(cam_idx)}.txt",
+ extrinsics_cam_to_world
+ )
+
+ # Intrinsics
+ intrinsics = np.array(calib_data['camera_intrinsic'])
+ # fx fy cx cy
+ Ks = [intrinsics[0, 0], intrinsics[1, 1], intrinsics[0, 2], intrinsics[1, 2], 0., 0., 0., 0., 0.]
+ np.savetxt(
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/intrinsics/"
+ f"{str(cam_idx)}.txt",
+ Ks
+ )
+
+ if curr_sample_record['next'] == '' or curr_sample_record['token'] == last_sample_token or key_frame_idx == max_frame_idx:
+ break
+ key_frame_idx += 1
+ curr_sample_record = self.nusc.get('sample', curr_sample_record['next'])
+
+ def save_calib_interpolated(self, scene_data, scene_idx, timestamps: np.array):
+ """Parse and save the interpolated calibration data."""
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ # Find the closest image tokens for each timestamp
+ closest_tokens = self.find_closest_img_tokens(scene_data, timestamps, cam_name)
+
+ for frame_idx, token in enumerate(closest_tokens):
+ cam_data = self.nusc.get('sample_data', token)
+ calib_data = self.nusc.get('calibrated_sensor', cam_data['calibrated_sensor_token'])
+
+ # Extrinsics (camera to ego)
+ extrinsics_cam_to_ego = np.eye(4)
+ extrinsics_cam_to_ego[:3, :3] = Quaternion(calib_data['rotation']).rotation_matrix
+ extrinsics_cam_to_ego[:3, 3] = np.array(calib_data['translation'])
+
+ # Get ego pose (ego to world)
+ ego_pose_data = self.nusc.get('ego_pose', cam_data['ego_pose_token'])
+ ego_to_world = np.eye(4)
+ ego_to_world[:3, :3] = Quaternion(ego_pose_data['rotation']).rotation_matrix
+ ego_to_world[:3, 3] = np.array(ego_pose_data['translation'])
+
+ # Transform camera extrinsics to world coordinates
+ extrinsics_cam_to_world = ego_to_world @ extrinsics_cam_to_ego
+
+ np.savetxt(
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/extrinsics/"
+ f"{str(frame_idx).zfill(3)}_{str(cam_idx)}.txt",
+ extrinsics_cam_to_world
+ )
+
+ # Intrinsics
+ intrinsics = np.array(calib_data['camera_intrinsic'])
+ # fx fy cx cy
+ Ks = [intrinsics[0, 0], intrinsics[1, 1], intrinsics[0, 2], intrinsics[1, 2], 0., 0., 0., 0., 0.]
+ np.savetxt(
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/intrinsics/"
+ f"{str(cam_idx)}.txt",
+ Ks
+ )
+
+ def save_lidar(self, scene_data, scene_idx, max_frame_idx):
+ """Parse and save the lidar data in bin format and lidar pose in world coordinates."""
+ first_sample_token, last_sample_token = scene_data['first_sample_token'], scene_data['last_sample_token']
+ curr_sample_record = self.nusc.get('sample', first_sample_token)
+ key_frame_idx = 0
+
+ while True:
+ lidar_token = curr_sample_record['data'][self.lidar_list[0]]
+ lidar_data = self.nusc.get('sample_data', lidar_token)
+ lidar_path = os.path.join(self.nusc.dataroot, lidar_data['filename'])
+
+ # Load point cloud
+ pc = LidarPointCloud.from_file(lidar_path)
+
+ # Get lidar extrinsics (lidar to ego)
+ calib_data = self.nusc.get('calibrated_sensor', lidar_data['calibrated_sensor_token'])
+ lidar_to_ego = np.eye(4)
+ lidar_to_ego[:3, :3] = Quaternion(calib_data['rotation']).rotation_matrix
+ lidar_to_ego[:3, 3] = np.array(calib_data['translation'])
+
+ # Save lidar points in ego frame
+ lidar_save_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/lidar/{str(key_frame_idx).zfill(3)}.bin"
+ pc.points.T.astype(np.float32).tofile(lidar_save_path)
+
+ # Get ego pose (ego to world)
+ ego_pose_data = self.nusc.get('ego_pose', lidar_data['ego_pose_token'])
+ ego_to_world = np.eye(4)
+ ego_to_world[:3, :3] = Quaternion(ego_pose_data['rotation']).rotation_matrix
+ ego_to_world[:3, 3] = np.array(ego_pose_data['translation'])
+
+ # Calculate lidar pose in world coordinates
+ lidar_to_world = ego_to_world @ lidar_to_ego
+
+ # Save lidar pose in world coordinates
+ np.savetxt(
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/lidar_pose/"
+ f"{str(key_frame_idx).zfill(3)}.txt",
+ lidar_to_world
+ )
+
+ if curr_sample_record['next'] == '' or curr_sample_record['token'] == last_sample_token or key_frame_idx == max_frame_idx:
+ break
+ key_frame_idx += 1
+ curr_sample_record = self.nusc.get('sample', curr_sample_record['next'])
+
+ def save_lidar_interpolated(self, scene_data, scene_idx, timestamps: np.array):
+ """Parse and save the interpolated lidar data in bin format and lidar pose in world coordinates."""
+ # Find the closest LiDAR tokens for each timestamp
+ closest_tokens = self.find_cloest_lidar_tokens(scene_data, timestamps)
+
+ for frame_idx, token in enumerate(closest_tokens):
+ lidar_data = self.nusc.get('sample_data', token)
+ lidar_path = os.path.join(self.nusc.dataroot, lidar_data['filename'])
+
+ # Load point cloud
+ pc = LidarPointCloud.from_file(lidar_path)
+
+ # Get lidar extrinsics (lidar to ego)
+ calib_data = self.nusc.get('calibrated_sensor', lidar_data['calibrated_sensor_token'])
+ lidar_to_ego = np.eye(4)
+ lidar_to_ego[:3, :3] = Quaternion(calib_data['rotation']).rotation_matrix
+ lidar_to_ego[:3, 3] = np.array(calib_data['translation'])
+
+ # Save lidar points in ego frame
+ lidar_save_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/lidar/{str(frame_idx).zfill(3)}.bin"
+ pc.points.T.astype(np.float32).tofile(lidar_save_path)
+
+ # Get ego pose (ego to world)
+ ego_pose_data = self.nusc.get('ego_pose', lidar_data['ego_pose_token'])
+ ego_to_world = np.eye(4)
+ ego_to_world[:3, :3] = Quaternion(ego_pose_data['rotation']).rotation_matrix
+ ego_to_world[:3, 3] = np.array(ego_pose_data['translation'])
+
+ # Calculate lidar pose in world coordinates
+ lidar_to_world = ego_to_world @ lidar_to_ego
+
+ # Save lidar pose in world coordinates
+ np.savetxt(
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/lidar_pose/"
+ f"{str(frame_idx).zfill(3)}.txt",
+ lidar_to_world
+ )
+
+ def save_dynamic_mask(self, scene_data, scene_idx, max_frame_idx, class_valid='all', img_shape=(900, 1600)):
+ """Parse and save the segmentation data."""
+ assert class_valid in ['all', 'human', 'vehicle'], "Invalid class valid"
+ if class_valid == 'all':
+ VALID_CLASSES = NUSCENES_DYNAMIC_CLASSES
+ elif class_valid == 'human':
+ VALID_CLASSES = NUSCENES_NONRIGID_DYNAMIC_CLASSES
+ elif class_valid == 'vehicle':
+ VALID_CLASSES = NUSCENES_RIGID_DYNAMIC_CLASSES
+ mask_dir = f"{self.save_dir}/{str(scene_idx).zfill(3)}/dynamic_masks/{class_valid}"
+ if not os.path.exists(mask_dir):
+ os.makedirs(mask_dir)
+
+ first_sample_token, last_sample_token = scene_data['first_sample_token'], scene_data['last_sample_token']
+ curr_sample_record = self.nusc.get('sample', first_sample_token)
+ key_frame_idx = 0
+
+ while True:
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ cam_data = self.nusc.get('sample_data', curr_sample_record['data'][cam_name])
+ # img_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/images/{str(key_frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+
+ # img = cv2.imread(img_path)
+ dynamic_mask = np.zeros(img_shape[:2], dtype=np.float32)
+
+ anns = [self.nusc.get('sample_annotation', token) for token in curr_sample_record['anns']]
+ valid_anns = [ann for ann in anns if ann['category_name'] in VALID_CLASSES]
+
+ # Get camera calibration data
+ cs_record = self.nusc.get('calibrated_sensor', cam_data['calibrated_sensor_token'])
+ camera_intrinsic = np.array(cs_record['camera_intrinsic'])
+
+ # Get ego pose
+ pose_record = self.nusc.get('ego_pose', cam_data['ego_pose_token'])
+
+ # Project 3D boxes to 2D and create mask
+ for ann in valid_anns:
+ # Create Box object
+ box = Box(ann['translation'], ann['size'], Quaternion(ann['rotation']))
+
+ # Move box to ego vehicle coordinate system
+ box.translate(-np.array(pose_record['translation']))
+ box.rotate(Quaternion(pose_record['rotation']).inverse)
+
+ # Move box to sensor coordinate system
+ box.translate(-np.array(cs_record['translation']))
+ box.rotate(Quaternion(cs_record['rotation']).inverse)
+
+ # Project 3D box to 2D
+ corners_3d = box.corners()
+ corners_2d = view_points(corners_3d, camera_intrinsic, normalize=True)
+
+ # Check if the object is in front of the camera and all corners are in the image
+ # NOTE: we use strict visibility check here, requiring all corners to be visible
+ in_front = np.all(corners_3d[2, :] > 0.1)
+ in_image = np.all(corners_2d[0, :] >= 0) & np.all(corners_2d[0, :] < img_shape[1]) & \
+ np.all(corners_2d[1, :] >= 0) & np.all(corners_2d[1, :] < img_shape[0])
+ if not (in_front and in_image):
+ continue
+
+ # If valid, extract x and y coordinates
+ corners_2d = corners_2d[:2, :]
+
+ # Fill the mask
+ u, v = corners_2d[0, :].astype(np.int32), corners_2d[1, :].astype(np.int32)
+ u = np.clip(u, 0, img_shape[1] - 1)
+ v = np.clip(v, 0, img_shape[0] - 1)
+
+ if u.max() - u.min() == 0 or v.max() - v.min() == 0:
+ continue
+
+ xy = (u.min(), v.min())
+ width = u.max() - u.min()
+ height = v.max() - v.min()
+
+ dynamic_mask[
+ int(xy[1]): int(xy[1] + height),
+ int(xy[0]): int(xy[0] + width)
+ ] = np.maximum(
+ dynamic_mask[
+ int(xy[1]): int(xy[1] + height),
+ int(xy[0]): int(xy[0] + width)
+ ],
+ 1
+ )
+
+ # Save dynamic mask
+ dynamic_mask = np.clip((dynamic_mask > 0.) * 255, 0, 255).astype(np.uint8)
+ dynamic_mask = Image.fromarray(dynamic_mask, "L")
+ dynamic_mask_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/dynamic_masks/{class_valid}/{str(key_frame_idx).zfill(3)}_{str(cam_idx)}.png"
+ dynamic_mask.save(dynamic_mask_path)
+
+ if curr_sample_record['next'] == '' or curr_sample_record['token'] == last_sample_token or key_frame_idx == max_frame_idx:
+ break
+ key_frame_idx += 1
+ curr_sample_record = self.nusc.get('sample', curr_sample_record['next'])
+
+ def save_dynamic_mask_interpolated(self, scene_data, scene_idx, timestamps, instances_info, class_valid='all'):
+ """Parse and save the interpolated segmentation data."""
+ assert class_valid in ['all', 'human', 'vehicle'], "Invalid class valid"
+ if class_valid == 'all':
+ VALID_CLASSES = NUSCENES_DYNAMIC_CLASSES
+ elif class_valid == 'human':
+ VALID_CLASSES = NUSCENES_NONRIGID_DYNAMIC_CLASSES
+ elif class_valid == 'vehicle':
+ VALID_CLASSES = NUSCENES_RIGID_DYNAMIC_CLASSES
+ mask_dir = f"{self.save_dir}/{str(scene_idx).zfill(3)}/dynamic_masks/{class_valid}"
+ if not os.path.exists(mask_dir):
+ os.makedirs(mask_dir)
+
+ for frame_idx, timestamp in enumerate(timestamps):
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ # Find the closest image token for this timestamp
+ closest_token = self.find_closest_img_tokens(scene_data, [timestamp], cam_name)[0]
+ cam_data = self.nusc.get('sample_data', closest_token)
+ img_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/images/{str(frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+
+ img = cv2.imread(img_path)
+ dynamic_mask = np.zeros(img.shape[:2], dtype=np.float32)
+
+ objects = [obj_id for obj_id, obj_info in instances_info.items()
+ if frame_idx in obj_info['frame_annotations']['frame_idx']
+ and obj_info['class_name'] in VALID_CLASSES]
+
+ # Get camera calibration data
+ cs_record = self.nusc.get('calibrated_sensor', cam_data['calibrated_sensor_token'])
+ camera_intrinsic = np.array(cs_record['camera_intrinsic'])
+
+ # Get ego pose
+ pose_record = self.nusc.get('ego_pose', cam_data['ego_pose_token'])
+
+ # Project 3D boxes to 2D and create mask
+ for obj_id in objects:
+ obj_info = instances_info[obj_id]
+ idx_in_obj = obj_info['frame_annotations']['frame_idx'].index(frame_idx)
+ o2w = np.array(obj_info['frame_annotations']['obj_to_world'][idx_in_obj])
+ length, width, height = obj_info['frame_annotations']['box_size'][idx_in_obj]
+
+ # Create Box object
+ box = Box(o2w[:3, 3], [width, length, height], Quaternion(matrix=o2w[:3, :3]))
+
+ # Move box to ego vehicle coordinate system
+ box.translate(-np.array(pose_record['translation']))
+ box.rotate(Quaternion(pose_record['rotation']).inverse)
+
+ # Move box to sensor coordinate system
+ box.translate(-np.array(cs_record['translation']))
+ box.rotate(Quaternion(cs_record['rotation']).inverse)
+
+ # Project 3D box to 2D
+ corners_3d = box.corners()
+ corners_2d = view_points(corners_3d, camera_intrinsic, normalize=True)
+
+ # Check if the object is in front of the camera and all corners are in the image
+ in_front = np.all(corners_3d[2, :] > 0.1)
+ in_image = np.all(corners_2d[0, :] >= 0) & np.all(corners_2d[0, :] < img.shape[1]) & \
+ np.all(corners_2d[1, :] >= 0) & np.all(corners_2d[1, :] < img.shape[0])
+ if not (in_front and in_image):
+ continue
+
+ # If valid, extract x and y coordinates
+ corners_2d = corners_2d[:2, :]
+
+ # Fill the mask
+ u, v = corners_2d[0, :].astype(np.int32), corners_2d[1, :].astype(np.int32)
+ u = np.clip(u, 0, img.shape[1] - 1)
+ v = np.clip(v, 0, img.shape[0] - 1)
+
+ if u.max() - u.min() == 0 or v.max() - v.min() == 0:
+ continue
+
+ xy = (u.min(), v.min())
+ width = u.max() - u.min()
+ height = v.max() - v.min()
+
+ dynamic_mask[
+ int(xy[1]): int(xy[1] + height),
+ int(xy[0]): int(xy[0] + width)
+ ] = np.maximum(
+ dynamic_mask[
+ int(xy[1]): int(xy[1] + height),
+ int(xy[0]): int(xy[0] + width)
+ ],
+ 1
+ )
+
+ # Save dynamic mask
+ dynamic_mask = np.clip((dynamic_mask > 0.) * 255, 0, 255).astype(np.uint8)
+ dynamic_mask = Image.fromarray(dynamic_mask, "L")
+ dynamic_mask_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/dynamic_masks/{class_valid}/{str(frame_idx).zfill(3)}_{str(cam_idx)}.png"
+ dynamic_mask.save(dynamic_mask_path)
+
+ def save_objects(self, scene_data, max_frame_idx):
+ """Parse and save the objects annotation data."""
+ first_sample_token, last_sample_token = scene_data['first_sample_token'], scene_data['last_sample_token']
+ curr_sample_record = self.nusc.get('sample', first_sample_token)
+ key_frame_idx = 0
+
+ instances_info, frame_instances = {}, {}
+ while True:
+ anns = [self.nusc.get('sample_annotation', token) for token in curr_sample_record['anns']]
+
+ for ann in anns:
+ if ann['category_name'] not in NUSCENES_DYNAMIC_CLASSES:
+ continue
+
+ instance_token = ann['instance_token']
+ if instance_token not in instances_info:
+ instances_info[instance_token] = {
+ 'id': instance_token,
+ 'class_name': ann['category_name'],
+ 'frame_annotations': {
+ 'frame_idx': [],
+ 'obj_to_world': [],
+ 'box_size': [],
+ }
+ }
+
+ # Object to world transformation
+ o2w = np.eye(4)
+ o2w[:3, :3] = Quaternion(ann['rotation']).rotation_matrix
+ o2w[:3, 3] = np.array(ann['translation'])
+
+ # Key frames are spaced (interpolate_N + 1) frames apart in the new sequence
+ obj_frame_idx = key_frame_idx * (self.interpolate_N + 1)
+ instances_info[instance_token]['frame_annotations']['frame_idx'].append(obj_frame_idx)
+ instances_info[instance_token]['frame_annotations']['obj_to_world'].append(o2w.tolist())
+ # convert wlh to lwh
+ lwh = [ann['size'][1], ann['size'][0], ann['size'][2]]
+ instances_info[instance_token]['frame_annotations']['box_size'].append(lwh)
+
+ if key_frame_idx not in frame_instances:
+ frame_instances[key_frame_idx] = []
+ frame_instances[key_frame_idx].extend([ann['instance_token'] for ann in anns if ann['category_name'] in NUSCENES_DYNAMIC_CLASSES])
+
+ if curr_sample_record['next'] == '' or curr_sample_record['token'] == last_sample_token or key_frame_idx == max_frame_idx:
+ break
+ key_frame_idx += 1
+ curr_sample_record = self.nusc.get('sample', curr_sample_record['next'])
+
+ # Correct ID mapping
+ id_map = {}
+ for i, (k, v) in enumerate(instances_info.items()):
+ id_map[v["id"]] = i
+
+ # Update keys in instances_info
+ new_instances_info = {}
+ for k, v in instances_info.items():
+ new_instances_info[id_map[v["id"]]] = v
+
+ # Update keys in frame_instances
+ new_frame_instances = {}
+ for k, v in frame_instances.items():
+ new_frame_instances[k] = [id_map[i] for i in v]
+
+ return new_instances_info, new_frame_instances
+
+ def interpolate_boxes(self, instances_info):
+ """Interpolate object positions and sizes between keyframes."""
+ new_instances_info = {}
+ new_frame_instances = {}
+
+ for obj_id, obj_info in instances_info.items():
+ frame_annotations = obj_info['frame_annotations']
+ keyframe_indices = frame_annotations['frame_idx']
+ obj_to_world_list = frame_annotations['obj_to_world']
+ box_size_list = frame_annotations['box_size']
+
+ new_frame_idx = []
+ new_obj_to_world = []
+ new_box_size = []
+
+ for i in range(len(keyframe_indices) - 1):
+ start_frame = keyframe_indices[i]
+ start_transform = np.array(obj_to_world_list[i])
+ end_transform = np.array(obj_to_world_list[i + 1])
+ start_quat = Quaternion(matrix=start_transform[:3, :3])
+ end_quat = Quaternion(matrix=end_transform[:3, :3])
+ start_size = np.array(box_size_list[i])
+ end_size = np.array(box_size_list[i + 1])
+
+ for j in range(self.interpolate_N + 1):
+ t = j / (self.interpolate_N + 1)
+ current_frame = start_frame + j
+
+ # Interpolate translation
+ translation = (1 - t) * start_transform[:3, 3] + t * end_transform[:3, 3]
+
+ # Interpolate rotation using Quaternions
+ current_quat = Quaternion.slerp(start_quat, end_quat, t)
+
+ # Construct interpolated transformation matrix
+ current_transform = np.eye(4)
+ current_transform[:3, :3] = current_quat.rotation_matrix
+ current_transform[:3, 3] = translation
+
+ # Interpolate box size
+ current_size = (1 - t) * start_size + t * end_size
+
+ new_frame_idx.append(current_frame)
+ new_obj_to_world.append(current_transform.tolist())
+ new_box_size.append(current_size.tolist())
+
+ # Add the last keyframe
+ new_frame_idx.append(keyframe_indices[-1])
+ new_obj_to_world.append(obj_to_world_list[-1])
+ new_box_size.append(box_size_list[-1])
+
+ # Update instance info
+ new_instances_info[obj_id] = {
+ 'id': obj_info['id'],
+ 'class_name': obj_info['class_name'],
+ 'frame_annotations': {
+ 'frame_idx': new_frame_idx,
+ 'obj_to_world': new_obj_to_world,
+ 'box_size': new_box_size,
+ }
+ }
+
+ # Update frame instances
+ for frame in new_frame_idx:
+ if frame not in new_frame_instances:
+ new_frame_instances[frame] = []
+ new_frame_instances[frame].append(obj_id)
+
+ return new_instances_info, new_frame_instances
+
+ def visualize_dynamic_objects(self, scene_data, scene_idx, instances_info, frame_instances):
+ """DEBUG: Visualize the dynamic objects' boxes with different colors on the image."""
+ output_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/instances/debug_vis"
+
+ # Get all sample tokens for this scene
+ first_sample_token, last_sample_token = scene_data['first_sample_token'], scene_data['last_sample_token']
+ curr_sample_record = self.nusc.get('sample', first_sample_token)
+ key_frame_idx = 0
+ while True:
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ cam_data = self.nusc.get('sample_data', curr_sample_record['data'][cam_name])
+ img_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/images/{str(key_frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+ canvas = np.array(Image.open(img_path))
+ objects = frame_instances[key_frame_idx]
+
+ if len(objects) > 0:
+ lstProj2d = []
+ color_list = []
+ for obj_id in objects:
+ idx_in_obj = instances_info[obj_id]['frame_annotations']['frame_idx'].index(key_frame_idx)
+ o2w = np.array(
+ instances_info[obj_id]['frame_annotations']['obj_to_world'][idx_in_obj]
+ )
+ length, width, height = instances_info[obj_id]['frame_annotations']['box_size'][idx_in_obj]
+ corners = get_corners(length, width, height)
+
+ # Transform corners to world coordinates
+ corners_world = (o2w[:3, :3] @ corners + o2w[:3, 3:4]).T
+
+ # Get camera calibration data
+ cs_record = self.nusc.get('calibrated_sensor', cam_data['calibrated_sensor_token'])
+ camera_intrinsic = np.array(cs_record['camera_intrinsic'])
+
+ # Get ego pose
+ pose_record = self.nusc.get('ego_pose', cam_data['ego_pose_token'])
+
+ # Transform from world to ego vehicle coordinates
+ corners_ego = corners_world - np.array(pose_record['translation'])
+ corners_ego = Quaternion(pose_record['rotation']).inverse.rotation_matrix @ corners_ego.T
+
+ # Transform from ego to camera coordinates
+ corners_cam = corners_ego - np.array(cs_record['translation']).reshape(3, 1)
+ corners_cam = Quaternion(cs_record['rotation']).inverse.rotation_matrix @ corners_cam
+
+ # Project to 2D
+ corners_2d = view_points(corners_cam, camera_intrinsic, normalize=True)
+
+ # Check if the object is in front of the camera and all corners are in the image
+ in_front = np.all(corners_cam[2, :] > 0.1)
+ in_image = np.all(corners_2d[0, :] >= 0) & np.all(corners_2d[0, :] < canvas.shape[1]) & \
+ np.all(corners_2d[1, :] >= 0) & np.all(corners_2d[1, :] < canvas.shape[0])
+ ok = in_front and in_image
+
+ if ok:
+ projected_points2d = corners_2d[:2, :].T
+ lstProj2d.append(projected_points2d)
+ color_list.append(color_mapper(str(obj_id)))
+
+ lstProj2d = np.asarray(lstProj2d)
+ img_plotted = dump_3d_bbox_on_image(coords=lstProj2d, img=canvas, color=color_list)
+ else:
+ img_plotted = canvas
+
+ img_path = (
+ f"{output_path}/"
+ + f"{str(key_frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+ )
+ Image.fromarray(img_plotted).save(img_path)
+
+ if curr_sample_record['next'] == '' or curr_sample_record['token'] == last_sample_token:
+ break
+ key_frame_idx += 1
+ curr_sample_record = self.nusc.get('sample', curr_sample_record['next'])
+
+ def visualize_dynamic_objects_interpolated(self, scene_data, scene_idx, timestamps, instances_info, frame_instances):
+ """DEBUG: Visualize the interpolated dynamic objects' boxes with different colors on the image."""
+ output_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/instances/debug_vis"
+
+ for frame_idx, timestamp in enumerate(timestamps):
+ for cam_idx, cam_name in enumerate(self.cam_list):
+ # Find the closest image token for this timestamp
+ closest_token = self.find_closest_img_tokens(scene_data, [timestamp], cam_name)[0]
+ cam_data = self.nusc.get('sample_data', closest_token)
+
+ img_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/images/{str(frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+ canvas = np.array(Image.open(img_path))
+ objects = frame_instances.get(frame_idx, [])
+
+ if len(objects) > 0:
+ lstProj2d = []
+ color_list = []
+ for obj_id in objects:
+ obj_info = instances_info[obj_id]
+ idx_in_obj = obj_info['frame_annotations']['frame_idx'].index(frame_idx)
+ o2w = np.array(obj_info['frame_annotations']['obj_to_world'][idx_in_obj])
+ length, width, height = obj_info['frame_annotations']['box_size'][idx_in_obj]
+ corners = get_corners(length, width, height)
+
+ # Transform corners to world coordinates
+ corners_world = (o2w[:3, :3] @ corners + o2w[:3, 3:4]).T
+
+ # Get camera calibration data
+ cs_record = self.nusc.get('calibrated_sensor', cam_data['calibrated_sensor_token'])
+ camera_intrinsic = np.array(cs_record['camera_intrinsic'])
+
+ # Get ego pose
+ pose_record = self.nusc.get('ego_pose', cam_data['ego_pose_token'])
+
+ # Transform from world to ego vehicle coordinates
+ corners_ego = corners_world - np.array(pose_record['translation'])
+ corners_ego = Quaternion(pose_record['rotation']).inverse.rotation_matrix @ corners_ego.T
+
+ # Transform from ego to camera coordinates
+ corners_cam = corners_ego - np.array(cs_record['translation']).reshape(3, 1)
+ corners_cam = Quaternion(cs_record['rotation']).inverse.rotation_matrix @ corners_cam
+
+ # Project to 2D
+ corners_2d = view_points(corners_cam, camera_intrinsic, normalize=True)
+
+ # Check if the object is in front of the camera and all corners are in the image
+ in_front = np.all(corners_cam[2, :] > 0.1)
+ in_image = np.all(corners_2d[0, :] >= 0) & np.all(corners_2d[0, :] < canvas.shape[1]) & \
+ np.all(corners_2d[1, :] >= 0) & np.all(corners_2d[1, :] < canvas.shape[0])
+ ok = in_front and in_image
+
+ if ok:
+ projected_points2d = corners_2d[:2, :].T
+ lstProj2d.append(projected_points2d)
+ color_list.append(color_mapper(str(obj_id)))
+
+ lstProj2d = np.asarray(lstProj2d)
+ img_plotted = dump_3d_bbox_on_image(coords=lstProj2d, img=canvas, color=color_list)
+ else:
+ img_plotted = canvas
+
+ img_path = f"{output_path}/{str(frame_idx).zfill(3)}_{str(cam_idx)}.jpg"
+ Image.fromarray(img_plotted).save(img_path)
+
+ def create_folder(self):
+ """Create folder for data preprocessing."""
+ if self.process_id_list is None:
+ id_list = range(len(self))
+ else:
+ id_list = self.process_id_list
+ for i in id_list:
+ if "images" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/images", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/sky_masks", exist_ok=True)
+ if "calib" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/extrinsics", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/intrinsics", exist_ok=True)
+ if "lidar" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/lidar", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/lidar_pose", exist_ok=True)
+ if "dynamic_masks" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/dynamic_masks/all", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/dynamic_masks/human", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/dynamic_masks/vehicle", exist_ok=True)
+ if "objects" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/instances", exist_ok=True)
+ if "objects_vis" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/instances/debug_vis", exist_ok=True)
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/nuscenes_gen/nuscenes_sourceloader.py b/worldbench/third_party/drivestudio/datasets/nuscenes_gen/nuscenes_sourceloader.py
new file mode 100644
index 0000000..7be6721
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/nuscenes_gen/nuscenes_sourceloader.py
@@ -0,0 +1,460 @@
+import json
+import logging
+import os
+from typing import Dict
+
+import joblib
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from pytorch3d.transforms import matrix_to_quaternion
+from torch import Tensor
+from tqdm import tqdm, trange
+
+from datasets.base.scene_dataset import ModelType
+from datasets.base.lidar_source import SceneLidarSource
+from datasets.base.pixel_source import CameraData, ScenePixelSource
+
+logger = logging.getLogger()
+
+# define each class's node type
+OBJECT_CLASS_NODE_MAPPING = {
+ # Rigid objects (vehicles)
+ "vehicle.bus.bendy": ModelType.RigidNodes,
+ "vehicle.bus.rigid": ModelType.RigidNodes,
+ "vehicle.car": ModelType.RigidNodes,
+ "vehicle.construction": ModelType.RigidNodes,
+ "vehicle.emergency.ambulance": ModelType.RigidNodes,
+ "vehicle.emergency.police": ModelType.RigidNodes,
+ "vehicle.motorcycle": ModelType.RigidNodes,
+ "vehicle.trailer": ModelType.RigidNodes,
+ "vehicle.truck": ModelType.RigidNodes,
+
+ # Humans (SMPL model)
+ "human.pedestrian.adult": ModelType.SMPLNodes,
+ "human.pedestrian.child": ModelType.SMPLNodes,
+ "human.pedestrian.construction_worker": ModelType.SMPLNodes,
+ "human.pedestrian.police_officer": ModelType.SMPLNodes,
+
+ # Potentially deformable objects
+ "human.pedestrian.personal_mobility": ModelType.DeformableNodes,
+ "human.pedestrian.stroller": ModelType.DeformableNodes,
+ "human.pedestrian.wheelchair": ModelType.DeformableNodes,
+ "animal": ModelType.DeformableNodes,
+ "vehicle.bicycle": ModelType.DeformableNodes
+}
+SMPLNODE_CLASSES = [
+ "human.pedestrian.adult",
+ "human.pedestrian.child",
+ "human.pedestrian.construction_worker",
+ "human.pedestrian.police_officer"
+]
+
+# OpenCV to Dataset coordinate transformation
+# opencv coordinate system: x right, y down, z front
+OPENCV2DATASET = np.array(
+ [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]
+)
+
+# NuScenes Camera List:
+# 0: CAM_FRONT
+# 1: CAM_FRONT_LEFT
+# 2: CAM_FRONT_RIGHT
+# 3: CAM_BACK_LEFT
+# 4: CAM_BACK_RIGHT
+# 5: CAM_BACK
+AVAILABLE_CAM_LIST = [0, 1, 2, 3, 4, 5]
+
+class NuScenesCameraData(CameraData):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ def load_calibrations(self):
+ cam_to_worlds, intrinsics, distortions = [], [], []
+
+ # Load the first camera (front) pose to align the world
+ camera_front_start = np.loadtxt(
+ os.path.join(self.data_path, "extrinsics", f"{self.start_timestep:03d}_0.txt")
+ )
+
+ for t in range(self.start_timestep, self.end_timestep):
+ # Load intrinsics
+ intrinsic = np.loadtxt(
+ os.path.join(self.data_path, "intrinsics", f"{self.cam_id}.txt")
+ )
+ fx, fy, cx, cy = intrinsic[0], intrinsic[1], intrinsic[2], intrinsic[3]
+ k1, k2, p1, p2, k3 = intrinsic[4], intrinsic[5], intrinsic[6], intrinsic[7], intrinsic[8]
+
+ # Scale intrinsics
+ fx, fy = (
+ fx * self.load_size[1] / self.original_size[1],
+ fy * self.load_size[0] / self.original_size[0],
+ )
+ cx, cy = (
+ cx * self.load_size[1] / self.original_size[1],
+ cy * self.load_size[0] / self.original_size[0],
+ )
+ intrinsic = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
+ intrinsics.append(intrinsic)
+ distortions.append(np.array([k1, k2, p1, p2, k3]))
+
+ # Load extrinsics
+ cam2world = np.loadtxt(
+ os.path.join(self.data_path, "extrinsics", f"{t:03d}_{self.cam_id}.txt")
+ )
+
+ # Align camera poses with the first camera pose
+ cam2world = np.linalg.inv(camera_front_start) @ cam2world
+
+ # Convert to OpenCV coordinate system
+ cam2world = cam2world @ OPENCV2DATASET
+ cam_to_worlds.append(cam2world)
+
+ self.intrinsics = torch.from_numpy(np.stack(intrinsics, axis=0)).float()
+ self.distortions = torch.from_numpy(np.stack(distortions, axis=0)).float()
+ self.cam_to_worlds = torch.from_numpy(np.stack(cam_to_worlds, axis=0)).float()
+
+ @classmethod
+ def get_camera2worlds(cls, data_path: str, cam_id: str, start_timestep: int, end_timestep: int) -> torch.Tensor:
+ """
+ Returns camera-to-world matrices for the specified camera and time range.
+
+ Args:
+ data_path (str): Path to the dataset.
+ cam_id (str): Camera ID.
+ start_timestep (int): Start timestep.
+ end_timestep (int): End timestep.
+
+ Returns:
+ torch.Tensor: Camera-to-world matrices of shape (num_frames, 4, 4).
+ """
+ # Load the first camera (front) pose to align the world
+ camera_front_start = np.loadtxt(
+ os.path.join(data_path, "extrinsics", f"{start_timestep:03d}_0.txt")
+ )
+
+ cam_to_worlds = []
+
+ for t in range(start_timestep, end_timestep):
+ # Load extrinsics for each timestep
+ cam2world = np.loadtxt(
+ os.path.join(data_path, "extrinsics", f"{t:03d}_{cam_id}.txt")
+ )
+
+ # Align camera poses with the first camera pose
+ cam2world = np.linalg.inv(camera_front_start) @ cam2world
+
+ # Convert to OpenCV coordinate system
+ cam2world = cam2world @ OPENCV2DATASET
+
+ cam_to_worlds.append(cam2world)
+
+ return torch.from_numpy(np.stack(cam_to_worlds, axis=0)).float()
+
+class NuScenesPixelSource(ScenePixelSource):
+ def __init__(
+ self,
+ dataset_name: str,
+ pixel_data_config: OmegaConf,
+ data_path: str,
+ start_timestep: int,
+ end_timestep: int,
+ device: torch.device = torch.device("cpu"),
+ ):
+ super().__init__(dataset_name, pixel_data_config, device=device)
+ self.data_path = data_path
+ self.start_timestep = start_timestep
+ self.end_timestep = end_timestep
+ self.load_data()
+
+ def load_cameras(self):
+ self._timesteps = torch.arange(self.start_timestep, self.end_timestep)
+ self.register_normalized_timestamps()
+
+ for idx, cam_id in enumerate(self.camera_list):
+ logger.info(f"Loading camera {cam_id}")
+ camera = NuScenesCameraData(
+ dataset_name=self.dataset_name,
+ data_path=self.data_path,
+ cam_id=cam_id,
+ start_timestep=self.start_timestep,
+ end_timestep=self.end_timestep,
+ load_dynamic_mask=self.data_cfg.load_dynamic_mask,
+ load_sky_mask=self.data_cfg.load_sky_mask,
+ downscale_when_loading=self.data_cfg.downscale_when_loading[idx],
+ undistort=self.data_cfg.undistort,
+ buffer_downscale=self.buffer_downscale,
+ device=self.device,
+ )
+ camera.load_time(self.normalized_time)
+ unique_img_idx = torch.arange(len(camera), device=self.device) * len(self.camera_list) + idx
+ camera.set_unique_ids(
+ unique_cam_idx = idx,
+ unique_img_idx = unique_img_idx
+ )
+ logger.info(f"Camera {camera.cam_name} loaded.")
+ self.camera_data[cam_id] = camera
+
+ def load_objects(self):
+ """
+ get ground truth bounding boxes of the dynamic objects
+
+ instances_info = {
+ "0": # simplified instance id
+ {
+ "id": str,
+ "class_name": str,
+ "frame_annotations": {
+ "frame_idx": List,
+ "obj_to_world": List,
+ "box_size": List,
+ },
+ ...
+ }
+ frame_instances = {
+ "0": # frame idx
+ List[int] # list of simplified instance ids
+ ...
+ }
+ """
+ instances_info_path = os.path.join(self.data_path, "instances", "instances_info.json")
+ frame_instances_path = os.path.join(self.data_path, "instances", "frame_instances.json")
+ with open(instances_info_path, "r") as f:
+ instances_info = json.load(f)
+ with open(frame_instances_path, "r") as f:
+ frame_instances = json.load(f)
+ # get pose of each instance at each frame
+ # shape (num_frames, num_instances, 4, 4)
+ num_instances = len(instances_info)
+ num_full_frames = len(frame_instances)
+ instances_pose = np.zeros((num_full_frames, num_instances, 4, 4))
+ instances_size = np.zeros((num_full_frames, num_instances, 3))
+ instances_true_id = np.arange(num_instances)
+ instances_model_types = np.ones(num_instances) * -1
+
+ # Load the first camera (front) pose to align the world
+ camera_front_start = np.loadtxt(
+ os.path.join(self.data_path, "extrinsics", f"{self.start_timestep:03d}_0.txt")
+ )
+ for k, v in instances_info.items():
+ instances_model_types[int(k)] = OBJECT_CLASS_NODE_MAPPING[v["class_name"]]
+ for frame_idx, obj_to_world, box_size in zip(v["frame_annotations"]["frame_idx"], v["frame_annotations"]["obj_to_world"], v["frame_annotations"]["box_size"]):
+ # the first ego pose as the origin of the world coordinate system.
+ obj_to_world = np.array(obj_to_world).reshape(4, 4)
+ obj_to_world = np.linalg.inv(camera_front_start) @ obj_to_world
+ instances_pose[frame_idx, int(k)] = np.array(obj_to_world)
+ instances_size[frame_idx, int(k)] = np.array(box_size)
+
+ # get frame valid instances
+ # shape (num_frames, num_instances)
+ per_frame_instance_mask = np.zeros((num_full_frames, num_instances))
+ for frame_idx, valid_instances in frame_instances.items():
+ per_frame_instance_mask[int(frame_idx), valid_instances] = 1
+
+ # select the frames that are in the range of start_timestep and end_timestep
+ instances_pose = torch.from_numpy(instances_pose[self.start_timestep:self.end_timestep]).float()
+ instances_size = torch.from_numpy(instances_size[self.start_timestep:self.end_timestep]).float()
+ instances_true_id = torch.from_numpy(instances_true_id).long()
+ instances_model_types = torch.from_numpy(instances_model_types).long()
+ per_frame_instance_mask = torch.from_numpy(per_frame_instance_mask[self.start_timestep:self.end_timestep]).bool()
+
+ # filter out the instances that are not visible in selected frames
+ ins_frame_cnt = per_frame_instance_mask.sum(dim=0)
+ instances_pose = instances_pose[:, ins_frame_cnt > 0]
+ instances_size = instances_size[:, ins_frame_cnt > 0]
+ instances_true_id = instances_true_id[ins_frame_cnt > 0]
+ instances_model_types = instances_model_types[ins_frame_cnt > 0]
+ per_frame_instance_mask = per_frame_instance_mask[:, ins_frame_cnt > 0]
+
+ # assign to the class
+ # (num_frames, num_instances, 4, 4)
+ self.instances_pose = instances_pose
+ # (num_instances, 3)
+ self.instances_size = instances_size.sum(0) / per_frame_instance_mask.sum(0).unsqueeze(-1)
+ # (num_frames, num_instances)
+ self.per_frame_instance_mask = per_frame_instance_mask
+ # (num_instances)
+ self.instances_true_id = instances_true_id
+ # (num_instances)
+ self.instances_model_types = instances_model_types
+
+ if self.data_cfg.load_smpl:
+ # Collect camera-to-world matrices for all available cameras
+ cam_to_worlds = {}
+ for cam_id in AVAILABLE_CAM_LIST:
+ cam_to_worlds[cam_id] = NuScenesCameraData.get_camera2worlds(
+ self.data_path,
+ str(cam_id),
+ self.start_timestep,
+ self.end_timestep
+ )
+
+ # load SMPL parameters
+ smpl_dict = joblib.load(os.path.join(self.data_path, "humanpose", "smpl.pkl"))
+ frame_num = self.end_timestep - self.start_timestep
+
+ smpl_human_all = {}
+ for fi in tqdm(range(self.start_timestep, self.end_timestep), desc="Loading SMPL"):
+ for instance_id, ins_smpl in smpl_dict.items():
+ if instance_id not in smpl_human_all:
+ smpl_human_all[instance_id] = {
+ "smpl_quats": torch.zeros((frame_num, 24, 4), dtype=torch.float32),
+ "smpl_trans": torch.zeros((frame_num, 3), dtype=torch.float32),
+ "smpl_betas": torch.zeros((frame_num, 10), dtype=torch.float32),
+ "frame_valid": torch.zeros((frame_num), dtype=torch.bool)
+ }
+ smpl_human_all[instance_id]["smpl_quats"][:, :, 0] = 1.0
+ if ins_smpl["valid_mask"][fi]:
+ betas = ins_smpl["smpl"]["betas"][fi]
+ smpl_human_all[instance_id]["smpl_betas"][fi - self.start_timestep] = betas
+
+ body_pose = ins_smpl["smpl"]["body_pose"][fi]
+ smpl_orient = ins_smpl["smpl"]["global_orient"][fi]
+ cam_depend = ins_smpl["selected_cam_idx"][fi].item()
+
+ c2w = cam_to_worlds[cam_depend][fi - self.start_timestep]
+ world_orient = c2w[:3, :3].to(smpl_orient.device) @ smpl_orient.squeeze()
+ smpl_quats = matrix_to_quaternion(
+ torch.cat([world_orient[None, ...], body_pose], dim=0)
+ )
+
+ ii = instances_info[str(instance_id)]['frame_annotations']["frame_idx"].index(fi)
+ o2w = np.array(
+ instances_info[str(instance_id)]['frame_annotations']["obj_to_world"][ii]
+ )
+ o2w = torch.from_numpy(
+ np.linalg.inv(camera_front_start) @ o2w
+ )
+ # box_size = instances_info[str(instance_id)]['frame_annotations']["box_size"][ii]
+
+ smpl_human_all[instance_id]["smpl_quats"][fi - self.start_timestep] = smpl_quats
+ smpl_human_all[instance_id]["smpl_trans"][fi - self.start_timestep] = o2w[:3, 3]
+ smpl_human_all[instance_id]["frame_valid"][fi - self.start_timestep] = True
+
+ self.smpl_human_all = smpl_human_all
+
+class NuScenesLiDARSource(SceneLidarSource):
+ def __init__(
+ self,
+ lidar_data_config: OmegaConf,
+ data_path: str,
+ start_timestep: int,
+ end_timestep: int,
+ device: torch.device = torch.device("cpu"),
+ ):
+ super().__init__(lidar_data_config, device=device)
+ self.data_path = data_path
+ self.start_timestep = start_timestep
+ self.end_timestep = end_timestep
+ self.create_all_filelist()
+ self.load_data()
+
+ def create_all_filelist(self):
+ lidar_filepaths = []
+ lidar_pose_filepaths = []
+ for t in range(self.start_timestep, self.end_timestep):
+ lidar_filepaths.append(
+ os.path.join(self.data_path, "lidar", f"{t:03d}.bin")
+ )
+ lidar_pose_filepaths.append(
+ os.path.join(self.data_path, "lidar_pose", f"{t:03d}.txt")
+ )
+ self.lidar_filepaths = np.array(lidar_filepaths)
+ self.lidar_pose_filepaths = np.array(lidar_pose_filepaths)
+
+ def load_calibrations(self):
+ lidar_to_worlds = []
+
+ # Load the first camera (front) pose to align the world
+ camera_front_start = np.loadtxt(
+ os.path.join(self.data_path, "extrinsics", f"{self.start_timestep:03d}_0.txt")
+ )
+
+ for pose_file in self.lidar_pose_filepaths:
+ lidar_to_world = np.loadtxt(pose_file)
+ # Align lidar poses with the first camera pose
+ lidar_to_world = np.linalg.inv(camera_front_start) @ lidar_to_world
+ lidar_to_worlds.append(lidar_to_world)
+
+ self.lidar_to_worlds = torch.from_numpy(
+ np.stack(lidar_to_worlds, axis=0)
+ ).float()
+
+ def load_lidar(self):
+ origins, directions, ranges = [], [], []
+ timesteps = []
+
+ accumulated_num_original_rays = 0
+ accumulated_num_rays = 0
+ for t in trange(
+ 0, len(self.lidar_filepaths), desc="Loading lidar", dynamic_ncols=True
+ ):
+ lidar_info = np.fromfile(self.lidar_filepaths[t], dtype=np.float32).reshape(-1, 4)
+ original_length = len(lidar_info)
+ accumulated_num_original_rays += original_length
+
+ lidar_points = torch.from_numpy(lidar_info[:, :3]).float()
+ lidar_origins = torch.zeros_like(lidar_points)
+
+ lidar_origins = (
+ self.lidar_to_worlds[t][:3, :3] @ lidar_origins.T
+ + self.lidar_to_worlds[t][:3, 3:4]
+ ).T
+ lidar_points = (
+ self.lidar_to_worlds[t][:3, :3] @ lidar_points.T
+ + self.lidar_to_worlds[t][:3, 3:4]
+ ).T
+ lidar_directions = lidar_points - lidar_origins
+ lidar_ranges = torch.norm(lidar_directions, dim=-1, keepdim=True)
+ lidar_directions = lidar_directions / lidar_ranges
+ lidar_timestamp = torch.ones_like(lidar_ranges).squeeze(-1) * t
+ accumulated_num_rays += len(lidar_ranges)
+
+ origins.append(lidar_origins)
+ directions.append(lidar_directions)
+ ranges.append(lidar_ranges)
+ timesteps.append(lidar_timestamp)
+
+ logger.info(
+ f"Number of lidar rays: {accumulated_num_rays} "
+ f"({accumulated_num_rays / accumulated_num_original_rays * 100:.2f}% of "
+ f"{accumulated_num_original_rays} original rays)"
+ )
+
+ self.origins = torch.cat(origins, dim=0)
+ self.directions = torch.cat(directions, dim=0)
+ self.ranges = torch.cat(ranges, dim=0)
+ self.visible_masks = torch.zeros_like(self.ranges).squeeze().bool()
+ self.colors = torch.ones_like(self.directions)
+
+ self._timesteps = torch.cat(timesteps, dim=0)
+ self.register_normalized_timestamps()
+
+ def get_lidar_rays(self, time_idx: int) -> Dict[str, Tensor]:
+ origins = self.origins[self.timesteps == time_idx]
+ directions = self.directions[self.timesteps == time_idx]
+ ranges = self.ranges[self.timesteps == time_idx]
+ normalized_time = self.normalized_time[self.timesteps == time_idx]
+ return {
+ "lidar_origins": origins,
+ "lidar_viewdirs": directions,
+ "lidar_ranges": ranges,
+ "lidar_normed_time": normalized_time,
+ "lidar_mask": self.timesteps == time_idx,
+ }
+
+ def delete_invisible_pts(self) -> None:
+ if self.visible_masks is not None:
+ num_bf = self.origins.shape[0]
+ self.origins = self.origins[self.visible_masks]
+ self.directions = self.directions[self.visible_masks]
+ self.ranges = self.ranges[self.visible_masks]
+ self._timesteps = self._timesteps[self.visible_masks]
+ self._normalized_time = self._normalized_time[self.visible_masks]
+ self.colors = self.colors[self.visible_masks]
+ logger.info(
+ f"[Lidar] {num_bf - self.visible_masks.sum()} out of {num_bf} points are cleared. {self.visible_masks.sum()} points left."
+ )
+ self.visible_masks = None
+ else:
+ logger.info("[Lidar] No unvisible points to clear.")
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/pandaset/pandaset_human_utils.py b/worldbench/third_party/drivestudio/datasets/pandaset/pandaset_human_utils.py
new file mode 100644
index 0000000..f82f56d
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/pandaset/pandaset_human_utils.py
@@ -0,0 +1,216 @@
+from typing import List
+import os
+import cv2
+import json
+import logging
+import numpy as np
+from tqdm import tqdm
+
+from utils.geometry import (
+ get_corners,
+ project_camera_points_to_image
+)
+from .pandaset_sourceloader import (
+ SMPLNODE_CLASSES,
+ OPENCV2DATASET,
+ AVAILABLE_CAM_LIST,
+)
+
+logger = logging.getLogger()
+
+CAMERA_LIST = AVAILABLE_CAM_LIST
+
+def project_human_boxes(
+ scene_dir: str,
+ camera_list: List[int],
+ save_temp=True,
+ verbose=False,
+ narrow_width_ratio=0.2,
+ fps=12
+):
+ """Project human boxes to 2D image space and save the results to pkl files
+
+ Args:
+ scene_dir: str, path to the scene directory
+ camera_list: List[int], a list of camera ids to be processed
+ save_temp: bool, whether to save the intermediate results
+ verbose: bool, whether to visualize the projected boxes
+ narrow_width_ratio: sometimes the projected boxes are too wide
+ we can narrow them with this ratio to get more accurate results
+ fps: int, FPS for the visualization video
+
+ Returns:
+ collector_all: dict, a dictionary containing the projected boxes for each camera
+ """
+ # check if the necessary directories exist
+ images_dir = f'{scene_dir}/images'
+ poses_dir = f'{scene_dir}/ego_pose'
+ extrinsics_dir = f'{scene_dir}/extrinsics'
+ intrinsics_dir = f'{scene_dir}/intrinsics'
+ instances_dir = f'{scene_dir}/instances'
+ valid_paths = [images_dir, poses_dir, extrinsics_dir, intrinsics_dir, instances_dir]
+ for path in valid_paths:
+ assert os.path.exists(path), \
+ f"Path {path} does not exist, you need to run waymo preprocess to generate the necessary files"
+
+ # create directories for saving the results
+ save_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks'
+ if not os.path.exists(save_dir):
+ os.makedirs(save_dir)
+
+ if verbose:
+ # create directories for saving the intermediate visualization results
+ video_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks/vis'
+ if not os.path.exists(video_dir):
+ os.makedirs(video_dir)
+ per_human_img_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks/vis/images'
+ if not os.path.exists(per_human_img_dir):
+ os.makedirs(per_human_img_dir)
+
+ # load instances and frame infos
+ frame_infos = json.load(
+ open(f'{instances_dir}/frame_instances.json')
+ )
+ instances_meta = json.load(
+ open(f'{instances_dir}/instances_info.json')
+ )
+
+ collector_all = {}
+ # iterate over each camera
+ for cam_id in camera_list:
+ # check if already processed
+ pkl_path = os.path.join(save_dir, f"{cam_id}.pkl")
+ if os.path.exists(pkl_path):
+ collector_all[cam_id] = json.load(open(pkl_path))
+ logger.info(f"Results for camera {cam_id} already exists at {pkl_path}")
+ continue
+
+ if verbose:
+ per_cam_vis_dir = os.path.join(per_human_img_dir, f"{cam_id}")
+ if not os.path.exists(per_cam_vis_dir):
+ os.makedirs(per_cam_vis_dir)
+
+ collector = {}
+ frames = []
+ for frame_id, frame_ins_list in frame_infos.items():
+ frame_id = int(frame_id)
+
+ # define empty instance collector for each frame
+ frame_collector = {
+ "gt_bbox": [],
+ "extra_data": {
+ "gt_track_id": [],
+ "gt_class": [],
+ }
+ }
+
+ # load extrinsic
+ cam2world = np.loadtxt(os.path.join(extrinsics_dir, f"{frame_id:03d}_{cam_id}.txt"))
+ cam2world = cam2world @ OPENCV2DATASET
+
+ # load intrinsic
+ Ks = np.loadtxt(os.path.join(intrinsics_dir, f"{cam_id}.txt"))
+ fx, fy, cx, cy = Ks[0], Ks[1], Ks[2], Ks[3]
+ # k1, k2, p1, p2, k3 = Ks[4], Ks[5], Ks[6], Ks[7], Ks[8]
+ intrinsic = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float32)
+
+ # load image
+ ori_image = cv2.imread(
+ os.path.join(images_dir, f"{str(frame_id).zfill(3)}_{cam_id}.jpg")
+ )
+ image = ori_image.copy()
+ H, W = image.shape[:2]
+
+ # iterate over each instance in the frame
+ if len(frame_ins_list) > 0:
+ # if there are pedestrians in this frame, project the boxes to the image
+ image_plotted = image.copy()
+ for instance_id in frame_ins_list:
+ ins = instances_meta[str(instance_id)]
+
+ if ins["class_name"] not in SMPLNODE_CLASSES:
+ continue
+
+ ins_anno = ins["frame_annotations"]
+ index = ins_anno['frame_idx'].index(frame_id)
+ obj_to_world = np.array(ins_anno['obj_to_world'][index])
+ l, w, h = ins_anno['box_size'][index]
+
+ # get box corners in object space
+ corners = get_corners(l, w, h)
+ # transform box corners to world space
+ corners_world = obj_to_world[:3, :3] @ corners + obj_to_world[:3, 3:4]
+ # transform box corners to image space
+ world2cam = np.linalg.inv(cam2world)
+ corners_cam = world2cam[:3, :3] @ corners_world + world2cam[:3, 3:4]
+ cam_points, depth = project_camera_points_to_image(corners_cam.T, intrinsic)
+
+ x_min, y_min = np.min(cam_points, axis=0)
+ x_max, y_max = np.max(cam_points, axis=0)
+ # clip left and right with this ratio
+ if narrow_width_ratio > 0.:
+ length = x_max - x_min
+ x_min += length * narrow_width_ratio
+ x_max -= length * narrow_width_ratio
+
+ # clip the box to the image
+ original_area = (x_max - x_min) * (y_max - y_min)
+ x_min, x_max = np.clip(x_min, 0, W), np.clip(x_max, 0, W)
+ y_min, y_max = np.clip(y_min, 0, H), np.clip(y_max, 0, H)
+ new_area = (x_max - x_min) * (y_max - y_min)
+
+ # filter out boxes that are too small or too large
+ behind = depth.max() < 0
+ too_small = new_area < W * H * (0.03)**2
+ too_large = new_area > W * H / 1.1
+ too_far = np.linalg.norm(obj_to_world[:3, 3] - cam2world[:3, 3]) > 40
+ clip_large = new_area / original_area < 1/3
+ if too_small or too_large or clip_large or behind or too_far:
+ continue
+
+ # gt box on image
+ gt_box = [x_min, y_min, x_max - x_min, y_max - y_min]
+
+ # save the projected box to the collector
+ frame_collector["gt_bbox"].append(gt_box)
+ frame_collector["extra_data"]["gt_track_id"].append(instance_id)
+ frame_collector["extra_data"]["gt_class"].append([0])
+
+ if verbose:
+ # visualize the projected boxes of ONE instance
+ raw_image = cv2.rectangle(
+ ori_image.copy(), (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 2
+ )
+ raw_image_path = os.path.join(per_cam_vis_dir, f"{frame_id}_{instance_id}.jpg")
+ cv2.imwrite(raw_image_path, raw_image)
+
+ # add this instance to the image
+ image_plotted = cv2.rectangle(
+ image, (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 2
+ )
+
+ if verbose:
+ frames.append(image_plotted)
+ else:
+ # if no instance in this frame, just save the original image
+ if verbose:
+ frames.append(ori_image)
+
+ collector[frame_id] = frame_collector
+
+ if verbose:
+ height, width = frames[0].shape[:2]
+ output_path = os.path.join(video_dir, f"cam_{cam_id}.mp4")
+ fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+ out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+ for frame in tqdm(frames, desc=f"Writing video for camera {cam_id}"):
+ out.write(frame)
+ out.release()
+
+ if save_temp:
+ # save collector to pkl
+ json.dump(collector, open(pkl_path, "w"))
+
+ collector_all[cam_id] = collector
+
+ return collector_all
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/pandaset/pandaset_preprocess.py b/worldbench/third_party/drivestudio/datasets/pandaset/pandaset_preprocess.py
new file mode 100644
index 0000000..028512d
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/pandaset/pandaset_preprocess.py
@@ -0,0 +1,603 @@
+import json
+import os
+
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+
+from pandaset import DataSet as PandaSet, geometry
+from pandaset.sequence import Sequence
+
+from datasets.tools.multiprocess_utils import track_parallel_progress
+from utils.visualization import color_mapper, dump_3d_bbox_on_image
+
+PANDA_LABELS = [
+ 'Animals - Other', 'Bicycle', 'Bus', 'Car',
+ 'Cones', 'Construction Signs', 'Emergency Vehicle', 'Medium-sized Truck',
+ 'Motorcycle', 'Motorized Scooter', 'Other Vehicle - Construction Vehicle', 'Other Vehicle - Pedicab',
+ 'Other Vehicle - Uncommon', 'Pedestrian', 'Pedestrian with Object', 'Personal Mobility Device',
+ 'Pickup Truck', 'Pylons', 'Road Barriers', 'Rolling Containers',
+ 'Semi-truck', 'Signs', 'Temporary Construction Barriers', 'Towed Object',
+ 'Train', 'Tram / Subway'
+]
+
+PANDA_NONRIGID_DYNAMIC_CLASSES = [
+ 'Pedestrian', 'Pedestrian with Object', 'Bicycle', 'Animals - Other'
+]
+
+PANDA_RIGID_DYNAMIC_CLASSES = [
+ 'Bus', 'Car', 'Emergency Vehicle', 'Medium-sized Truck',
+ 'Motorcycle', 'Motorized Scooter', 'Other Vehicle - Construction Vehicle', 'Other Vehicle - Pedicab',
+ 'Other Vehicle - Uncommon', 'Personal Mobility Device', 'Pickup Truck',
+ 'Semi-truck', 'Train', 'Tram / Subway'
+]
+
+PANDA_DYNAMIC_CLASSES = PANDA_NONRIGID_DYNAMIC_CLASSES + PANDA_RIGID_DYNAMIC_CLASSES
+
+class PandaSetProcessor(object):
+ """Process PandaSet.
+
+ Args:
+ load_dir (str): Directory to load data.
+ save_dir (str): Directory to save data in KITTI format.
+ prefix (str): Prefix of filename.
+ workers (int, optional): Number of workers for the parallel process.
+ Defaults to 64.
+ Defaults to False.
+ save_cam_sync_labels (bool, optional): Whether to save cam sync labels.
+ Defaults to True.
+ """
+
+ def __init__(
+ self,
+ load_dir,
+ save_dir,
+ process_keys=[
+ "images",
+ "lidar",
+ "calib",
+ "pose",
+ "dynamic_masks",
+ "objects"
+ ],
+ process_id_list=None,
+ workers=64,
+ ):
+ self.process_id_list = process_id_list
+ self.process_keys = process_keys
+ print("will process keys: ", self.process_keys)
+
+ # PandaSet Provides 6 cameras and 2 lidars
+ self.cam_list = [ # {frame_idx}_{cam_id}.jpg
+ "front_camera", # "xxx_0.jpg"
+ "front_left_camera", # "xxx_1.jpg"
+ "front_right_camera", # "xxx_2.jpg"
+ "left_camera", # "xxx_3.jpg"
+ "right_camera", # "xxx_4.jpg"
+ "back_camera" # "xxx_5.jpg"
+ ]
+ # 0: mechanical 360° LiDAR, 1: front-facing LiDAR, -1: All LiDARs
+ self.lidar_list = [-1]
+
+ self.load_dir = load_dir
+ self.save_dir = f"{save_dir}"
+ self.workers = int(workers)
+ self.pandaset = PandaSet(load_dir)
+ self.create_folder()
+
+ def convert(self):
+ """Convert action."""
+ print("Start converting ...")
+ if self.process_id_list is None:
+ id_list = range(len(self))
+ else:
+ id_list = self.process_id_list
+ track_parallel_progress(self.convert_one, id_list, self.workers)
+ print("\nFinished ...")
+
+ def convert_one(self, scene_idx):
+ """Convert action for single file.
+
+ Args:
+ scene_idx (str): Scene index.
+ """
+ scene_data = self.pandaset[scene_idx]
+ scene_data.load()
+ num_frames = sum(1 for _ in scene_data.timestamps)
+ for frame_idx in tqdm(range(num_frames), desc=f"File {scene_idx}", total=num_frames, dynamic_ncols=True):
+ if "images" in self.process_keys:
+ self.save_image(scene_data, scene_idx, frame_idx)
+ if "calib" in self.process_keys:
+ self.save_calib(scene_data, scene_idx, frame_idx)
+ if "lidar" in self.process_keys:
+ self.save_lidar(scene_data, scene_idx, frame_idx)
+ if "pose" in self.process_keys:
+ self.save_pose(scene_data, scene_idx, frame_idx)
+ if "3dbox_vis" in self.process_keys:
+ # visualize 3d box, debug usage
+ self.visualize_3dbox(scene_data, scene_idx, frame_idx)
+ if "dynamic_masks" in self.process_keys:
+ self.save_dynamic_mask(scene_data, scene_idx, frame_idx, class_valid='all')
+ self.save_dynamic_mask(scene_data, scene_idx, frame_idx, class_valid='human')
+ self.save_dynamic_mask(scene_data, scene_idx, frame_idx, class_valid='vehicle')
+
+ # save instances info
+ if "objects" in self.process_keys:
+ instances_info = self.save_objects(scene_data, num_frames)
+
+ # solve duplicated objects from different lidars
+ duplicated_id_pairs = []
+ for k, v in instances_info.items():
+ if v["sibling_id"] != '-':
+ # find if the pair is already in the list
+ if (v["id"], v["sibling_id"]) in duplicated_id_pairs or (v["sibling_id"], v["id"]) in duplicated_id_pairs:
+ continue
+ else:
+ duplicated_id_pairs.append((v["id"], v["sibling_id"]))
+
+ for pair in duplicated_id_pairs:
+ # check if all in the pair are in the instances_info
+ if pair[0] not in instances_info:
+ # print(f"WARN: {pair[0]} not in instances_info")
+ continue
+ elif pair[1] not in instances_info:
+ # print(f"WARN: {pair[1]} not in instances_info")
+ continue
+ else:
+ # keep the longer one in pairs
+ if len(instances_info[pair[0]]['frame_annotations']['frame_idx']) > \
+ len(instances_info[pair[1]]['frame_annotations']['frame_idx']):
+ instances_info.pop(pair[1])
+ else:
+ instances_info.pop(pair[0])
+
+ # rough filter stationary objects
+ # if all the annotations of an object are stationary, remove it
+ static_ids = []
+ for k, v in instances_info.items():
+ if all(v['frame_annotations']['stationary']):
+ static_ids.append(v['id'])
+ print(f"INFO: {len(static_ids)} static objects removed")
+ for static_id in static_ids:
+ instances_info.pop(static_id)
+ print(f"INFO: Final number of objects: {len(instances_info)}")
+
+ frame_instances = {}
+ # update frame_instances
+ for frame_idx in range(num_frames):
+ # must ceate a object for each frame
+ frame_instances[frame_idx] = []
+ for k, v in instances_info.items():
+ if frame_idx in v['frame_annotations']['frame_idx']:
+ frame_instances[frame_idx].append(v["id"])
+
+ # verbose: visualize the instances on the image (Debug Usage)
+ if "objects_vis" in self.process_keys:
+ self.visualize_dynamic_objects(
+ scene_data, scene_idx,
+ instances_info=instances_info,
+ frame_instances=frame_instances
+ )
+
+ # correct id
+ id_map = {}
+ for i, (k, v) in enumerate(instances_info.items()):
+ id_map[v["id"]] = i
+ # update keys in instances_info
+ new_instances_info = {}
+ for k, v in instances_info.items():
+ new_instances_info[id_map[v["id"]]] = v
+ # update keys in frame_instances
+ new_frame_instances = {}
+ for k, v in frame_instances.items():
+ new_frame_instances[k] = [id_map[i] for i in v]
+
+ # write as json
+ instances_dir = f"{self.save_dir}/{str(scene_idx).zfill(3)}/instances"
+ with open(f"{instances_dir}/instances_info.json", "w") as fp:
+ json.dump(new_instances_info, fp, indent=4)
+ with open(f"{instances_dir}/frame_instances.json", "w") as fp:
+ json.dump(new_frame_instances, fp, indent=4)
+
+ def __len__(self):
+ """Length of the filename list."""
+ return len(self.process_id_list)
+
+ def save_image(self, scene_data: Sequence, scene_idx, frame_idx):
+ """Parse and save the images in jpg format.
+
+ Args:
+ scene_data (:obj:`Sequence`): PandaSet sequence.
+ scene_idx (str): Current file index.
+ frame_idx (int): Current frame index.
+ """
+ for idx, cam in enumerate(self.cam_list):
+ img_path = (
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/images/"
+ + f"{str(frame_idx).zfill(3)}_{str(idx)}.jpg"
+ )
+ # write PIL Image to jpg
+ image = scene_data.camera[cam][frame_idx]
+ image.save(img_path)
+
+ def save_calib(self, scene_data: Sequence, scene_idx, frame_idx):
+ """Parse and save the calibration data.
+
+ Args:
+ scene_data (:obj:`Sequence`): PandaSet sequence.
+ scene_idx (str): Current file index.
+ frame_idx (int): Current frame index.
+ """
+ for idx, cam in enumerate(self.cam_list):
+ camera = scene_data.camera[cam]
+ poses = camera.poses[frame_idx]
+ c2w = geometry._heading_position_to_mat(poses['heading'], poses['position'])
+ K = camera.intrinsics
+ intrinsics = [K.fx, K.fy, K.cx, K.cy, 0.0, 0.0, 0.0, 0.0, 0.0]
+
+ np.savetxt(
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/extrinsics/"
+ + f"{str(frame_idx).zfill(3)}_{str(idx)}.txt",
+ c2w,
+ )
+ np.savetxt(
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/intrinsics/"
+ + f"{str(idx)}.txt",
+ intrinsics,
+ )
+
+ def save_lidar(self, scene_data: Sequence, scene_idx, frame_idx):
+ """Parse and save the lidar data in psd format.
+
+ Args:
+ scene_data (:obj:`Sequence`): PandaSet sequence.
+ scene_idx (str): Current file index.
+ frame_idx (int): Current frame index.
+ """
+ pc_world = scene_data.lidar[frame_idx].to_numpy()
+ # index x y z i t d
+ # 0 -75.131138 -79.331690 3.511804 7.0 1.557540e+09 0
+ # 1 -112.588306 -118.666002 1.423499 31.0 1.557540e+09 0
+ # - `i`: `float`: Reflection intensity in a range `[0,255]`
+ # - `t`: `float`: Recorded timestamp for specific point
+ # - `d`: `int`: Sensor ID. `0` -> mechnical 360° LiDAR, `1` -> forward-facing LiDAR
+ lidar_poses = scene_data.lidar.poses[frame_idx]
+
+ # save lidar pts in ego coordinate system
+ pcd_ego = geometry.lidar_points_to_ego(
+ pc_world[:, :3], lidar_poses
+ )
+ intensity = pc_world[:, 3]
+ laser_ids = pc_world[:, 5]
+
+ point_cloud = np.column_stack(
+ (
+ pcd_ego,
+ intensity,
+ laser_ids
+ )
+ )
+
+ pc_path = (
+ f"{self.save_dir}/"
+ + f"{str(scene_idx).zfill(3)}/lidar/{str(frame_idx).zfill(3)}.bin"
+ )
+ point_cloud.astype(np.float32).tofile(pc_path)
+
+ def save_pose(self, scene_data: Sequence, scene_idx, frame_idx):
+ """Parse and save the pose data.
+
+ Since pandaset does not provide the ego pose, we use the lidar pose as the ego pose.
+
+ Args:
+ scene_data (:obj:`Sequence`): PandaSet sequence.
+ scene_idx (str): Current file index.
+ frame_idx (int): Current frame index.
+ """
+ lidar_poses = scene_data.lidar.poses[frame_idx]
+ lidar_to_world = geometry._heading_position_to_mat(lidar_poses['heading'], lidar_poses['position'])
+
+ np.savetxt(
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/ego_pose/"
+ + f"{str(frame_idx).zfill(3)}.txt",
+ lidar_to_world,
+ )
+
+ def visualize_3dbox(self, scene_data: Sequence, scene_idx, frame_idx):
+ """DEBUG: Visualize the 3D bounding box on the image.
+ Visualize the 3D bounding box all with the same COLOR.
+ If you want to visualize the 3D bounding box with different colors, please use the `visualize_dynamic_objects` function.
+
+ Args:
+ scene_data (:obj:`Sequence`): PandaSet sequence.
+ scene_idx (str): Current file index.
+ frame_idx (int): Current frame index.
+ """
+ for idx, cam in enumerate(self.cam_list):
+ img_path = (
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/images/"
+ + f"{str(frame_idx).zfill(3)}_{str(idx)}.jpg"
+ )
+ canvas = np.array(Image.open(img_path))
+
+ camera = scene_data.camera[cam]
+ cuboids = scene_data.cuboids[frame_idx]
+
+ lstProj2d = []
+ recorded_id = []
+ for _, row in cuboids.iterrows():
+ if row["label"] not in PANDA_DYNAMIC_CLASSES or row["stationary"]:
+ continue
+ if not row["cuboids.sensor_id"] == -1:
+ if row["cuboids.sibling_id"] in recorded_id:
+ continue
+ recorded_id.append(row["uuid"])
+ box = [
+ row[ "position.x"], row[ "position.y"], row[ "position.z"],
+ row["dimensions.x"], row["dimensions.y"], row["dimensions.z"],
+ row["yaw"]
+ ]
+ corners = geometry.center_box_to_corners(box)
+
+ projected_points2d, _, _ = geometry.projection(
+ lidar_points=corners,
+ camera_data=camera[frame_idx],
+ camera_pose=camera.poses[frame_idx],
+ camera_intrinsics=camera.intrinsics,
+ filter_outliers=True
+ )
+ projected_points2d = projected_points2d.tolist()
+ if len(projected_points2d) == 8:
+ lstProj2d.append(projected_points2d)
+
+ lstProj2d = np.asarray(lstProj2d)
+
+ img_plotted = dump_3d_bbox_on_image(coords=lstProj2d, img=canvas, color=(255,0,0))
+
+ # save
+ img_path = (
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/3dbox_vis/"
+ + f"{str(frame_idx).zfill(3)}_{str(idx)}.jpg"
+ )
+ Image.fromarray(img_plotted).save(img_path)
+
+ def visualize_dynamic_objects(
+ self, scene_data: Sequence, scene_idx,
+ instances_info: dict, frame_instances: dict
+ ):
+ """DEBUG: Visualize the dynamic objects'box with different colors on the image.
+
+ Args:
+ scene_data (:obj:`Sequence`): PandaSet sequence.
+ scene_idx (str): Current file index.
+ instances_info (dict): Instances information.
+ frame_instances (dict): Frame instances.
+ """
+ output_path = f"{self.save_dir}/{str(scene_idx).zfill(3)}/instances/debug_vis"
+
+ num_frames = sum(1 for _ in scene_data.timestamps)
+ for frame_idx in range(num_frames):
+ for idx, cam in enumerate(self.cam_list):
+ img_path = (
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/images/"
+ + f"{str(frame_idx).zfill(3)}_{str(idx)}.jpg"
+ )
+ canvas = np.array(Image.open(img_path))
+
+ camera = scene_data.camera[cam]
+ if frame_idx in frame_instances:
+ objects = frame_instances[frame_idx]
+
+ lstProj2d = []
+ color_list = []
+ for obj_id in objects:
+ idx_in_obj = instances_info[obj_id]['frame_annotations']['frame_idx'].index(frame_idx)
+ o2w = instances_info[obj_id]['frame_annotations']['obj_to_world'][idx_in_obj]
+ o2w = np.array(o2w)
+ length, width, height = instances_info[obj_id]['frame_annotations']['box_size'][idx_in_obj]
+ half_dim_x, half_dim_y, half_dim_z = length/2.0, width/2.0, height/2.0
+ corners = np.array(
+ [[half_dim_x, half_dim_y, -half_dim_z],
+ [half_dim_x, -half_dim_y, -half_dim_z],
+ [-half_dim_x, -half_dim_y, -half_dim_z],
+ [-half_dim_x, half_dim_y, -half_dim_z],
+ [half_dim_x, half_dim_y, half_dim_z],
+ [half_dim_x, -half_dim_y, half_dim_z],
+ [-half_dim_x, -half_dim_y, half_dim_z],
+ [-half_dim_x, half_dim_y, half_dim_z]]
+ )
+ corners = (o2w[:3, :3] @ corners.T + o2w[:3, [3]]).T
+
+ projected_points2d, _, _ = geometry.projection(
+ lidar_points=corners,
+ camera_data=camera[frame_idx],
+ camera_pose=camera.poses[frame_idx],
+ camera_intrinsics=camera.intrinsics,
+ filter_outliers=True
+ )
+ projected_points2d = projected_points2d.tolist()
+ if len(projected_points2d) == 8:
+ lstProj2d.append(projected_points2d)
+ color_list.append(color_mapper(obj_id))
+
+ lstProj2d = np.asarray(lstProj2d)
+ img_plotted = dump_3d_bbox_on_image(coords=lstProj2d, img=canvas, color=color_list)
+
+ img_path = (
+ f"{output_path}/"
+ + f"{str(frame_idx).zfill(3)}_{str(idx)}.jpg"
+ )
+ Image.fromarray(img_plotted).save(img_path)
+
+ def save_dynamic_mask(self, scene_data: Sequence, scene_idx, frame_idx, class_valid='all'):
+ """Parse and save the segmentation data.
+
+ Args:
+ scene_data (:obj:`Sequence`): PandaSet sequence.
+ scene_idx (str): Current file index.
+ frame_idx (int): Current frame index.
+ class_valid (str): Class valid for dynamic mask.
+ """
+ assert class_valid in ['all', 'human', 'vehicle'], "Invalid class valid"
+ if class_valid == 'all':
+ VALID_CLASSES = PANDA_DYNAMIC_CLASSES
+ elif class_valid == 'human':
+ VALID_CLASSES = PANDA_NONRIGID_DYNAMIC_CLASSES
+ elif class_valid == 'vehicle':
+ VALID_CLASSES = PANDA_RIGID_DYNAMIC_CLASSES
+ mask_dir = f"{self.save_dir}/{str(scene_idx).zfill(3)}/dynamic_masks/{class_valid}"
+ if not os.path.exists(mask_dir):
+ os.makedirs(mask_dir)
+
+ for idx, cam in enumerate(self.cam_list):
+ # dynamic_mask
+ img_path = (
+ f"{self.save_dir}/{str(scene_idx).zfill(3)}/images/"
+ + f"{str(frame_idx).zfill(3)}_{str(idx)}.jpg"
+ )
+ img_shape = np.array(Image.open(img_path))
+ dynamic_mask = np.zeros_like(img_shape, dtype=np.float32)[..., 0]
+
+ camera = scene_data.camera[cam]
+ cuboids = scene_data.cuboids[frame_idx]
+
+ recorded_id = []
+ for _, row in cuboids.iterrows():
+ if row["label"] not in VALID_CLASSES or row["stationary"]:
+ continue
+ if not row["cuboids.sensor_id"] == -1:
+ if row["cuboids.sibling_id"] in recorded_id:
+ continue
+ recorded_id.append(row["uuid"])
+
+ box = [
+ row[ "position.x"], row[ "position.y"], row[ "position.z"],
+ row["dimensions.x"], row["dimensions.y"], row["dimensions.z"],
+ row["yaw"]
+ ]
+ corners = geometry.center_box_to_corners(box)
+
+ projected_points2d, _, _ = geometry.projection(
+ lidar_points=corners,
+ camera_data=camera[frame_idx],
+ camera_pose=camera.poses[frame_idx],
+ camera_intrinsics=camera.intrinsics,
+ filter_outliers=True
+ )
+ # Skip object if any corner projection failed. Note that this is very
+ # strict and can lead to exclusion of some partially visible objects.
+ if not len(projected_points2d) == 8:
+ continue
+ u, v= projected_points2d[:, 0], projected_points2d[:, 1]
+ u = u.astype(np.int32)
+ v = v.astype(np.int32)
+
+ # Clip box to image bounds.
+ u = np.clip(u, 0, camera[frame_idx].size[0])
+ v = np.clip(v, 0, camera[frame_idx].size[1])
+
+ if u.max() - u.min() == 0 or v.max() - v.min() == 0:
+ continue
+
+ # Draw projected 2D box onto the image.
+ xy = (u.min(), v.min())
+ width = u.max() - u.min()
+ height = v.max() - v.min()
+ # max pooling
+ dynamic_mask[
+ int(xy[1]) : int(xy[1] + height),
+ int(xy[0]) : int(xy[0] + width),
+ ] = np.maximum(
+ dynamic_mask[
+ int(xy[1]) : int(xy[1] + height),
+ int(xy[0]) : int(xy[0] + width),
+ ],
+ 1.,
+ )
+ dynamic_mask = np.clip((dynamic_mask > 0.) * 255, 0, 255).astype(np.uint8)
+ dynamic_mask = Image.fromarray(dynamic_mask, "L")
+ dynamic_mask_path = os.path.join(mask_dir, f"{str(frame_idx).zfill(3)}_{str(idx)}.png")
+ dynamic_mask.save(dynamic_mask_path)
+
+ def save_objects(self, scene_data: Sequence, num_frames):
+ """Parse and save the objects annotation data.
+
+ Args:
+ scene_data (:obj:`Sequence`): PandaSet sequence.
+ num_frames (int): Number of frames.
+ """
+ instances_info = {}
+
+ for frame_idx in range(num_frames):
+ cuboids = scene_data.cuboids[frame_idx]
+ for _, row in cuboids.iterrows():
+ str_id = row["uuid"]
+ label = row["label"]
+ if label not in PANDA_DYNAMIC_CLASSES:
+ continue
+
+ if str_id not in instances_info:
+ instances_info[str_id] = dict(
+ id=str_id,
+ class_name=row["label"],
+ sibling_id=row["cuboids.sibling_id"],
+ frame_annotations={
+ "frame_idx": [],
+ "obj_to_world": [],
+ "box_size": [],
+ "stationary": [],
+ }
+ )
+
+ # Box coordinates in vehicle frame.
+ tx, ty, tz = row["position.x"], row["position.y"], row["position.z"]
+
+ # The heading of the bounding box (in radians). The heading is the angle
+ # required to rotate +x to the surface normal of the box front face. It is
+ # normalized to [-pi, pi).
+ c = np.math.cos(row["yaw"])
+ s = np.math.sin(row["yaw"])
+
+ # [object to world] transformation matrix
+ o2w = np.array([
+ [ c, -s, 0, tx],
+ [ s, c, 0, ty],
+ [ 0, 0, 1, tz],
+ [ 0, 0, 0, 1]])
+
+ # Dimensions of the box. length: dim x. width: dim y. height: dim z.
+ # length: dim_x: along heading; dim_y: verticle to heading; dim_z: verticle up
+ dimension = [row["dimensions.x"], row["dimensions.y"], row["dimensions.z"]]
+
+ instances_info[str_id]['frame_annotations']['frame_idx'].append(frame_idx)
+ instances_info[str_id]['frame_annotations']['obj_to_world'].append(o2w.tolist())
+ instances_info[str_id]['frame_annotations']['box_size'].append(dimension)
+ instances_info[str_id]['frame_annotations']['stationary'].append(row["stationary"])
+
+ return instances_info
+
+ def create_folder(self):
+ """Create folder for data preprocessing."""
+ if self.process_id_list is None:
+ id_list = range(len(self))
+ else:
+ id_list = self.process_id_list
+ for i in id_list:
+ if "images" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/images", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/sky_masks", exist_ok=True)
+ if "calib" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/extrinsics", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/intrinsics", exist_ok=True)
+ if "pose" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/ego_pose", exist_ok=True)
+ if "lidar" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/lidar", exist_ok=True)
+ if "3dbox_vis" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/3dbox_vis", exist_ok=True)
+ if "dynamic_masks" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/dynamic_masks", exist_ok=True)
+ if "objects" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/instances", exist_ok=True)
+ if "objects_vis" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/instances/debug_vis", exist_ok=True)
diff --git a/worldbench/third_party/drivestudio/datasets/pandaset/pandaset_preview.py b/worldbench/third_party/drivestudio/datasets/pandaset/pandaset_preview.py
new file mode 100644
index 0000000..006a0e7
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/pandaset/pandaset_preview.py
@@ -0,0 +1,219 @@
+import os
+import pickle
+
+import cv2
+import numpy as np
+import pandas as pd
+from moviepy.editor import VideoFileClip, clips_array
+from tqdm import tqdm
+
+def count_labels(
+ base_path: str = 'data/pandaset/raw'
+):
+ combined_label_counts = pd.Series(dtype=int)
+
+ # Traverse all scene directories
+ for scene_dir in os.listdir(base_path):
+ scene_path = os.path.join(base_path, scene_dir, 'annotations', 'cuboids')
+ if os.path.isdir(scene_path):
+ # Traverse all PKL files in each scene
+ for pkl_file in os.listdir(scene_path):
+ if pkl_file.endswith('.pkl'):
+ pkl_file_path = os.path.join(scene_path, pkl_file)
+ # Read the PKL file
+ with open(pkl_file_path, 'rb') as f:
+ data = pickle.load(f)
+ # Count labels in the current PKL file
+ label_counts = data['label'].value_counts()
+ # Merge current label counts into the total counts
+ combined_label_counts = combined_label_counts.add(label_counts, fill_value=0)
+
+ total_label_counts_df = combined_label_counts.reset_index()
+ total_label_counts_df.columns = ['Label', 'Count']
+
+ print(total_label_counts_df)
+
+ unique_labels = total_label_counts_df['Label'].tolist()
+
+ return unique_labels
+
+def preview_camera_captures(
+ input_root: str = 'data/pandaset/processed',
+ output_root: str = 'data/pandaset/preview/camera_captures',
+ frame_width: int = 960 * 5, # Concatenate five camera images horizontally
+ frame_height: int = 540,
+ fps: int = 16
+):
+ if not os.path.exists(output_root):
+ os.makedirs(output_root)
+
+ # Iterate over each scene directory in the input root directory
+ for scene_id in tqdm(sorted(os.listdir(input_root))):
+ scene_folder = os.path.join(input_root, scene_id, "images")
+
+ # Define the output video file path
+ output_video_path = os.path.join(output_root, f'{scene_id}.mp4')
+
+ # Get paths to each camera image folder
+ images_all = sorted([f for f in sorted(os.listdir(scene_folder)) if f.lower().endswith('.jpg')])
+ num_images = len(images_all)
+ num_frames = num_images // 6
+
+ out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
+ for frame_id in range(num_frames):
+ front = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6]))
+ front_left = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6 + 1]))
+ front_right = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6 + 2]))
+ left = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6 + 3]))
+ right = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6 + 4]))
+
+ combined_img = np.hstack((left, front_left, front, front_right, right))
+ combined_img = cv2.resize(combined_img, (frame_width, frame_height))
+
+ out.write(combined_img)
+ out.release()
+
+def preview_3dbox_projection(
+ input_root: str = 'data/pandaset/processed',
+ output_root: str = 'data/pandaset/preview/3dbox_vis',
+ frame_width: int = 960 * 5,
+ frame_height: int = 540,
+ fps: int = 24
+):
+ if not os.path.exists(output_root):
+ os.makedirs(output_root)
+
+ # Iterate over each scene directory in the input root directory
+ for scene_id in tqdm(sorted(os.listdir(input_root))):
+ scene_folder = os.path.join(input_root, scene_id, "3dbox_vis")
+
+ # Define the output video file path
+ output_video_path = os.path.join(output_root, f'{scene_id}.mp4')
+
+ # Get paths to each camera image folder
+ images_all = sorted([f for f in sorted(os.listdir(scene_folder)) if f.lower().endswith('.jpg')])
+ num_images = len(images_all)
+ num_frames = num_images // 6
+
+ out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
+ for frame_id in range(num_frames):
+ front = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6]))
+ front_left = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6 + 1]))
+ front_right = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6 + 2]))
+ left = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6 + 3]))
+ right = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6 + 4]))
+
+ combined_img = np.hstack((left, front_left, front, front_right, right))
+ combined_img = cv2.resize(combined_img, (frame_width, frame_height))
+
+ out.write(combined_img)
+ out.release()
+
+def preview_dynamic_masks(
+ input_root: str = 'data/pandaset/processed',
+ output_root: str = 'data/pandaset/preview/dynamic_masks',
+ frame_width: int = 960 * 5,
+ frame_height: int = 540,
+ fps: int = 24,
+ dynamic_class: str = 'all' # 'all' or 'vehicle' or 'human'
+):
+ if not os.path.exists(output_root):
+ os.makedirs(output_root)
+
+ # Iterate over each scene directory in the input root directory
+ for scene_id in tqdm(sorted(os.listdir(input_root))):
+ scene_folder = os.path.join(input_root, scene_id, "dynamic_masks", dynamic_class)
+
+ # Define the output video file path
+ output_video_path = os.path.join(output_root, f'{scene_id}.mp4')
+
+ # Get paths to each camera image folder
+ images_all = sorted([f for f in sorted(os.listdir(scene_folder)) if f.lower().endswith('.png')])
+ num_images = len(images_all)
+ num_frames = num_images // 6
+
+ out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
+ for frame_id in range(num_frames):
+ front = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6]))
+ front_left = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6 + 1]))
+ front_right = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6 + 2]))
+ left = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6 + 3]))
+ right = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6 + 4]))
+
+ combined_img = np.hstack((left, front_left, front, front_right, right))
+ combined_img = cv2.resize(combined_img, (frame_width, frame_height))
+
+ out.write(combined_img)
+ out.release()
+
+def preview_instances(
+ input_root: str = 'data/pandaset/processed',
+ output_root: str = 'data/pandaset/preview/instances',
+ frame_width: int = 960 * 5,
+ frame_height: int = 540,
+ fps: int = 24,
+):
+ if not os.path.exists(output_root):
+ os.makedirs(output_root)
+
+ # Iterate over each scene directory in the input root directory
+ for scene_id in tqdm(sorted(os.listdir(input_root))):
+ scene_folder = os.path.join(input_root, scene_id, "instances", "debug_vis")
+
+ # Define the output video file path
+ output_video_path = os.path.join(output_root, f'{scene_id}.mp4')
+
+ # Get paths to each camera image folder
+ images_all = sorted([f for f in sorted(os.listdir(scene_folder)) if f.lower().endswith('.jpg')])
+ num_images = len(images_all)
+ num_frames = num_images // 6
+
+ out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
+ for frame_id in range(num_frames):
+ front = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6]))
+ front_left = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6 + 1]))
+ front_right = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6 + 2]))
+ left = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6 + 3]))
+ right = cv2.imread(os.path.join(scene_folder, images_all[frame_id * 6 + 4]))
+
+ combined_img = np.hstack((left, front_left, front, front_right, right))
+ combined_img = cv2.resize(combined_img, (frame_width, frame_height))
+
+ out.write(combined_img)
+ out.release()
+
+def cat_all_videos(
+ keys: list,
+ output_root: str = 'data/pandaset/preview/catted_videos',
+ scene_id: str = 'all',
+ fps: int = 24
+):
+ if not os.path.exists(output_root):
+ os.makedirs(output_root)
+
+ if scene_id == 'all':
+ scene_ids = sorted(os.listdir('data/pandaset/processed'))
+ else:
+ scene_ids = [scene_id]
+
+ for scene_id in scene_ids:
+ video_paths = []
+ for key in keys:
+ video_paths.append(f'data/pandaset/preview/{key}/{scene_id}.mp4')
+
+ output_video_path = os.path.join(output_root, f'{scene_id}.mp4')
+
+ # use moviepy to concatenate videos
+ clips = [VideoFileClip(path) for path in video_paths]
+ # cat videos from top to bottom, vertically
+ final_clip = clips_array([[clip] for clip in clips])
+ final_clip.write_videofile(output_video_path, fps=fps)
+ final_clip.close()
+
+if __name__ == '__main__':
+ # unique_labels = count_labels()
+ preview_camera_captures(fps=16)
+ preview_3dbox_projection(fps=16)
+ preview_dynamic_masks(fps=16)
+ preview_instances(fps=16)
+ cat_all_videos(keys=['camera_captures', '3dbox_vis', 'dynamic_masks', 'instances'], scene_id='all', fps=16)
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/pandaset/pandaset_sourceloader.py b/worldbench/third_party/drivestudio/datasets/pandaset/pandaset_sourceloader.py
new file mode 100644
index 0000000..d113734
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/pandaset/pandaset_sourceloader.py
@@ -0,0 +1,504 @@
+import json
+import logging
+import os
+from typing import Dict
+
+import joblib
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from pytorch3d.transforms import matrix_to_quaternion
+from torch import Tensor
+from tqdm import tqdm, trange
+
+from datasets.base.scene_dataset import ModelType
+from datasets.base.lidar_source import SceneLidarSource
+from datasets.base.pixel_source import CameraData, ScenePixelSource
+
+logger = logging.getLogger()
+
+# define each class's node type
+OBJECT_CLASS_NODE_MAPPING = {
+ # Rigid objects (vehicles)
+ "Bus": ModelType.RigidNodes,
+ "Car": ModelType.RigidNodes,
+ "Emergency Vehicle": ModelType.RigidNodes,
+ "Medium-sized Truck": ModelType.RigidNodes,
+ "Motorcycle": ModelType.RigidNodes,
+ "Motorized Scooter": ModelType.RigidNodes,
+ "Other Vehicle - Construction Vehicle": ModelType.RigidNodes,
+ "Other Vehicle - Pedicab": ModelType.RigidNodes,
+ "Other Vehicle - Uncommon": ModelType.RigidNodes,
+ "Pickup Truck": ModelType.RigidNodes,
+ "Semi-truck": ModelType.RigidNodes,
+ "Train": ModelType.RigidNodes,
+ "Tram / Subway": ModelType.RigidNodes,
+
+ # Humans (SMPL model)
+ "Pedestrian": ModelType.SMPLNodes,
+ "Pedestrian with Object": ModelType.SMPLNodes,
+
+ # Potentially deformable objects
+ "Bicycle": ModelType.DeformableNodes,
+ "Animals - Other": ModelType.DeformableNodes,
+ "Personal Mobility Device": ModelType.DeformableNodes,
+}
+SMPLNODE_CLASSES = ["Pedestrian", "Pedestrian with Object"]
+
+# OpenCV to Dataset coordinate transformation
+# opencv coordinate system: x right, y down, z front
+OPENCV2DATASET = np.array(
+ [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]
+)
+
+# PandaSet Camera List:
+# 0: front_camera
+# 1: front_left_camera
+# 2: front_right_camera
+# 3: left_camera
+# 4: right_camera
+# 5: back_camera
+AVAILABLE_CAM_LIST = [0, 1, 2, 3, 4, 5]
+
+class PandaCameraData(CameraData):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ def load_calibrations(self):
+ """
+ Load the camera intrinsics, extrinsics, timestamps, etc.
+ Compute the camera-to-world matrices, ego-to-world matrices, etc.
+ """
+ # to store per-camera intrinsics and extrinsics
+ ego_to_world_start = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{self.start_timestep:03d}.txt")
+ )
+ # compute per-image poses and intrinsics
+ cam_to_worlds, intrinsics, distortions = [], [], []
+
+ for t in range(self.start_timestep, self.end_timestep):
+ # load camera intrinsics
+ # 1d Array of [f_u, f_v, c_u, c_v, k{1, 2}, p{1, 2}, k{3}].
+ # ====!! we did not use distortion parameters for simplicity !!====
+ # to be improved!!
+ intrinsic = np.loadtxt(
+ os.path.join(self.data_path, "intrinsics", f"{self.cam_id}.txt")
+ )
+ fx, fy, cx, cy = intrinsic[0], intrinsic[1], intrinsic[2], intrinsic[3]
+ k1, k2, p1, p2, k3 = intrinsic[4], intrinsic[5], intrinsic[6], intrinsic[7], intrinsic[8]
+ # scale intrinsics w.r.t. load size
+ fx, fy = (
+ fx * self.load_size[1] / self.original_size[1],
+ fy * self.load_size[0] / self.original_size[0],
+ )
+ cx, cy = (
+ cx * self.load_size[1] / self.original_size[1],
+ cy * self.load_size[0] / self.original_size[0],
+ )
+ intrinsic = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
+ intrinsics.append(intrinsic)
+ distortions.append(np.array([k1, k2, p1, p2, k3]))
+
+ # load camera extrinsics
+ cam2world = np.loadtxt(
+ os.path.join(self.data_path, "extrinsics", f"{t:03d}_{self.cam_id}.txt")
+ )
+ # we transform the camera poses w.r.t. the first timestep to make the translation vector of
+ # the first ego pose as the origin of the world coordinate system.
+ cam2world = np.linalg.inv(ego_to_world_start) @ cam2world
+ cam2world = cam2world @ OPENCV2DATASET
+ cam_to_worlds.append(cam2world)
+
+ self.intrinsics = torch.from_numpy(np.stack(intrinsics, axis=0)).float()
+ self.distortions = torch.from_numpy(np.stack(distortions, axis=0)).float()
+ self.cam_to_worlds = torch.from_numpy(np.stack(cam_to_worlds, axis=0)).float()
+
+ @classmethod
+ def get_camera2worlds(cls, data_path: str, cam_id: str, start_timestep: int, end_timestep: int) -> torch.Tensor:
+ """
+ Returns camera-to-world matrices for the specified camera and time range.
+
+ Args:
+ data_path (str): Path to the dataset.
+ cam_id (str): Camera ID.
+ start_timestep (int): Start timestep.
+ end_timestep (int): End timestep.
+
+ Returns:
+ torch.Tensor: Camera-to-world matrices of shape (num_frames, 4, 4).
+ """
+ # Load the ego pose for the start timestep to use as reference
+ ego_to_world_start = np.loadtxt(
+ os.path.join(data_path, "ego_pose", f"{start_timestep:03d}.txt")
+ )
+
+ cam_to_worlds = []
+
+ for t in range(start_timestep, end_timestep):
+ # Load camera extrinsics for each timestep
+ cam2world = np.loadtxt(
+ os.path.join(data_path, "extrinsics", f"{t:03d}_{cam_id}.txt")
+ )
+
+ # Transform the camera poses w.r.t. the first timestep
+ cam2world = np.linalg.inv(ego_to_world_start) @ cam2world
+
+ # Apply OPENCV2DATASET transformation
+ cam2world = cam2world @ OPENCV2DATASET
+
+ cam_to_worlds.append(cam2world)
+
+ return torch.from_numpy(np.stack(cam_to_worlds, axis=0)).float()
+
+class PandaPixelSource(ScenePixelSource):
+ def __init__(
+ self,
+ dataset_name: str,
+ pixel_data_config: OmegaConf,
+ data_path: str,
+ start_timestep: int,
+ end_timestep: int,
+ device: torch.device = torch.device("cpu"),
+ ):
+ super().__init__(dataset_name, pixel_data_config, device=device)
+ self.data_path = data_path
+ self.start_timestep = start_timestep
+ self.end_timestep = end_timestep
+ self.load_data()
+
+ def load_cameras(self):
+ self._timesteps = torch.arange(self.start_timestep, self.end_timestep)
+ self.register_normalized_timestamps()
+
+ for idx, cam_id in enumerate(self.camera_list):
+ logger.info(f"Loading camera {cam_id}")
+ camera = PandaCameraData(
+ dataset_name=self.dataset_name,
+ data_path=self.data_path,
+ cam_id=cam_id,
+ start_timestep=self.start_timestep,
+ end_timestep=self.end_timestep,
+ load_dynamic_mask=self.data_cfg.load_dynamic_mask,
+ load_sky_mask=self.data_cfg.load_sky_mask,
+ downscale_when_loading=self.data_cfg.downscale_when_loading[idx],
+ undistort=self.data_cfg.undistort,
+ buffer_downscale=self.buffer_downscale,
+ device=self.device,
+ )
+ camera.load_time(self.normalized_time)
+ unique_img_idx = torch.arange(len(camera), device=self.device) * len(self.camera_list) + idx
+ camera.set_unique_ids(
+ unique_cam_idx = idx,
+ unique_img_idx = unique_img_idx
+ )
+ logger.info(f"Camera {camera.cam_name} loaded.")
+ self.camera_data[cam_id] = camera
+
+ def load_objects(self):
+ """
+ get ground truth bounding boxes of the dynamic objects
+
+ instances_info = {
+ "0": # simplified instance id
+ {
+ "id": str,
+ "class_name": str,
+ "frame_annotations": {
+ "frame_idx": List,
+ "obj_to_world": List,
+ "box_size": List,
+ },
+ ...
+ }
+ frame_instances = {
+ "0": # frame idx
+ List[int] # list of simplified instance ids
+ ...
+ }
+ """
+ instances_info_path = os.path.join(self.data_path, "instances", "instances_info.json")
+ frame_instances_path = os.path.join(self.data_path, "instances", "frame_instances.json")
+ with open(instances_info_path, "r") as f:
+ instances_info = json.load(f)
+ with open(frame_instances_path, "r") as f:
+ frame_instances = json.load(f)
+ # get pose of each instance at each frame
+ # shape (num_frames, num_instances, 4, 4)
+ num_instances = len(instances_info)
+ num_full_frames = len(frame_instances)
+ instances_pose = np.zeros((num_full_frames, num_instances, 4, 4))
+ instances_size = np.zeros((num_full_frames, num_instances, 3))
+ instances_true_id = np.arange(num_instances)
+ instances_model_types = np.ones(num_instances) * -1
+
+ ego_to_world_start = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{self.start_timestep:03d}.txt")
+ )
+ for k, v in instances_info.items():
+ instances_model_types[int(k)] = OBJECT_CLASS_NODE_MAPPING[v["class_name"]]
+ for frame_idx, obj_to_world, box_size in zip(v["frame_annotations"]["frame_idx"], v["frame_annotations"]["obj_to_world"], v["frame_annotations"]["box_size"]):
+ # the first ego pose as the origin of the world coordinate system.
+ obj_to_world = np.array(obj_to_world).reshape(4, 4)
+ obj_to_world = np.linalg.inv(ego_to_world_start) @ obj_to_world
+ instances_pose[frame_idx, int(k)] = np.array(obj_to_world)
+ instances_size[frame_idx, int(k)] = np.array(box_size)
+
+ # get frame valid instances
+ # shape (num_frames, num_instances)
+ per_frame_instance_mask = np.zeros((num_full_frames, num_instances))
+ for frame_idx, valid_instances in frame_instances.items():
+ per_frame_instance_mask[int(frame_idx), valid_instances] = 1
+
+ # select the frames that are in the range of start_timestep and end_timestep
+ instances_pose = torch.from_numpy(instances_pose[self.start_timestep:self.end_timestep]).float()
+ instances_size = torch.from_numpy(instances_size[self.start_timestep:self.end_timestep]).float()
+ instances_true_id = torch.from_numpy(instances_true_id).long()
+ instances_model_types = torch.from_numpy(instances_model_types).long()
+ per_frame_instance_mask = torch.from_numpy(per_frame_instance_mask[self.start_timestep:self.end_timestep]).bool()
+
+ # filter out the instances that are not visible in selected frames
+ ins_frame_cnt = per_frame_instance_mask.sum(dim=0)
+ instances_pose = instances_pose[:, ins_frame_cnt > 0]
+ instances_size = instances_size[:, ins_frame_cnt > 0]
+ instances_true_id = instances_true_id[ins_frame_cnt > 0]
+ instances_model_types = instances_model_types[ins_frame_cnt > 0]
+ per_frame_instance_mask = per_frame_instance_mask[:, ins_frame_cnt > 0]
+
+ # assign to the class
+ # (num_frames, num_instances, 4, 4)
+ self.instances_pose = instances_pose
+ # (num_instances, 3)
+ self.instances_size = instances_size.sum(0) / per_frame_instance_mask.sum(0).unsqueeze(-1)
+ # (num_frames, num_instances)
+ self.per_frame_instance_mask = per_frame_instance_mask
+ # (num_instances)
+ self.instances_true_id = instances_true_id
+ # (num_instances)
+ self.instances_model_types = instances_model_types
+
+ if self.data_cfg.load_smpl:
+ # Collect camera-to-world matrices for all available cameras
+ cam_to_worlds = {}
+ for cam_id in AVAILABLE_CAM_LIST:
+ cam_to_worlds[cam_id] = PandaCameraData.get_camera2worlds(
+ self.data_path,
+ str(cam_id),
+ self.start_timestep,
+ self.end_timestep
+ )
+
+ # load SMPL parameters
+ smpl_dict = joblib.load(os.path.join(self.data_path, "humanpose", "smpl.pkl"))
+ frame_num = self.end_timestep - self.start_timestep
+
+ smpl_human_all = {}
+ for fi in tqdm(range(self.start_timestep, self.end_timestep), desc="Loading SMPL"):
+ for instance_id, ins_smpl in smpl_dict.items():
+ if instance_id not in smpl_human_all:
+ smpl_human_all[instance_id] = {
+ "smpl_quats": torch.zeros((frame_num, 24, 4), dtype=torch.float32),
+ "smpl_trans": torch.zeros((frame_num, 3), dtype=torch.float32),
+ "smpl_betas": torch.zeros((frame_num, 10), dtype=torch.float32),
+ "frame_valid": torch.zeros((frame_num), dtype=torch.bool)
+ }
+ smpl_human_all[instance_id]["smpl_quats"][:, :, 0] = 1.0
+ if ins_smpl["valid_mask"][fi]:
+ betas = ins_smpl["smpl"]["betas"][fi]
+ smpl_human_all[instance_id]["smpl_betas"][fi - self.start_timestep] = betas
+
+ body_pose = ins_smpl["smpl"]["body_pose"][fi]
+ smpl_orient = ins_smpl["smpl"]["global_orient"][fi]
+ cam_depend = ins_smpl["selected_cam_idx"][fi].item()
+
+ c2w = cam_to_worlds[cam_depend][fi - self.start_timestep]
+ world_orient = c2w[:3, :3].to(smpl_orient.device) @ smpl_orient.squeeze()
+ smpl_quats = matrix_to_quaternion(
+ torch.cat([world_orient[None, ...], body_pose], dim=0)
+ )
+
+ ii = instances_info[str(instance_id)]['frame_annotations']["frame_idx"].index(fi)
+ o2w = np.array(
+ instances_info[str(instance_id)]['frame_annotations']["obj_to_world"][ii]
+ )
+ o2w = torch.from_numpy(
+ np.linalg.inv(ego_to_world_start) @ o2w
+ )
+ # box_size = instances_info[str(instance_id)]['frame_annotations']["box_size"][ii]
+
+ smpl_human_all[instance_id]["smpl_quats"][fi - self.start_timestep] = smpl_quats
+ smpl_human_all[instance_id]["smpl_trans"][fi - self.start_timestep] = o2w[:3, 3]
+ smpl_human_all[instance_id]["frame_valid"][fi - self.start_timestep] = True
+
+ self.smpl_human_all = smpl_human_all
+
+class PandaLiDARSource(SceneLidarSource):
+ def __init__(
+ self,
+ lidar_data_config: OmegaConf,
+ data_path: str,
+ start_timestep: int,
+ end_timestep: int,
+ device: torch.device = torch.device("cpu"),
+ ):
+ super().__init__(lidar_data_config, device=device)
+ self.data_path = data_path
+ self.start_timestep = start_timestep
+ self.end_timestep = end_timestep
+ self.create_all_filelist()
+ self.load_data()
+
+ def create_all_filelist(self):
+ """
+ Create a list of all the files in the dataset.
+ e.g., a list of all the lidar scans in the dataset.
+ """
+ lidar_filepaths = []
+ for t in range(self.start_timestep, self.end_timestep):
+ lidar_filepaths.append(
+ os.path.join(self.data_path, "lidar", f"{t:03d}.bin")
+ )
+ self.lidar_filepaths = np.array(lidar_filepaths)
+
+ def load_calibrations(self):
+ """
+ Load the calibration files of the dataset.
+ e.g., lidar to world transformation matrices.
+ """
+ # Note that in the Waymo Open Dataset, the lidar coordinate system is the same
+ # as the vehicle coordinate system
+ lidar_to_worlds = []
+
+ # we tranform the poses w.r.t. the first timestep to make the origin of the
+ # first ego pose as the origin of the world coordinate system.
+ ego_to_world_start = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{self.start_timestep:03d}.txt")
+ )
+ for t in range(self.start_timestep, self.end_timestep):
+ ego_to_world_current = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{t:03d}.txt")
+ )
+ # compute ego_to_world transformation
+ lidar_to_world = np.linalg.inv(ego_to_world_start) @ ego_to_world_current
+ lidar_to_worlds.append(lidar_to_world)
+
+ self.lidar_to_worlds = torch.from_numpy(
+ np.stack(lidar_to_worlds, axis=0)
+ ).float()
+
+ def load_lidar(self):
+ """
+ Load the lidar data of the dataset from the filelist.
+ """
+ origins, directions, ranges, laser_ids = [], [], [], []
+ # we simplify timestamps as the time indices
+ timesteps = []
+
+ accumulated_num_original_rays = 0
+ accumulated_num_rays = 0
+ for t in trange(
+ 0, len(self.lidar_filepaths), desc="Loading lidar", dynamic_ncols=True
+ ):
+ # each lidar_info contains an Nx5 array
+ # from left to right:
+ # points: 3d, intensity: 1d, laser_id: 1d
+ lidar_info = np.memmap(
+ self.lidar_filepaths[t],
+ dtype=np.float32,
+ mode="r",
+ ).reshape(-1, 5)
+ original_length = len(lidar_info)
+ accumulated_num_original_rays += original_length
+
+ # select lidar points based on the laser id
+ if self.data_cfg.only_use_360_lidar:
+ # laser_ids: 0: mechanical 360° LiDAR, 1: front-facing LiDAR
+ lidar_info = lidar_info[lidar_info[:, 4] == 0]
+
+ lidar_points = torch.from_numpy(lidar_info[:, :3]).float()
+ # convert the lidar points to opencv coordinate system
+ # since lidar points already include the information of origins
+ # NOTE: here we use the pesudo origins, TODO: consider a more accurate way
+ lidar_origins = torch.zeros_like(lidar_points)
+ lidar_ids = torch.from_numpy(lidar_info[:, 4]).float()
+
+ # transform lidar points from lidar coordinate system to world coordinate system
+ lidar_origins = (
+ self.lidar_to_worlds[t][:3, :3] @ lidar_origins.T
+ + self.lidar_to_worlds[t][:3, 3:4]
+ ).T
+ lidar_points = (
+ self.lidar_to_worlds[t][:3, :3] @ lidar_points.T
+ + self.lidar_to_worlds[t][:3, 3:4]
+ ).T
+ # compute lidar directions
+ lidar_directions = lidar_points - lidar_origins
+ lidar_ranges = torch.norm(lidar_directions, dim=-1, keepdim=True)
+ lidar_directions = lidar_directions / lidar_ranges
+ # we use time indices as the timestamp
+ lidar_timestamp = torch.ones_like(lidar_ranges).squeeze(-1) * t
+ accumulated_num_rays += len(lidar_ranges)
+
+ origins.append(lidar_origins)
+ directions.append(lidar_directions)
+ ranges.append(lidar_ranges)
+ laser_ids.append(lidar_ids)
+ # we use time indices as the timestamp
+ timesteps.append(lidar_timestamp)
+
+ logger.info(
+ f"Number of lidar rays: {accumulated_num_rays} "
+ f"({accumulated_num_rays / accumulated_num_original_rays * 100:.2f}% of "
+ f"{accumulated_num_original_rays} original rays)"
+ )
+ logger.info("Filter condition:")
+ logger.info(f" only_use_360_lidar: {self.data_cfg.only_use_360_lidar}")
+
+ self.origins = torch.cat(origins, dim=0)
+ self.directions = torch.cat(directions, dim=0)
+ self.ranges = torch.cat(ranges, dim=0)
+ self.laser_ids = torch.cat(laser_ids, dim=0)
+ self.visible_masks = torch.zeros_like(self.ranges).squeeze().bool()
+ self.colors = torch.ones_like(self.directions)
+
+ # the underscore here is important.
+ self._timesteps = torch.cat(timesteps, dim=0)
+ self.register_normalized_timestamps()
+
+ def get_lidar_rays(self, time_idx: int) -> Dict[str, Tensor]:
+ """
+ Get the of rays for rendering at the given timestep.
+ Args:
+ time_idx: the index of the lidar scan to render.
+ Returns:
+ a dict of the sampled rays.
+ """
+ origins = self.origins[self.timesteps == time_idx]
+ directions = self.directions[self.timesteps == time_idx]
+ ranges = self.ranges[self.timesteps == time_idx]
+ normalized_time = self.normalized_time[self.timesteps == time_idx]
+ return {
+ "lidar_origins": origins,
+ "lidar_viewdirs": directions,
+ "lidar_ranges": ranges,
+ "lidar_normed_time": normalized_time,
+ "lidar_mask": self.timesteps == time_idx,
+ }
+
+ def delete_invisible_pts(self) -> None:
+ """
+ Clear the unvisible points.
+ """
+ if self.visible_masks is not None:
+ num_bf = self.origins.shape[0]
+ self.origins = self.origins[self.visible_masks]
+ self.directions = self.directions[self.visible_masks]
+ self.ranges = self.ranges[self.visible_masks]
+ self._timesteps = self._timesteps[self.visible_masks]
+ self._normalized_time = self._normalized_time[self.visible_masks]
+ self.colors = self.colors[self.visible_masks]
+ logger.info(
+ f"[Lidar] {num_bf - self.visible_masks.sum()} out of {num_bf} points are cleared. {self.visible_masks.sum()} points left."
+ )
+ self.visible_masks = None
+ else:
+ logger.info("[Lidar] No unvisible points to clear.")
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/preprocess.py b/worldbench/third_party/drivestudio/datasets/preprocess.py
new file mode 100644
index 0000000..595e3a0
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/preprocess.py
@@ -0,0 +1,274 @@
+import argparse
+import numpy as np
+
+if __name__ == "__main__":
+ """
+ Unified Dataset preprocessing script
+ ===========================
+
+ This script facilitates the preprocessing of datasets:
+ - Waymo
+ - Argoverse
+ - NuScenes
+ - KITTI
+ - NUPlan
+ - PandaSet
+
+ Usage:
+ ------
+ python datasets/preprocess.py \
+ --data_root \
+ --dataset \
+ --split \
+ --target_dir \
+ --workers \
+ --scene_ids \
+ --split_file \
+ --start_idx \
+ --num_scenes \
+ --process_keys
+
+ Example:
+ --------
+ Waymo:
+ python datasets/preprocess.py \
+ --data_root data/waymo/raw/ \
+ --target_dir data/waymo/processed \
+ --split training \
+ --process_keys images lidar calib pose dynamic_masks objects \
+ --workers 8 \
+ --scene_ids 23 114 327 621 703 172 552 788
+
+ PandaSet:
+ python datasets/preprocess.py \
+ --data_root data/pandaset/raw \
+ --dataset pandaset \
+ --split_file data/pandaset_example_scenes.txt \
+ --target_dir data/pandaset/processed \
+ --workers 32 \
+ --process_keys images lidar calib pose dynamic_masks objects
+
+ Please refer to the documentation for more information on the available options.
+
+ Arguments:
+ ----------
+ --data_root (str):
+ The root directory where the Waymo dataset is stored. This is a required argument.
+
+ --split (str):
+ Specifies the name of the data split. Default is set to "training".
+
+ --target_dir (str):
+ Designates the directory where the processed data will be saved. This is a mandatory argument.
+
+ --workers (int):
+ The number of processing threads. Default is set to 4.
+
+ --scene_ids (list[int]):
+ List of specific scene IDs for processing. Should be integers separated by spaces.
+
+ --split_file (str):
+ If provided, indicates the path to a file located in `data/waymo_splits` that contains the desired scene IDs.
+
+ --start_idx (int):
+ Used in conjunction with `num_scenes` to generate a list of scene IDs when neither `scene_ids` nor `split_file` are provided.
+
+ --num_scenes (int):
+ The total number of scenes to be processed.
+
+ --process_keys (list[str]):
+ Denotes the types of data components to be processed. Options include but aren't limited to "images", "lidar", "calib", "pose", etc.
+
+ Notes:
+ ------
+ The logic of the script ensures that if specific scene IDs (`scene_ids`) are provided, they are prioritized.
+ If a split file (`split_file`) is indicated, it is utilized next.
+ If neither is available, the script uses the `start_idx` and `num_scenes` parameters to determine the scene IDs.
+ """
+ parser = argparse.ArgumentParser(description="Data converter arg parser")
+ parser.add_argument(
+ "--data_root", type=str, required=True, help="root path of dataset"
+ )
+ parser.add_argument("--dataset", type=str, default="waymo", help="dataset name")
+ parser.add_argument(
+ "--split",
+ type=str,
+ default="training",
+ help="split of the dataset, e.g. training, validation, testing, please specify the split name for different dataset",
+ )
+ parser.add_argument(
+ "--target_dir",
+ type=str,
+ required=True,
+ help="output directory of processed data",
+ )
+ parser.add_argument(
+ "--workers", type=int, default=4, help="number of threads to be used"
+ )
+ # priority: scene_ids > split_file > start_idx + num_scenes
+ parser.add_argument(
+ "--scene_ids",
+ default=None,
+ type=int,
+ nargs="+",
+ help="scene ids to be processed, a list of integers separated by space. Range: [0, 798] for training, [0, 202] for validation",
+ )
+ parser.add_argument(
+ "--split_file", type=str, default=None, help="Split file in data/waymo_splits"
+ )
+ parser.add_argument(
+ "--start_idx",
+ type=int,
+ default=0,
+ help="If no scene id or split_file is given, use start_idx and num_scenes to generate scene_ids_list",
+ )
+ parser.add_argument(
+ "--num_scenes",
+ type=int,
+ default=200,
+ help="number of scenes to be processed",
+ )
+ parser.add_argument(
+ "--max_frame_limit",
+ type=int,
+ default=300,
+ help="maximum number of frames to be processed in a dataset, in nuplan dataset, \
+ the scene duration super long, we can limit the number of frames to be processed, \
+ this argument is used only for nuplan dataset",
+ )
+ parser.add_argument(
+ "--start_frame_idx",
+ type=int,
+ default=1000,
+ help="We skip the first start_frame_idx frames to avoid ego static frames",
+ )
+ parser.add_argument(
+ "--interpolate_N",
+ type=int,
+ default=0,
+ help="Interpolate to get frames at higher frequency, this is only used for nuscene dataset",
+ )
+ parser.add_argument(
+ "--data_source",
+ type=str,
+ default="gt"
+ )
+ parser.add_argument(
+ "--process_keys",
+ nargs="+",
+ default=[
+ "images",
+ "lidar",
+ "calib",
+ "pose",
+ "dynamic_masks",
+ "objects"
+ ],
+ )
+ args = parser.parse_args()
+ if (args.dataset != 'nuscenes' and args.dataset != 'nuscenes_gen') and args.interpolate_N > 0:
+ parser.error("interpolate_N > 0 is only allowed when dataset is 'nuscenes'")
+
+ if args.scene_ids is not None:
+ scene_ids_list = args.scene_ids
+ elif args.split_file is not None:
+ # parse the split file
+ split_file = open(args.split_file, "r").readlines()[1:]
+ scene_ids_list = [line.strip().split(",")[0] for line in split_file]
+ else:
+ scene_ids_list = np.arange(args.start_idx, args.start_idx + args.num_scenes)
+
+ if args.dataset == "waymo":
+ from datasets.waymo.waymo_preprocess import WaymoProcessor
+
+ scene_ids_list = [int(scene_id) for scene_id in scene_ids_list]
+ dataset_processor = WaymoProcessor(
+ load_dir=args.data_root,
+ save_dir=args.target_dir,
+ prefix=args.split,
+ process_keys=args.process_keys,
+ process_id_list=scene_ids_list,
+ workers=args.workers,
+ )
+ elif args.dataset == "pandaset":
+ from datasets.pandaset.pandaset_preprocess import PandaSetProcessor
+
+ scene_ids_list = [str(scene_id).zfill(3) for scene_id in scene_ids_list]
+ dataset_processor = PandaSetProcessor(
+ load_dir=args.data_root,
+ save_dir=args.target_dir,
+ process_keys=args.process_keys,
+ process_id_list=scene_ids_list,
+ workers=args.workers,
+ )
+ elif args.dataset == "argoverse":
+ from datasets.argoverse.argoverse_preprocess import ArgoVerseProcessor
+
+ scene_ids_list = [int(scene_id) for scene_id in scene_ids_list]
+ dataset_processor = ArgoVerseProcessor(
+ load_dir=args.data_root,
+ save_dir=args.target_dir,
+ process_keys=args.process_keys,
+ process_id_list=scene_ids_list,
+ workers=args.workers,
+ )
+ elif args.dataset == "nuscenes":
+ from datasets.nuscenes.nuscenes_preprocess import NuScenesProcessor
+
+ scene_ids_list = [int(scene_id) for scene_id in scene_ids_list]
+ dataset_processor = NuScenesProcessor(
+ load_dir=args.data_root,
+ save_dir=args.target_dir,
+ split=args.split,
+ interpolate_N=args.interpolate_N,
+ process_keys=args.process_keys,
+ process_id_list=scene_ids_list,
+ workers=args.workers,
+ )
+ elif args.dataset == "kitti":
+ from datasets.kitti.kitti_preprocess import KittiProcessor
+
+ dataset_processor = KittiProcessor(
+ load_dir=args.data_root,
+ save_dir=args.target_dir,
+ prefix=args.split,
+ process_keys=args.process_keys,
+ process_id_list=scene_ids_list,
+ workers=args.workers,
+ )
+ elif args.dataset == "nuplan":
+ from datasets.nuplan.nuplan_preprocess import NuPlanProcessor
+
+ dataset_processor = NuPlanProcessor(
+ load_dir=args.data_root,
+ save_dir=args.target_dir,
+ prefix=args.split,
+ start_frame_idx=args.start_frame_idx,
+ max_frame_limit=args.max_frame_limit,
+ process_keys=args.process_keys,
+ process_id_list=scene_ids_list,
+ workers=args.workers,
+ )
+ elif args.dataset == "nuscenes_gen":
+ from datasets.nuscenes_gen.nuscenes_preprocess import NuScenesProcessor
+
+ scene_ids_list = [int(scene_id) for scene_id in scene_ids_list]
+ dataset_processor = NuScenesProcessor(
+ load_dir=args.data_root,
+ save_dir=args.target_dir,
+ split=args.split,
+ interpolate_N=args.interpolate_N,
+ process_keys=args.process_keys,
+ process_id_list=scene_ids_list,
+ workers=args.workers,
+ num_scenes=args.num_scenes,
+ data_source=args.data_source
+ )
+ else:
+ raise ValueError(f"Unknown dataset {args.dataset}, please choose from waymo, pandaset, argoverse, nuscenes, kitti, nuplan")
+
+ if args.scene_ids is not None and args.workers == 1:
+ for scene_id in args.scene_ids:
+ dataset_processor.convert_one(scene_id)
+ else:
+ dataset_processor.convert()
diff --git a/worldbench/third_party/drivestudio/datasets/tools/extract_masks.py b/worldbench/third_party/drivestudio/datasets/tools/extract_masks.py
new file mode 100644
index 0000000..f8c3587
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/tools/extract_masks.py
@@ -0,0 +1,195 @@
+"""
+@file extract_masks.py
+@author Jianfei Guo, Shanghai AI Lab
+@brief Extract semantic mask
+
+Using SegFormer, 2021. Cityscapes 83.2%
+Relies on timm==0.3.2 & pytorch 1.8.1 (buggy on pytorch >= 1.9)
+
+Installation:
+ NOTE: mmcv-full==1.2.7 requires another pytorch version & conda env.
+ Currently mmcv-full==1.2.7 does not support pytorch>=1.9;
+ will raise AttributeError: 'super' object has no attribute '_specify_ddp_gpu_num'
+ Hence, a seperate conda env is needed.
+
+ git clone https://github.com/NVlabs/SegFormer
+
+ conda create -n segformer python=3.8
+ conda activate segformer
+ # conda install pytorch==1.8.1 torchvision==0.9.1 torchaudio==0.8.1 cudatoolkit=11.3 -c pytorch -c conda-forge
+ pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html
+
+ pip install timm==0.3.2 pylint debugpy opencv-python attrs ipython tqdm imageio scikit-image omegaconf
+ pip install mmcv-full==1.2.7 --no-cache-dir
+
+ cd SegFormer
+ pip install .
+
+Usage:
+ Direct run this script in the newly set conda env.
+"""
+
+
+from mmseg.apis import inference_segmentor, init_segmentor, show_result_pyplot
+from mmseg.core.evaluation import get_palette
+
+semantic_classes = [
+ 'road', 'sidewalk', 'building', 'wall', 'fence', 'pole',
+ 'traffic light', 'traffic sign', 'vegetation', 'terrain', 'sky',
+ 'person', 'rider', 'car', 'truck', 'bus', 'train', 'motorcycle',
+ 'bicycle'
+]
+dataset_classes_in_sematic = {
+ 'Vehicle': [13, 14, 15], # 'car', 'truck', 'bus'
+ 'human': [11, 12, 17, 18], # 'person', 'rider', 'motorcycle', 'bicycle'
+}
+
+if __name__ == "__main__":
+ import os
+ import imageio
+ import numpy as np
+ from glob import glob
+ from tqdm import tqdm
+ from argparse import ArgumentParser
+ parser = ArgumentParser()
+ # Custom configs
+ parser.add_argument('--data_root', type=str, default='data/waymo/processed/training')
+ parser.add_argument(
+ "--scene_ids",
+ default=None,
+ type=int,
+ nargs="+",
+ help="scene ids to be processed, a list of integers separated by space. Range: [0, 798] for training, [0, 202] for validation",
+ )
+ parser.add_argument(
+ "--split_file", type=str, default=None, help="Split file in data/waymo_splits"
+ )
+ parser.add_argument(
+ "--start_idx",
+ type=int,
+ default=0,
+ help="If no scene id or split_file is given, use start_idx and num_scenes to generate scene_ids_list",
+ )
+ parser.add_argument(
+ "--num_scenes",
+ type=int,
+ default=200,
+ help="number of scenes to be processed",
+ )
+ parser.add_argument(
+ '--process_dynamic_mask',
+ action='store_true',
+ help="Whether to process dynamic masks",
+ )
+ parser.add_argument('--verbose', action='store_true')
+ parser.add_argument('--ignore_existing', action='store_true')
+ parser.add_argument('--no_compress', action='store_true')
+ parser.add_argument('--rgb_dirname', type=str, default="images")
+ parser.add_argument('--mask_dirname', type=str, default="fine_dynamic_masks")
+
+ # Algorithm configs
+ parser.add_argument('--segformer_path', type=str, default='/home/guojianfei/ai_ws/SegFormer')
+ parser.add_argument('--config', help='Config file', type=str, default=None)
+ parser.add_argument('--checkpoint', help='Checkpoint file', type=str, default=None)
+ parser.add_argument('--device', default='cuda:0', help='Device used for inference')
+ parser.add_argument('--palette', default='cityscapes', help='Color palette used for segmentation map')
+
+ args = parser.parse_args()
+ if args.config is None:
+ args.config = os.path.join(args.segformer_path, 'local_configs', 'segformer', 'B5', 'segformer.b5.1024x1024.city.160k.py')
+ if args.checkpoint is None:
+ args.checkpoint = os.path.join(args.segformer_path, 'pretrained', 'segformer.b5.1024x1024.city.160k.pth')
+
+ if args.scene_ids is not None:
+ scene_ids_list = args.scene_ids
+ elif args.split_file is not None:
+ # parse the split file
+ split_file = open(args.split_file, "r").readlines()[1:]
+ # NOTE: small hack here, to be refined in the futher (TODO)
+ if "kitti" in args.split_file or "nuplan" in args.split_file:
+ scene_ids_list = [line.strip().split(",")[0] for line in split_file]
+ else:
+ scene_ids_list = [int(line.strip().split(",")[0]) for line in split_file]
+ else:
+ scene_ids_list = sorted(os.listdir(args.data_root))[:args.num_scenes]
+ # scene_ids_list = np.arange(args.start_idx, args.start_idx + args.num_scenes)
+
+ model = init_segmentor(args.config, args.checkpoint, device=args.device)
+
+ for scene_i, scene_id in enumerate(tqdm(scene_ids_list, f'Extracting Masks ...')):
+ scene_id = str(scene_id).zfill(3)
+ img_dir = os.path.join(args.data_root, scene_id, args.rgb_dirname)
+
+ wait_flag = False
+ while not os.path.exists(os.path.join(args.data_root, scene_id, "instances", "instances_info.json")):
+ if not wait_flag:
+ print(f"Scene {scene_id} has less than 16*6 images, waiting ...")
+ wait_flag = True
+ wait_flag = False
+
+ # create mask dir
+ sky_mask_dir = os.path.join(args.data_root, scene_id, "sky_masks")
+ if not os.path.exists(sky_mask_dir):
+ os.makedirs(sky_mask_dir)
+
+ # create dynamic mask dir
+ if args.process_dynamic_mask:
+ rough_human_mask_dir = os.path.join(args.data_root, scene_id, "dynamic_masks", "human")
+ rough_vehicle_mask_dir = os.path.join(args.data_root, scene_id, "dynamic_masks", "vehicle")
+
+ all_mask_dir = os.path.join(args.data_root, scene_id, "fine_dynamic_masks", "all")
+ if not os.path.exists(all_mask_dir):
+ os.makedirs(all_mask_dir)
+ human_mask_dir = os.path.join(args.data_root, scene_id, "fine_dynamic_masks", "human")
+ if not os.path.exists(human_mask_dir):
+ os.makedirs(human_mask_dir)
+ vehicle_mask_dir = os.path.join(args.data_root, scene_id, "fine_dynamic_masks", "vehicle")
+ if not os.path.exists(vehicle_mask_dir):
+ os.makedirs(vehicle_mask_dir)
+
+ flist = sorted(glob(os.path.join(img_dir, '*')))
+ for fpath in tqdm(flist, f'scene[{scene_id}]'):
+ fbase = os.path.splitext(os.path.basename(os.path.normpath(fpath)))[0]
+
+ # if args.no_compress:
+ # mask_fpath = os.path.join(mask_dir, f"{fbase}.npy")
+ # else:
+ # mask_fpath = os.path.join(mask_dir, f"{fbase}.npz")
+
+ if args.ignore_existing and os.path.exists(os.path.join(args.data_root, scene_id, "fine_dynamic_masks")) and \
+ len(os.listdir(os.path.join(args.data_root, scene_id, 'fine_dynamic_masks', 'all'))) == len(os.listdir(os.path.join(args.data_root, scene_id, 'images'))) and \
+ len(os.listdir(os.path.join(args.data_root, scene_id, 'fine_dynamic_masks', 'human'))) == len(os.listdir(os.path.join(args.data_root, scene_id, 'images'))) and \
+ len(os.listdir(os.path.join(args.data_root, scene_id, 'fine_dynamic_masks', 'vehicle'))) == len(os.listdir(os.path.join(args.data_root, scene_id, 'images'))):
+ print(f"Scene {scene_id} already processed, skip ...")
+ continue
+
+ #---- Inference and save outputs
+ result = inference_segmentor(model, fpath)
+ mask = result[0].astype(np.uint8) # NOTE: in the settings of "cityscapes", there are 19 classes at most.
+ # if args.no_compress:
+ # np.save(mask_fpath, mask)
+ # else:
+ # np.savez_compressed(mask_fpath, mask) # NOTE: compressed files are 100x smaller.
+
+ # save sky mask
+ sky_mask = np.isin(mask, [10])
+ imageio.imwrite(os.path.join(sky_mask_dir, f"{fbase}.png"), sky_mask.astype(np.uint8)*255)
+
+ if args.process_dynamic_mask:
+ # save human masks
+ rough_human_mask_path = os.path.join(rough_human_mask_dir, f"{fbase}.png")
+ rough_human_mask = (imageio.imread(rough_human_mask_path) > 0)
+ huamn_mask = np.isin(mask, dataset_classes_in_sematic['human'])
+ valid_human_mask = np.logical_and(huamn_mask, rough_human_mask)
+ imageio.imwrite(os.path.join(human_mask_dir, f"{fbase}.png"), valid_human_mask.astype(np.uint8)*255)
+
+ # save vehicle mask
+ rough_vehicle_mask_path = os.path.join(rough_vehicle_mask_dir, f"{fbase}.png")
+ rough_vehicle_mask = (imageio.imread(rough_vehicle_mask_path) > 0)
+ vehicle_mask = np.isin(mask, dataset_classes_in_sematic['Vehicle'])
+ valid_vehicle_mask = np.logical_and(vehicle_mask, rough_vehicle_mask)
+ imageio.imwrite(os.path.join(vehicle_mask_dir, f"{fbase}.png"), valid_vehicle_mask.astype(np.uint8)*255)
+
+ # save dynamic mask
+ valid_all_mask = np.logical_or(valid_human_mask, valid_vehicle_mask)
+ imageio.imwrite(os.path.join(all_mask_dir, f"{fbase}.png"), valid_all_mask.astype(np.uint8)*255)
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/tools/extract_masks_mmsegnew.py b/worldbench/third_party/drivestudio/datasets/tools/extract_masks_mmsegnew.py
new file mode 100644
index 0000000..07de5fe
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/tools/extract_masks_mmsegnew.py
@@ -0,0 +1,224 @@
+"""
+@file extract_masks.py
+@author Jianfei Guo, Shanghai AI Lab
+@brief Extract semantic mask
+
+Using SegFormer, 2021. Cityscapes 83.2%
+Relies on timm==0.3.2 & pytorch 1.8.1 (buggy on pytorch >= 1.9)
+
+Installation:
+ NOTE: mmcv-full==1.2.7 requires another pytorch version & conda env.
+ Currently mmcv-full==1.2.7 does not support pytorch>=1.9;
+ will raise AttributeError: 'super' object has no attribute '_specify_ddp_gpu_num'
+ Hence, a seperate conda env is needed.
+
+ git clone https://github.com/NVlabs/SegFormer
+
+ conda create -n segformer python=3.8
+ conda activate segformer
+ # conda install pytorch==1.8.1 torchvision==0.9.1 torchaudio==0.8.1 cudatoolkit=11.3 -c pytorch -c conda-forge
+ pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html
+
+ pip install timm==0.3.2 pylint debugpy opencv-python attrs ipython tqdm imageio scikit-image omegaconf
+ pip install mmcv-full==1.2.7 --no-cache-dir
+
+ cd SegFormer
+ pip install .
+
+Usage:
+ Direct run this script in the newly set conda env.
+"""
+import torch
+from mmseg.apis import MMSegInferencer
+
+#from mmseg.apis import inference_segmentor, init_segmentor, show_result_pyplot
+#from mmseg.core.evaluation import get_palette
+
+semantic_classes = [
+ 'road', 'sidewalk', 'building', 'wall', 'fence', 'pole',
+ 'traffic light', 'traffic sign', 'vegetation', 'terrain', 'sky',
+ 'person', 'rider', 'car', 'truck', 'bus', 'train', 'motorcycle',
+ 'bicycle'
+]
+dataset_classes_in_sematic = {
+ 'Vehicle': [13, 14, 15], # 'car', 'truck', 'bus'
+ 'human': [11, 12, 17, 18], # 'person', 'rider', 'motorcycle', 'bicycle'
+}
+
+############ Segmentation Related ############
+@torch.inference_mode()
+def inference_mmseg(video_numpy_list, inferenecer):
+ """
+ Inference the video numpy list using mmseg
+
+ Args:
+ video_numpy_list: list,
+ list of numpy array with shape [H, W, 3],
+ each numpy array is a frame of the video
+
+ inferencer: MMSegInferencer
+
+ Returns:
+ list: list of semantic segmentation results for each frame
+ """
+ pred_list = []
+
+ # video_numpy_list can be too long for one inference, so we split it into several parts
+ chunking_num = 15
+ chunk_index = torch.linspace(0, len(video_numpy_list), chunking_num + 1).long()
+ print("Inference mmseg...")
+
+ for idx in range(chunking_num):
+ # print(f"Processing segformer chunk {idx}/{chunking_num}")
+ if chunk_index[idx] == chunk_index[idx+1]:
+ continue
+ video_numpy_chunk = video_numpy_list[chunk_index[idx]:chunk_index[idx+1]]
+ semantic_map = inferenecer(video_numpy_chunk)['predictions']
+ if isinstance(semantic_map, np.ndarray):
+ semantic_map = [semantic_map]
+
+ pred_list.extend(semantic_map)
+
+ return pred_list
+
+if __name__ == "__main__":
+ import os
+ import imageio
+ import numpy as np
+ from glob import glob
+ from tqdm import tqdm
+ from argparse import ArgumentParser
+ parser = ArgumentParser()
+ # Custom configs
+ parser.add_argument('--data_root', type=str, default='data/waymo/processed/training')
+ parser.add_argument(
+ "--scene_ids",
+ default=None,
+ type=int,
+ nargs="+",
+ help="scene ids to be processed, a list of integers separated by space. Range: [0, 798] for training, [0, 202] for validation",
+ )
+ parser.add_argument(
+ "--split_file", type=str, default=None, help="Split file in data/waymo_splits"
+ )
+ parser.add_argument(
+ "--start_idx",
+ type=int,
+ default=0,
+ help="If no scene id or split_file is given, use start_idx and num_scenes to generate scene_ids_list",
+ )
+ parser.add_argument(
+ "--num_scenes",
+ type=int,
+ default=200,
+ help="number of scenes to be processed",
+ )
+ parser.add_argument(
+ '--process_dynamic_mask',
+ action='store_true',
+ help="Whether to process dynamic masks",
+ )
+ parser.add_argument('--verbose', action='store_true')
+ parser.add_argument('--ignore_existing', action='store_true')
+ parser.add_argument('--no_compress', action='store_true')
+ parser.add_argument('--rgb_dirname', type=str, default="images")
+ parser.add_argument('--mask_dirname', type=str, default="fine_dynamic_masks")
+
+ # Algorithm configs
+ parser.add_argument('--segformer_path', type=str, default='/home/guojianfei/ai_ws/SegFormer')
+ parser.add_argument('--config', help='Config file', type=str, default=None)
+ parser.add_argument('--checkpoint', help='Checkpoint file', type=str, default=None)
+ parser.add_argument('--device', default='cuda:0', help='Device used for inference')
+ parser.add_argument('--palette', default='cityscapes', help='Color palette used for segmentation map')
+
+ args = parser.parse_args()
+ if args.config is None:
+ args.config = os.path.join(args.segformer_path, 'local_configs', 'segformer', 'B5', 'segformer.b5.1024x1024.city.160k.py')
+ if args.checkpoint is None:
+ args.checkpoint = os.path.join(args.segformer_path, 'pretrained', 'segformer.b5.1024x1024.city.160k.pth')
+
+ if args.scene_ids is not None:
+ scene_ids_list = args.scene_ids
+ elif args.split_file is not None:
+ # parse the split file
+ split_file = open(args.split_file, "r").readlines()[1:]
+ # NOTE: small hack here, to be refined in the futher (TODO)
+ if "kitti" in args.split_file or "nuplan" in args.split_file:
+ scene_ids_list = [line.strip().split(",")[0] for line in split_file]
+ else:
+ scene_ids_list = [int(line.strip().split(",")[0]) for line in split_file]
+ else:
+ scene_ids_list = np.arange(args.start_idx, args.start_idx + args.num_scenes)
+
+ # model = init_segmentor(args.config, args.checkpoint, device=args.device)
+ inferenecer = MMSegInferencer(model='segformer_mit-b5_8xb1-160k_cityscapes-1024x1024', device='cuda')
+
+ for scene_i, scene_id in enumerate(tqdm(scene_ids_list, f'Extracting Masks ...')):
+ scene_id = str(scene_id).zfill(3)
+ img_dir = os.path.join(args.data_root, scene_id, args.rgb_dirname)
+
+ # create mask dir
+ sky_mask_dir = os.path.join(args.data_root, scene_id, "sky_masks")
+ if not os.path.exists(sky_mask_dir):
+ os.makedirs(sky_mask_dir)
+
+ # create dynamic mask dir
+ if args.process_dynamic_mask:
+ rough_human_mask_dir = os.path.join(args.data_root, scene_id, "dynamic_masks", "human")
+ rough_vehicle_mask_dir = os.path.join(args.data_root, scene_id, "dynamic_masks", "vehicle")
+
+ all_mask_dir = os.path.join(args.data_root, scene_id, "fine_dynamic_masks", "all")
+ if not os.path.exists(all_mask_dir):
+ os.makedirs(all_mask_dir)
+ human_mask_dir = os.path.join(args.data_root, scene_id, "fine_dynamic_masks", "human")
+ if not os.path.exists(human_mask_dir):
+ os.makedirs(human_mask_dir)
+ vehicle_mask_dir = os.path.join(args.data_root, scene_id, "fine_dynamic_masks", "vehicle")
+ if not os.path.exists(vehicle_mask_dir):
+ os.makedirs(vehicle_mask_dir)
+
+ flist = sorted(glob(os.path.join(img_dir, '*')))
+ results = inference_mmseg(flist, inferenecer)
+ for i, fpath in enumerate(tqdm(flist, f'scene[{scene_id}]')):
+ fbase = os.path.splitext(os.path.basename(os.path.normpath(fpath)))[0]
+
+ # if args.no_compress:
+ # mask_fpath = os.path.join(mask_dir, f"{fbase}.npy")
+ # else:
+ # mask_fpath = os.path.join(mask_dir, f"{fbase}.npz")
+
+ if args.ignore_existing and os.path.exists(os.path.join(args.data_root, scene_id, "fine_dynamic_masks")):
+ continue
+
+ #---- Inference and save outputs
+ #result = inference_segmentor(model, fpath)
+ result = results[i]
+ mask = result.astype(np.uint8) # NOTE: in the settings of "cityscapes", there are 19 classes at most.
+
+ # if args.no_compress:
+ # np.save(mask_fpath, mask)
+ # else:
+ # np.savez_compressed(mask_fpath, mask) # NOTE: compressed files are 100x smaller.
+
+ # save sky mask
+ sky_mask = np.isin(mask, [10])
+ imageio.imwrite(os.path.join(sky_mask_dir, f"{fbase}.png"), sky_mask.astype(np.uint8)*255)
+
+ if args.process_dynamic_mask:
+ # save human masks
+ rough_human_mask_path = os.path.join(rough_human_mask_dir, f"{fbase}.png")
+ rough_human_mask = (imageio.imread(rough_human_mask_path) > 0)
+ huamn_mask = np.isin(mask, dataset_classes_in_sematic['human'])
+ valid_human_mask = np.logical_and(huamn_mask, rough_human_mask)
+ imageio.imwrite(os.path.join(human_mask_dir, f"{fbase}.png"), valid_human_mask.astype(np.uint8)*255)
+
+ # save vehicle mask
+ rough_vehicle_mask_path = os.path.join(rough_vehicle_mask_dir, f"{fbase}.png")
+ rough_vehicle_mask = (imageio.imread(rough_vehicle_mask_path) > 0)
+ vehicle_mask = np.isin(mask, dataset_classes_in_sematic['Vehicle'])
+ valid_vehicle_mask = np.logical_and(vehicle_mask, rough_vehicle_mask)
+ imageio.imwrite(os.path.join(vehicle_mask_dir, f"{fbase}.png"), valid_vehicle_mask.astype(np.uint8)*255)
+
+ # save dynamic mask
+ valid_all_mask = np.logical_or(valid_human_mask, valid_vehicle_mask)
+ imageio.imwrite(os.path.join(all_mask_dir, f"{fbase}.png"), valid_all_mask.astype(np.uint8)*255)
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/tools/extract_smpl.py b/worldbench/third_party/drivestudio/datasets/tools/extract_smpl.py
new file mode 100644
index 0000000..7641280
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/tools/extract_smpl.py
@@ -0,0 +1,288 @@
+import os
+import warnings
+import logging
+from glob import glob
+from pathlib import Path
+from typing import Optional, Tuple, List
+from dataclasses import dataclass
+
+import cv2
+import hydra
+import joblib
+import numpy as np
+import torch
+from hydra.core.config_store import ConfigStore
+from hydra.core.global_hydra import GlobalHydra
+
+from phalp.configs.base import FullConfig, CACHE_DIR
+from phalp.models.hmar.hmr import HMR2018Predictor
+from phalp.trackers.PHALP import PHALP
+from phalp.utils import get_pylogger
+from third_party.Humans4D.hmr2.datasets.utils import expand_bbox_to_aspect_ratio
+
+warnings.filterwarnings('ignore')
+
+logger = logging.getLogger()
+
+class HMR2Predictor(HMR2018Predictor):
+ def __init__(self, cfg) -> None:
+ super().__init__(cfg)
+ # Setup our new model
+ from third_party.Humans4D.hmr2.models import download_models, load_hmr2
+
+ # Download and load checkpoints
+ download_models()
+ model, _ = load_hmr2()
+
+ self.model = model
+ self.model.eval()
+
+ def forward(self, x):
+ hmar_out = self.hmar_old(x)
+ batch = {
+ 'img': x[:,:3,:,:],
+ 'mask': (x[:,3,:,:]).clip(0,1),
+ }
+ model_out = self.model(batch)
+ out = hmar_out | {
+ 'pose_smpl': model_out['pred_smpl_params'],
+ 'pred_cam': model_out['pred_cam'],
+ }
+ return out
+
+class HMR2023TextureSampler(HMR2Predictor):
+ def __init__(self, cfg) -> None:
+ super().__init__(cfg)
+
+ # Model's all set up. Now, load tex_bmap and tex_fmap
+ # Texture map atlas
+ bmap = np.load(os.path.join(CACHE_DIR, 'phalp/3D/bmap_256.npy'))
+ fmap = np.load(os.path.join(CACHE_DIR, 'phalp/3D/fmap_256.npy'))
+ self.register_buffer('tex_bmap', torch.tensor(bmap, dtype=torch.float))
+ self.register_buffer('tex_fmap', torch.tensor(fmap, dtype=torch.long))
+
+ self.img_size = 256 #self.cfg.MODEL.IMAGE_SIZE
+ self.focal_length = 5000. #self.cfg.EXTRA.FOCAL_LENGTH
+
+ import neural_renderer as nr
+ self.neural_renderer = nr.Renderer(dist_coeffs=None, orig_size=self.img_size,
+ image_size=self.img_size,
+ light_intensity_ambient=1,
+ light_intensity_directional=0,
+ anti_aliasing=False)
+
+ def forward(self, x):
+ batch = {
+ 'img': x[:,:3,:,:],
+ 'mask': (x[:,3,:,:]).clip(0,1),
+ }
+ model_out = self.model(batch)
+
+ # from hmr2.models.prohmr_texture import unproject_uvmap_to_mesh
+
+ def unproject_uvmap_to_mesh(bmap, fmap, verts, faces):
+ # bmap: 256,256,3
+ # fmap: 256,256
+ # verts: B,V,3
+ # faces: F,3
+ valid_mask = (fmap >= 0)
+
+ fmap_flat = fmap[valid_mask] # N
+ bmap_flat = bmap[valid_mask,:] # N,3
+
+ face_vids = faces[fmap_flat, :] # N,3
+ face_verts = verts[:, face_vids, :] # B,N,3,3
+
+ bs = face_verts.shape
+ map_verts = torch.einsum('bnij,ni->bnj', face_verts, bmap_flat) # B,N,3
+
+ return map_verts, valid_mask
+
+ pred_verts = model_out['pred_vertices'] + model_out['pred_cam_t'].unsqueeze(1)
+ device = pred_verts.device
+ face_tensor = torch.tensor(self.smpl.faces.astype(np.int64), dtype=torch.long, device=device)
+ map_verts, valid_mask = unproject_uvmap_to_mesh(self.tex_bmap, self.tex_fmap, pred_verts, face_tensor) # B,N,3
+
+ # Project map_verts to image using K,R,t
+ # map_verts_view = einsum('bij,bnj->bni', R, map_verts) + t # R=I t=0
+ focal = self.focal_length / (self.img_size / 2)
+ map_verts_proj = focal * map_verts[:, :, :2] / map_verts[:, :, 2:3] # B,N,2
+ map_verts_depth = map_verts[:, :, 2] # B,N
+
+ # Render Depth. Annoying but we need to create this
+ K = torch.eye(3, device=device)
+ K[0, 0] = K[1, 1] = self.focal_length
+ K[1, 2] = K[0, 2] = self.img_size / 2 # Because the neural renderer only support squared images
+ K = K.unsqueeze(0)
+ R = torch.eye(3, device=device).unsqueeze(0)
+ t = torch.zeros(3, device=device).unsqueeze(0)
+ rend_depth = self.neural_renderer(pred_verts,
+ face_tensor[None].expand(pred_verts.shape[0], -1, -1).int(),
+ # textures=texture_atlas_rgb,
+ mode='depth',
+ K=K, R=R, t=t)
+
+ rend_depth_at_proj = torch.nn.functional.grid_sample(rend_depth[:,None,:,:], map_verts_proj[:,None,:,:]) # B,1,1,N
+ rend_depth_at_proj = rend_depth_at_proj.squeeze(1).squeeze(1) # B,N
+
+ img_rgba = torch.cat([batch['img'], batch['mask'][:,None,:,:]], dim=1) # B,4,H,W
+ img_rgba_at_proj = torch.nn.functional.grid_sample(img_rgba, map_verts_proj[:,None,:,:]) # B,4,1,N
+ img_rgba_at_proj = img_rgba_at_proj.squeeze(2) # B,4,N
+
+ visibility_mask = map_verts_depth <= (rend_depth_at_proj + 1e-4) # B,N
+ img_rgba_at_proj[:,3,:][~visibility_mask] = 0
+
+ # Paste image back onto square uv_image
+ uv_image = torch.zeros((batch['img'].shape[0], 4, 256, 256), dtype=torch.float, device=device)
+ uv_image[:, :, valid_mask] = img_rgba_at_proj
+
+ out = {
+ 'uv_image': uv_image,
+ 'uv_vector' : self.hmar_old.process_uv_image(uv_image),
+ 'pose_smpl': model_out['pred_smpl_params'],
+ 'pred_cam': model_out['pred_cam'],
+ }
+ return out
+
+class HMR2_4dhuman(PHALP):
+ def __init__(self, cfg):
+ super().__init__(cfg)
+
+ def setup_hmr(self):
+ self.HMAR = HMR2023TextureSampler(self.cfg)
+
+ def get_detections(self, image, frame_name, t_, additional_data=None, measurments=None):
+ (
+ pred_bbox, pred_bbox, pred_masks, pred_scores, pred_classes,
+ ground_truth_track_id, ground_truth_annotations
+ ) = super().get_detections(image, frame_name, t_, additional_data, measurments)
+
+ # Pad bounding boxes
+ pred_bbox_padded = expand_bbox_to_aspect_ratio(pred_bbox, self.cfg.expand_bbox_shape)
+
+ return (
+ pred_bbox, pred_bbox_padded, pred_masks, pred_scores, pred_classes,
+ ground_truth_track_id, ground_truth_annotations
+ )
+
+
+@dataclass
+class Human4DConfig(FullConfig):
+ # override defaults if needed
+ expand_bbox_shape: Optional[Tuple[int]] = (192,256)
+ pass
+
+cs = ConfigStore.instance()
+cs.store(name="config", node=Human4DConfig)
+
+def initialize_config(config_name="config"):
+ hydra.initialize(version_base="1.2", config_path=".")
+ cfg = hydra.compose(config_name=config_name)
+ return cfg
+
+def run_4DHumans(
+ scene_dir: str, camera_list: List[int], save_temp: bool=True, verbose: bool=False, fps: int=12
+) -> Optional[float]:
+ """Main function for running the PHALP tracker.
+
+ Args:
+ scene_dir: Path to the scene directory containing the images files
+ camera_list: List of camera IDs to run the tracker on
+ save_temp: Whether to save temporary files or not,
+ recommended to set to True if you have enough disk space
+ verbose: Whether to visualize the smpl mesh in video or not
+ fps: Frames per second of the input images
+
+ Returns:
+ pred_tracks_allcam: Dictionary containing the predicted tracks for each camera
+ """
+ # make sure the scene directory exists
+ assert os.path.exists(scene_dir), f"Scene directory {scene_dir} does not exist"
+ images_dir = os.path.join(scene_dir, 'images')
+ assert os.path.exists(images_dir), f"Images directory {images_dir} does not exist"
+ temp_dir = os.path.join(scene_dir, 'humanpose', 'temp')
+
+ # create a flag to check if the results already exist
+ already_done = False
+
+ for cam_id in camera_list:
+ results_path = os.path.join(temp_dir, 'phalp_output', f'cam_{cam_id}.pkl')
+ if os.path.exists(results_path):
+ logger.info(f"Results for camera {cam_id} already exists at {results_path}")
+ already_done = True
+ continue
+
+ cfg = initialize_config()
+
+ cached_video_path = os.path.join(temp_dir, 'raw_videos', f'{cam_id}.mp4')
+ if not os.path.exists(cached_video_path):
+ logger.info(f"Cached video not found at {cached_video_path}, we will create it")
+ os.makedirs(os.path.dirname(cached_video_path), exist_ok=True)
+
+ image_paths = sorted(glob(os.path.join(images_dir, f"*_{cam_id}.*")))
+
+ # create video
+ height, width = cv2.imread(image_paths[0]).shape[:2]
+ fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+ out = cv2.VideoWriter(cached_video_path, fourcc, fps, (width, height))
+ for image_path in image_paths:
+ frame = cv2.imread(image_path)
+ out.write(frame)
+ out.release()
+ cfg.video.source = cached_video_path
+
+ output_dir = os.path.join(temp_dir, 'phalp_output')
+ if not os.path.exists(output_dir):
+ os.makedirs(output_dir)
+ cfg.video.output_dir = output_dir
+
+ phalp_tracker = HMR2_4dhuman(cfg)
+ phalp_tracker.track()
+
+ # reset global hydra
+ GlobalHydra.instance().clear()
+
+ # remove temporary or useless files
+ dirs_to_remove = [
+ os.path.join(temp_dir, 'phalp_output', '_DEMO'),
+ os.path.join(temp_dir, 'phalp_output', '_TMP'),
+ os.path.join(temp_dir, 'phalp_output', 'results_tracks')
+ ]
+ if not save_temp:
+ # remove temporary files
+ dirs_to_remove.append(os.path.join(temp_dir, 'raw_videos'))
+ for d in dirs_to_remove:
+ if os.path.exists(d):
+ os.system(f"rm -rf {d}")
+
+ if not verbose:
+ # remove the video files generated by PHALP
+ files_to_remove = [
+ os.path.join(temp_dir, 'phalp_output', f'PHALP_{i}.mp4') for i in camera_list
+ ]
+ for f in files_to_remove:
+ if os.path.exists(f):
+ os.remove(f)
+
+ # load those pickle files to return
+ pred_tracks_allcam = {}
+ for i, cam_id in enumerate(camera_list):
+ if already_done:
+ pred_tracks_allcam[cam_id] = joblib.load(
+ os.path.join(temp_dir, 'phalp_output', f'cam_{cam_id}.pkl')
+ )
+ else:
+ pred_tracks_allcam[cam_id] = joblib.load(
+ os.path.join(temp_dir, 'phalp_output', 'results', f'demo_{i}.pkl')
+ )
+
+ # move and rename saved pickle files if save_temp
+ if save_temp and not already_done:
+ for i, cam_id in enumerate(camera_list):
+ os.rename(
+ os.path.join(temp_dir, 'phalp_output', 'results', f'demo_{i}.pkl'),
+ os.path.join(temp_dir, 'phalp_output', f'cam_{cam_id}.pkl')
+ )
+ os.system(f"rm -rf {os.path.join(temp_dir, 'phalp_output', 'results')}")
+
+ return pred_tracks_allcam
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/tools/humanpose_process.py b/worldbench/third_party/drivestudio/datasets/tools/humanpose_process.py
new file mode 100644
index 0000000..26a7146
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/tools/humanpose_process.py
@@ -0,0 +1,159 @@
+from typing import List, Callable
+import os
+import joblib
+import logging
+import argparse
+import numpy as np
+from hydra.core.global_hydra import GlobalHydra
+
+
+from datasets.tools.extract_smpl import run_4DHumans
+from datasets.tools.postprocess import match_and_postprocess
+
+logger = logging.getLogger()
+
+def extract_humanpose(
+ scene_dir,
+ projection_fn: Callable,
+ camera_list: List[str],
+ save_temp: bool=True,
+ verbose: bool=False,
+ fps: int=12
+):
+ """Extract human pose from the waymo dataset
+
+ Args:
+ scene_dir: str, path to the scene directory
+ save_temp: bool, whether to save the intermediate results
+ verbose: bool, whether to visualize debug images
+ fps: int, FPS for the visualization video
+ """
+ # project human boxes to 2D image space
+ GTTracks_meta = projection_fn(
+ scene_dir, camera_list=camera_list,
+ save_temp=save_temp, verbose=verbose,
+ narrow_width_ratio=0.2, fps=fps
+ )
+
+ # run 4DHuman to get predicted human tracks with SMPL parameters
+ PredTracks_meta = run_4DHumans(
+ scene_dir, camera_list=camera_list,
+ save_temp=save_temp, verbose=verbose, fps=fps
+ )
+
+ # match the predicted tracks with the ground truth tracks
+ smpl_meta = match_and_postprocess(
+ scene_dir, camera_list=camera_list,
+ GTTracksDict=GTTracks_meta, PredTracksDict=PredTracks_meta,
+ save_temp=save_temp, verbose=verbose, fps=fps
+ )
+
+ joblib.dump(
+ smpl_meta,
+ os.path.join(scene_dir, "humanpose", "smpl.pkl")
+ )
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(description="Data converter arg parser")
+ parser.add_argument("--data_root", type=str, required=True, help="root path of waymo dataset")
+ parser.add_argument("--dataset", type=str, default="waymo", help="dataset name")
+ parser.add_argument(
+ "--scene_ids",
+ default=None,
+ type=int,
+ nargs="+",
+ help="scene ids to be processed, a list of integers separated by space. Range: [0, 798] for training, [0, 202] for validation",
+ )
+ parser.add_argument(
+ "--split_file", type=str, default=None, help="Split file in data/waymo_splits"
+ )
+ parser.add_argument(
+ "--start_idx",
+ type=int,
+ default=0,
+ help="If no scene id or split_file is given, use start_idx and num_scenes to generate scene_ids_list",
+ )
+ parser.add_argument(
+ "--num_scenes",
+ type=int,
+ default=200,
+ help="number of scenes to be processed",
+ )
+ parser.add_argument(
+ "--save_temp",
+ action="store_true",
+ help="Whether to save the intermediate results",
+ )
+ parser.add_argument(
+ "--verbose",
+ action="store_true",
+ help="Whether to visualize the intermediate results",
+ )
+ parser.add_argument(
+ "--fps",
+ type=int,
+ default=12,
+ help="FPS for the visualization video if verbose is True",
+ )
+ args = parser.parse_args()
+
+ if args.dataset == "waymo":
+ from datasets.waymo.waymo_human_utils import project_human_boxes, CAMERA_LIST
+ elif args.dataset == "pandaset":
+ from datasets.pandaset.pandaset_human_utils import project_human_boxes, CAMERA_LIST
+ elif args.dataset == "argoverse":
+ from datasets.argoverse.argoverse_human_utils import project_human_boxes, CAMERA_LIST
+ elif args.dataset == "nuscenes":
+ from datasets.nuscenes.nuscenes_human_utils import project_human_boxes, CAMERA_LIST
+ elif args.dataset == "kitti":
+ from datasets.kitti.kitti_human_utils import project_human_boxes, CAMERA_LIST
+ elif args.dataset == "nuplan":
+ from datasets.nuplan.nuplan_human_utils import project_human_boxes, CAMERA_LIST
+ else:
+ raise ValueError(f"Unknown dataset {args.dataset}, please choose from waymo, pandaset, argoverse, nuscenes, kitti, nuplan")
+
+ if args.scene_ids is not None:
+ scene_ids_list = args.scene_ids
+ elif args.split_file is not None:
+ # parse the split file
+ split_file = open(args.split_file, "r").readlines()[1:]
+ try:
+ # Waymo, Pandaset, Argoverse, NuScenes
+ scene_ids_list = [int(line.strip().split(",")[0]) for line in split_file]
+ except:
+ # KITTI
+ scene_ids_list = [line.strip().split(" ")[0] for line in split_file]
+ else:
+ # scene_ids_list = np.arange(args.start_idx, args.start_idx + args.num_scenes)
+ scene_ids_list = sorted(os.listdir(args.data_root))
+ scene_ids_list = scene_ids_list[:args.num_scenes]
+
+ for scene_id in scene_ids_list:
+ GlobalHydra.instance().clear()
+ try:
+ scene_dir = f'{args.data_root}/{str(scene_id).zfill(3)}'
+
+ wait_flag = False
+ while not os.path.exists(os.path.join(scene_dir, "instances", "instances_info.json")):
+ if not wait_flag:
+ print(f"Scene {scene_id} has less than 16*6 images, waiting ...")
+ wait_flag = True
+ wait_flag = False
+
+ # check exists
+ if os.path.exists(os.path.join(scene_dir, 'humanpose', 'smpl.pkl')):
+ print(f"Scene {scene_id} already processed, skip it.")
+ continue
+
+ extract_humanpose(
+ scene_dir=scene_dir,
+ projection_fn=project_human_boxes,
+ camera_list=CAMERA_LIST,
+ save_temp=args.save_temp,
+ verbose=args.verbose,
+ fps=args.fps
+ )
+ logger.info(f"Finished processing scene {scene_id}")
+ except Exception as e:
+ logger.error(f"Error processing scene {scene_id}: {e}")
+ continue
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/tools/multiprocess_utils.py b/worldbench/third_party/drivestudio/datasets/tools/multiprocess_utils.py
new file mode 100644
index 0000000..dd3bde6
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/tools/multiprocess_utils.py
@@ -0,0 +1,150 @@
+"""
+@brief Parallel processing utilities
+@note Borrowed from: https://github.com/NVlabs/EmerNeRF
+"""
+import sys
+import time
+from collections.abc import Iterable
+from multiprocessing import Pool
+from shutil import get_terminal_size
+
+class ProgressBar:
+ """A progress bar which can print the progress."""
+
+ def __init__(self, task_num=0, bar_width=50, start=True, file=sys.stdout):
+ self.task_num = task_num
+ self.bar_width = bar_width
+ self.completed = 0
+ self.file = file
+ if start:
+ self.start()
+
+ @property
+ def terminal_width(self):
+ width, _ = get_terminal_size()
+ return width
+
+ def start(self):
+ if self.task_num > 0:
+ self.file.write(
+ f'[{" " * self.bar_width}] 0/{self.task_num}, ' "elapsed: 0s, ETA:"
+ )
+ else:
+ self.file.write("completed: 0, elapsed: 0s")
+ self.file.flush()
+ self.start_time = time.time()
+
+ def update(self, num_tasks=1):
+ assert num_tasks > 0
+ self.completed += num_tasks
+ elapsed = time.time() - self.start_time
+ if elapsed > 0:
+ fps = self.completed / elapsed
+ else:
+ fps = float("inf")
+ if self.task_num > 0:
+ percentage = self.completed / float(self.task_num)
+ eta = int(elapsed * (1 - percentage) / percentage + 0.5)
+ msg = (
+ f"\r[{{}}] {self.completed}/{self.task_num}, "
+ f"{fps:.1f} task/s, elapsed: {int(elapsed + 0.5)}s, "
+ f"ETA: {eta:5}s"
+ )
+
+ bar_width = min(
+ self.bar_width,
+ int(self.terminal_width - len(msg)) + 2,
+ int(self.terminal_width * 0.6),
+ )
+ bar_width = max(2, bar_width)
+ mark_width = int(bar_width * percentage)
+ bar_chars = ">" * mark_width + " " * (bar_width - mark_width)
+ self.file.write(msg.format(bar_chars))
+ else:
+ self.file.write(
+ f"completed: {self.completed}, elapsed: {int(elapsed + 0.5)}s,"
+ f" {fps:.1f} tasks/s"
+ )
+ self.file.flush()
+
+
+def init_pool(process_num, initializer=None, initargs=None):
+ if initializer is None:
+ return Pool(process_num)
+ elif initargs is None:
+ return Pool(process_num, initializer)
+ else:
+ if not isinstance(initargs, tuple):
+ raise TypeError('"initargs" must be a tuple')
+ return Pool(process_num, initializer, initargs)
+
+
+def track_parallel_progress(
+ func,
+ tasks,
+ nproc,
+ initializer=None,
+ initargs=None,
+ bar_width=50,
+ chunksize=1,
+ skip_first=False,
+ keep_order=True,
+ file=sys.stdout,
+):
+ """Track the progress of parallel task execution with a progress bar.
+
+ The built-in :mod:`multiprocessing` module is used for process pools and
+ tasks are done with :func:`Pool.map` or :func:`Pool.imap_unordered`.
+
+ Args:
+ func (callable): The function to be applied to each task.
+ tasks (list or tuple[Iterable, int]): A list of tasks or
+ (tasks, total num).
+ nproc (int): Process (worker) number.
+ initializer (None or callable): Refer to :class:`multiprocessing.Pool`
+ for details.
+ initargs (None or tuple): Refer to :class:`multiprocessing.Pool` for
+ details.
+ chunksize (int): Refer to :class:`multiprocessing.Pool` for details.
+ bar_width (int): Width of progress bar.
+ skip_first (bool): Whether to skip the first sample for each worker
+ when estimating fps, since the initialization step may takes
+ longer.
+ keep_order (bool): If True, :func:`Pool.imap` is used, otherwise
+ :func:`Pool.imap_unordered` is used.
+
+ Returns:
+ list: The task results.
+ """
+ if isinstance(tasks, tuple):
+ assert len(tasks) == 2
+ assert isinstance(tasks[0], Iterable)
+ assert isinstance(tasks[1], int)
+ task_num = tasks[1]
+ tasks = tasks[0]
+ elif isinstance(tasks, Iterable):
+ task_num = len(tasks)
+ else:
+ raise TypeError('"tasks" must be an iterable object or a (iterator, int) tuple')
+ pool = init_pool(nproc, initializer, initargs)
+ start = not skip_first
+ task_num -= nproc * chunksize * int(skip_first)
+ prog_bar = ProgressBar(task_num, bar_width, start, file=file)
+ results = []
+ if keep_order:
+ gen = pool.imap(func, tasks, chunksize)
+ else:
+ gen = pool.imap_unordered(func, tasks, chunksize)
+ for result in gen:
+ results.append(result)
+ if skip_first:
+ if len(results) < nproc * chunksize:
+ continue
+ elif len(results) == nproc * chunksize:
+ prog_bar.start()
+ continue
+ prog_bar.update()
+ prog_bar.file.write("\n")
+ pool.close()
+ pool.join()
+ return results
diff --git a/worldbench/third_party/drivestudio/datasets/tools/postprocess.py b/worldbench/third_party/drivestudio/datasets/tools/postprocess.py
new file mode 100644
index 0000000..fb5b694
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/tools/postprocess.py
@@ -0,0 +1,602 @@
+from typing import List, Dict
+import os
+import joblib
+
+import torch
+from pytorch3d.transforms import matrix_to_quaternion, quaternion_to_matrix
+
+def interpolate_matrix(matrix1, matrix2, fraction):
+ # use quaternion slerp
+ q1 = matrix_to_quaternion(matrix1)
+ q2 = matrix_to_quaternion(matrix2)
+ dot = (q1 * q2).sum(dim=-1)
+ dot = torch.clamp(dot, -1, 1)
+
+ neg_mask = dot < 0
+ q2[neg_mask] = -q2[neg_mask]
+ dot[neg_mask] = -dot[neg_mask]
+
+ similar_mask = dot > 0.9995
+ q_interp_similar = q1 + fraction * (q2 - q1)
+
+ theta_0 = torch.acos(dot)
+ theta = theta_0 * fraction
+
+ sin_theta = torch.sin(theta)
+ sin_theta_0 = torch.sin(theta_0)
+
+ s1 = torch.cos(theta) - dot * sin_theta / sin_theta_0
+ s2 = sin_theta / sin_theta_0
+
+ q_interp = (s1[:, None] * q1) + (s2[:, None] * q2)
+
+ final_q_interp = torch.zeros_like(q1)
+ final_q_interp[similar_mask] = q_interp_similar[similar_mask]
+ final_q_interp[~similar_mask] = q_interp[~similar_mask]
+ return quaternion_to_matrix(final_q_interp)
+
+def compute_iou(bbox1, bbox2):
+ x1, y1, w1, h1 = bbox1
+ x2, y2, w2, h2 = bbox2
+ xA = max(x1, x2)
+ yA = max(y1, y2)
+ xB = min(x1 + w1, x2 + w2)
+ yB = min(y1 + h1, y2 + h2)
+ interArea = max(0, xB - xA) * max(0, yB - yA)
+ boxAArea = w1 * h1
+ boxBArea = w2 * h2
+ union_area = boxAArea + boxBArea - interArea
+ iou = (interArea / union_area) if union_area > 0 else 0
+ return iou
+
+def interpolate_features(mask, features, is_rot_mat=False):
+ features = features.clone()
+ n = features.shape[0]
+ for i in range(n):
+ if not mask[i]:
+ prev_index = i - 1
+ next_index = i + 1
+
+ while prev_index >= 0 and not mask[prev_index]:
+ prev_index -= 1
+
+ while next_index < n and not mask[next_index]:
+ next_index += 1
+
+ if prev_index >= 0 and next_index < n:
+ interp_factor = (i - prev_index) / (next_index - prev_index)
+ if not is_rot_mat:
+ features[i] = (1 - interp_factor) * features[prev_index] + interp_factor * features[next_index]
+ else:
+ features[i] = interpolate_matrix(features[prev_index], features[next_index], interp_factor)
+
+ elif prev_index >= 0:
+ features[i] = features[prev_index]
+ elif next_index < n:
+ features[i] = features[next_index]
+ return features
+
+def detect_breaks_mask(bool_sequence):
+ mask = [False] * len(bool_sequence)
+ in_true_sequence = False
+ start_false_sequence = False
+ last_true_index = -1 # 跟踪最后一个True的位置
+
+ for i, value in enumerate(bool_sequence):
+ if value:
+ if start_false_sequence:
+ start_false_sequence = False
+ if last_true_index != -1 and i < len(bool_sequence) - 1:
+ mask[last_true_index + 1:i] = [True] * (i - last_true_index - 1)
+ in_true_sequence = True
+ last_true_index = i
+ else:
+ if in_true_sequence:
+ start_false_sequence = True
+ in_true_sequence = False
+
+ if start_false_sequence and last_true_index < len(bool_sequence) - 1:
+ mask[last_true_index + 1:len(bool_sequence)] = [False] * (len(bool_sequence) - last_true_index - 1)
+
+ return mask
+
+def match_and_postprocess(
+ scene_dir: str,
+ GTTracksDict: Dict[int, Dict],
+ PredTracksDict: Dict[int, Dict],
+ camera_list: List[int],
+ save_temp=True,
+ verbose=False,
+ fps=12
+):
+ """Match tha predicted tracks to the GT tracks, and postprocess the matched
+ tracks: complete the missing frames, interpolate the missing values, etc.
+
+ Args:
+ scene_dir: str, the path to the scene directory
+ GTTracksDict: Dict[int, Dict], the dictionary of GT tracks
+ PredTracksDict: Dict[int, Dict], the dictionary of predicted tracks
+ camera_list: List[int], the list of camera IDs
+ save_temp: bool, whether to save the temporary results
+ verbose: bool, whether to print the verbose information
+ fps: int, the frame rate of the video, if save the video, the fps will be used
+
+ Returns:
+ merged_collector: Dict[int, Dict], the dictionary of proprocessed SMPL parameters
+ """
+ # number of cameras
+ num_c = len(camera_list)
+
+ if save_temp:
+ temp_dir = os.path.join(scene_dir, "humanpose", "postprocess")
+ if not os.path.exists(temp_dir):
+ os.makedirs(temp_dir)
+
+ # --------------------------------------------------------------------------
+ # Parse ground truth and predicted tracks
+ # --------------------------------------------------------------------------
+ parsed_pred_track_all, parsed_gt_track_all = {}, {}
+ for cam_id in camera_list:
+ parsed_pred_track, parsed_gt_track = {}, {}
+
+ raw_pred_track = PredTracksDict[cam_id]
+ raw_gt_track = GTTracksDict[cam_id]
+ assert len(raw_pred_track) == len(raw_gt_track), \
+ f"len(raw_pred_track) {len(raw_pred_track)} != len(raw_gt_track) {len(raw_gt_track)}"
+ num_f = len(raw_pred_track) # number of frames
+
+ for fi, pred_track_f in enumerate(raw_pred_track.values()):
+ for i_, tid in enumerate(pred_track_f["tid"]):
+ if pred_track_f["tracked_time"][i_] != 0:
+ continue
+ if tid not in parsed_pred_track:
+ parsed_pred_track[tid] = {
+ "valids": torch.zeros(num_f, dtype=torch.bool),
+ "tracked_bbox": torch.zeros(num_f, 4),
+ "smpl": {
+ "global_orient": torch.zeros(num_f, 1, 3, 3),
+ "body_pose": torch.zeros(num_f, 23, 3, 3),
+ "betas": torch.zeros(num_f, 10),
+ },
+ "camera": torch.zeros(num_f, 3),
+ }
+ parsed_pred_track[tid]["valids"][fi] = True
+ parsed_pred_track[tid]["tracked_bbox"][fi] = torch.from_numpy(pred_track_f["bbox"][i_])
+ parsed_pred_track[tid]["smpl"]["global_orient"][fi] = torch.tensor(pred_track_f["smpl"][i_]["global_orient"])
+ parsed_pred_track[tid]["smpl"]["body_pose"][fi] = torch.tensor(pred_track_f["smpl"][i_]["body_pose"])
+ parsed_pred_track[tid]["smpl"]["betas"][fi] = torch.tensor(pred_track_f["smpl"][i_]["betas"])
+ parsed_pred_track[tid]["camera"][fi] = torch.tensor(pred_track_f["camera"][i_])
+
+ # collect GT track ids
+ for fi, gt_track_f in enumerate(raw_gt_track.values()):
+ for i_, tid in enumerate(gt_track_f["extra_data"]["gt_track_id"]):
+ if tid not in parsed_gt_track:
+ parsed_gt_track[tid] = {
+ "valids": torch.zeros(num_f, dtype=torch.bool),
+ "tracked_bbox": torch.zeros(num_f, 4),
+ }
+ parsed_gt_track[tid]["valids"][fi] = True
+ parsed_gt_track[tid]["tracked_bbox"][fi] = torch.tensor(gt_track_f["gt_bbox"][i_])
+
+ parsed_pred_track_all[cam_id] = parsed_pred_track
+ parsed_gt_track_all[cam_id] = parsed_gt_track
+
+ # --------------------------------------------------------------------------
+ # Find GT tracks that have at least one 4D-Humans prediction
+ # --------------------------------------------------------------------------
+ # For each camera:
+ # 1. Compare predicted tracks to GT tracks using IoU (Intersection over Union)
+ # 2. Link each predicted track to its best-matching GT track
+ #
+ # Key points:
+ # - Some GT tracks may not have predictions (e.g., far-away pedestrians)
+ # - A GT track might match multiple predicted tracks
+ # - We keep only GT tracks that have at least one prediction in any camera
+ #
+ # Result: valid_gt_tids list contains GT tracks visible to 4D-Humans in at least one camera
+ valid_gt_tids = []
+ match_pred2gt = {}
+ for cam_id in camera_list:
+ parsed_pred_track = parsed_pred_track_all[cam_id]
+ parsed_gt_track = parsed_gt_track_all[cam_id]
+ _matches = {}
+ for pred_tid in parsed_pred_track.keys():
+ max_iou = 0
+ matched_gt_tid = -1
+ for gt_tid in parsed_gt_track.keys():
+ iou = 0
+ for fi in range(num_f):
+ if parsed_pred_track[pred_tid]["valids"][fi] and parsed_gt_track[gt_tid]["valids"][fi]:
+ iou += compute_iou(parsed_pred_track[pred_tid]["tracked_bbox"][fi], parsed_gt_track[gt_tid]["tracked_bbox"][fi])
+ iou /= num_f
+ if iou > max_iou:
+ max_iou = iou
+ matched_gt_tid = gt_tid
+ if matched_gt_tid == -1:
+ continue
+ _matches[pred_tid] = matched_gt_tid
+
+ print(f"cam_id {cam_id}: {len(_matches)} pred tracks matched")
+ print("_matches:", _matches)
+
+ valid_gt_tids += list(_matches.values())
+ match_pred2gt[cam_id] = _matches
+
+ # get unique valid GT IDs
+ valid_gt_tids = list(set(valid_gt_tids))
+ print(f"valid_gt_tids: {valid_gt_tids}")
+
+ # --------------------------------------------------------------------------
+ # Link GT tracks to predicted tracks
+ # --------------------------------------------------------------------------
+ # For each camera:
+ # 1. Calculate average IoU between GT and predicted tracks
+ # 2. Match each GT track to its best predicted track (if any)
+ #
+ # Key points:
+ # - A GT track can match different predicted tracks across cameras
+ # - The same GT Track may have different predicted track IDs across cameras
+ # due to independent processing of each camera's video
+ # - It connects the consistent GT ID to camera-specific predicted IDs
+ #
+ # Result: We create a map from GT tracks to their best predicted tracks (if any) in each camera
+ match_gt2pred = {}
+ for cam_id in camera_list:
+ parsed_pred_track = parsed_pred_track_all[cam_id]
+ parsed_gt_track = parsed_gt_track_all[cam_id]
+ _matches = {}
+ for gt_tid in valid_gt_tids:
+ if gt_tid not in parsed_gt_track:
+ continue
+ max_iou = 0
+ matched_pred_tid = -1
+ for pred_tid in parsed_pred_track.keys():
+ iou = 0
+ for fi in range(num_f):
+ if parsed_pred_track[pred_tid]["valids"][fi] and parsed_gt_track[gt_tid]["valids"][fi]:
+ iou += compute_iou(parsed_pred_track[pred_tid]["tracked_bbox"][fi], parsed_gt_track[gt_tid]["tracked_bbox"][fi])
+ iou /= num_f
+ if iou > max_iou:
+ max_iou = iou
+ matched_pred_tid = pred_tid
+ _matches[gt_tid] = matched_pred_tid
+
+ print(f"cam_id {cam_id}: {len(_matches)} GT tracks matched")
+ print("_matches:", _matches)
+ match_gt2pred[cam_id] = _matches
+
+ # --------------------------------------------------------------------------
+ # Collect SMPL parameters for each valid GT track
+ # --------------------------------------------------------------------------
+ # This step processes the results of GT-to-pred track matching across multiple cameras:
+ # 1. For each valid GT track ID:
+ # a. Initialize data structures to store SMPL parameters, camera info, and validity masks
+ # b. For each camera:
+ # - Record when the GT track appears in 2D (2DBox_appear_mask)
+ # - If matched to a pred track:
+ # * Collect SMPL parameters (global_orient, body_pose, betas) from the matched pred track
+ # * Collect camera parameters
+ # * Mark frames where data is available (matched_mask)
+ # 2. The collected data retains the multi-camera structure, allowing for later analysis
+ # of track consistency across different views
+ #
+ # This approach allows us to:
+ # - Maintain the GT track ID as the primary identifier
+ # - Collect all available pose data for each GT track from matched predictions
+ # - Preserve the multi-camera nature of the data for subsequent processing steps
+
+ collector = {}
+ # Initialize the collector
+ for gt_tid in valid_gt_tids:
+ collector[gt_tid] = {
+ "smpl": {
+ "global_orient": torch.zeros(num_c, num_f, 1, 3, 3),
+ "body_pose": torch.zeros(num_c, num_f, 23, 3, 3),
+ "betas": torch.zeros(num_c, num_f, 10),
+ },
+ "camera": torch.zeros(num_c, num_f, 3),
+ "2DBox_appear_mask": torch.zeros(num_c, num_f).bool(),
+ "area": torch.zeros(num_c, num_f),
+ "matched_mask": torch.zeros(num_c, num_f).bool(),
+ }
+
+ # collect data
+ for cam_id in camera_list:
+ parsed_pred_track = parsed_pred_track_all[cam_id]
+ parsed_gt_track = parsed_gt_track_all[cam_id]
+ _matches = match_gt2pred[cam_id]
+
+ if len(_matches) > 0:
+ # collect 2DBox_appear_mask
+ for gt_tid in valid_gt_tids:
+ if gt_tid not in parsed_gt_track:
+ continue
+ for fi in range(num_f):
+ if parsed_gt_track[gt_tid]["valids"][fi]:
+ collector[gt_tid]["2DBox_appear_mask"][cam_id, fi] = 1
+ box = parsed_gt_track[gt_tid]["tracked_bbox"][fi]
+ collector[gt_tid]["area"][cam_id, fi] = box[2] * box[3]
+
+ # collect SMPL data
+ for gt_tid in valid_gt_tids:
+ if gt_tid not in parsed_gt_track:
+ continue
+ matched_pred_tid = _matches[gt_tid]
+ if matched_pred_tid == -1:
+ continue
+ for fi in range(num_f):
+ if parsed_gt_track[gt_tid]["valids"][fi] and parsed_pred_track[matched_pred_tid]["valids"][fi]:
+ collector[gt_tid]["smpl"]["global_orient"][cam_id, fi] = parsed_pred_track[matched_pred_tid]["smpl"]["global_orient"][fi].clone()
+ collector[gt_tid]["smpl"]["body_pose"][cam_id, fi] = parsed_pred_track[matched_pred_tid]["smpl"]["body_pose"][fi].clone()
+ collector[gt_tid]["smpl"]["betas"][cam_id, fi] = parsed_pred_track[matched_pred_tid]["smpl"]["betas"][fi].clone()
+ collector[gt_tid]["camera"][cam_id, fi] = parsed_pred_track[matched_pred_tid]["camera"][fi].clone()
+ collector[gt_tid]["matched_mask"][cam_id, fi] = 1
+
+ # update valid mask
+ for gt_tid in valid_gt_tids:
+ appear_mask = collector[gt_tid]["2DBox_appear_mask"]
+ matched_mask = collector[gt_tid]["matched_mask"]
+ # valid mask: 2DBox_appear_mask & matched_mask
+ valid_mask = appear_mask & matched_mask
+ collector[gt_tid]["valid_mask"] = valid_mask
+
+ # Save matched data in a format compatible with 4D-Humans visualization
+ if save_temp:
+ for cam_id in camera_list:
+ pkl_dict = {}
+ for fi in range(num_f):
+ fi_info = {
+ "time": fi,
+ "shot": 0,
+ "tid": [],
+ "tracked_time": [],
+ "mask": None,
+ "bbox": None,
+ "smpl": [],
+ "camera": [],
+ }
+ for gt_tid in valid_gt_tids:
+ if collector[gt_tid]["valid_mask"][cam_id, fi]:
+ fi_info["tid"].append(gt_tid)
+ fi_info["tracked_time"].append(0)
+ fi_info["smpl"].append({
+ "global_orient": collector[gt_tid]["smpl"]["global_orient"][cam_id, fi].numpy(),
+ "body_pose": collector[gt_tid]["smpl"]["body_pose"][cam_id, fi].numpy(),
+ "betas": collector[gt_tid]["smpl"]["betas"][cam_id, fi].numpy(),
+ })
+ fi_info["camera"].append(collector[gt_tid]["camera"][cam_id, fi].numpy())
+ pkl_dict[fi] = fi_info
+ joblib.dump(
+ pkl_dict, os.path.join(temp_dir, f"{cam_id}_matched.pkl")
+ )
+
+ # --------------------------------------------------------------------------
+ # Complete data for all frames where GT 2D boxes exist
+ # --------------------------------------------------------------------------
+ # For each camera and GT track:
+ # 1. Check each frame where a GT 2D box exists (human is visible in scene)
+ # 2. If no matched data exists for that frame:
+ # - This often occurs due to occlusion or tracking failures
+ # - We have the 2D box (know human's position) but lack SMPL params
+ # 3. Interpolate missing data (SMPL params, camera info) from surrounding frames
+ # 4. This ensures continuity in human pose data even when direct detection fails
+ # 5. Result: All frames with 2DBox_appear_mask==True will have complete data
+ #
+ # This step is crucial for handling temporary occlusions or detection misses,
+ # maintaining a consistent track even when the pose predictor fails momentarily.
+ for cam_id in camera_list:
+ for gt_tid in valid_gt_tids:
+ need_fill = collector[gt_tid]["2DBox_appear_mask"][cam_id] & (~collector[gt_tid]["matched_mask"][cam_id])
+ if need_fill.any() and not collector[gt_tid]["matched_mask"][cam_id].sum() == 0:
+ appear_mask = collector[gt_tid]["2DBox_appear_mask"][cam_id]
+ mask = collector[gt_tid]["valid_mask"][cam_id][appear_mask]
+ collector[gt_tid]["matched_mask"][cam_id] = collector[gt_tid]["matched_mask"][cam_id] | appear_mask
+
+ feature_masked = collector[gt_tid]["smpl"]["global_orient"][cam_id, appear_mask]
+ feature_full = interpolate_features(mask, feature_masked, is_rot_mat=True)
+ collector[gt_tid]["smpl"]["global_orient"][cam_id, appear_mask] = feature_full
+
+ feature_masked = collector[gt_tid]["smpl"]["body_pose"][cam_id, appear_mask]
+ feature_full = interpolate_features(mask, feature_masked, is_rot_mat=True)
+ collector[gt_tid]["smpl"]["body_pose"][cam_id, appear_mask] = feature_full
+
+ feature_masked = collector[gt_tid]["smpl"]["betas"][cam_id, appear_mask]
+ feature_full = interpolate_features(mask, feature_masked)
+ collector[gt_tid]["smpl"]["betas"][cam_id, appear_mask] = feature_full
+
+ feature_masked = collector[gt_tid]["camera"][cam_id, appear_mask]
+ feature_full = interpolate_features(mask, feature_masked)
+ collector[gt_tid]["camera"][cam_id, appear_mask] = feature_full
+
+ # Save completed data in a format compatible with 4D-Humans visualization
+ if save_temp:
+ for cam_id in camera_list:
+ pkl_dict = {}
+ for fi in range(num_f):
+ fi_info = {
+ "time": fi,
+ "shot": 0,
+ "tid": [],
+ "tracked_time": [],
+ "mask": None,
+ "bbox": None,
+ "smpl": [],
+ "camera": [],
+ }
+ for gt_tid in valid_gt_tids:
+ if collector[gt_tid]["2DBox_appear_mask"][cam_id, fi]: # the difference is we use 2DBox_appear_mask here
+ fi_info["tid"].append(gt_tid)
+ fi_info["tracked_time"].append(0)
+ fi_info["smpl"].append({
+ "global_orient": collector[gt_tid]["smpl"]["global_orient"][cam_id, fi].numpy(),
+ "body_pose": collector[gt_tid]["smpl"]["body_pose"][cam_id, fi].numpy(),
+ "betas": collector[gt_tid]["smpl"]["betas"][cam_id, fi].numpy(),
+ })
+ fi_info["camera"].append(collector[gt_tid]["camera"][cam_id, fi].numpy())
+ pkl_dict[fi] = fi_info
+ joblib.dump(
+ pkl_dict, os.path.join(temp_dir, f"{cam_id}_completed.pkl")
+ )
+
+ # --------------------------------------------------------------------------
+ # Merge and refine SMPL data across all cameras
+ # --------------------------------------------------------------------------
+ # 1. Initialize merged_collector:
+ # - For each GT track, create a structure to hold combined data from all cameras
+ # 2. Merge data for each GT track:
+ # a. Combine matched_mask from all cameras
+ # b. Select best SMPL parameters for each frame:
+ # - If only one camera detects the person, use that camera's data
+ # - If multiple cameras detect (e.g., person moving across camera views):
+ # * Choose the camera with largest detection area
+ # * This typically provides the most accurate pose estimation
+ # c. Set must_appear_mask where the person is visible in any camera
+ # d. Store selected SMPL parameters, camera info, and 2D box appearance masks
+ merged_collector = {}
+ for gt_tid in valid_gt_tids:
+ merged_collector[gt_tid] = {
+ "smpl": {
+ "global_orient": torch.zeros(num_f, 1, 3, 3),
+ "body_pose": torch.zeros(num_f, 23, 3, 3),
+ "betas": torch.zeros(num_f, 10),
+ },
+ "camera": torch.zeros(3, num_f, 3), # 3 cameras
+ "2DBox_appear_mask": torch.zeros(3, num_f).bool(),
+ "matched_mask": torch.zeros(num_f).bool(), # matched_mask for all cameras
+ "must_appear_mask": torch.zeros(num_f).bool(),
+ "selected_cam_idx": torch.zeros(num_f).long() - 1, # -1 means not selected
+ }
+ matched_sum = collector[gt_tid]["matched_mask"].sum(0)
+ merged_collector[gt_tid]["matched_mask"] = matched_sum > 0
+
+ # if matched_sum == 1 means only one cam get the smpl, we just select the smpl from that cam
+ matched_cam_id = collector[gt_tid]["matched_mask"].float().argmax(dim=0)
+ # if matched_sum > 1 means multiple cams get the smpl, we select the smpl from the cam with the largest area
+ area = collector[gt_tid]["area"]
+ max_cam_id = area.argmax(0)
+ best_cam_id = matched_cam_id * (matched_sum == 1)+ max_cam_id * (matched_sum > 1)
+
+ mask = collector[gt_tid]["2DBox_appear_mask"].sum(0) > 0
+ merged_collector[gt_tid]["must_appear_mask"] = mask
+ masked_best_cam_id = best_cam_id[mask]
+
+ merged_collector[gt_tid]["smpl"]["global_orient"][mask] = collector[gt_tid]["smpl"]["global_orient"][:, mask][masked_best_cam_id, torch.arange(masked_best_cam_id.shape[0])]
+ merged_collector[gt_tid]["smpl"]["body_pose"][mask] = collector[gt_tid]["smpl"]["body_pose"][:, mask][masked_best_cam_id, torch.arange(masked_best_cam_id.shape[0])]
+ merged_collector[gt_tid]["smpl"]["betas"][mask] = collector[gt_tid]["smpl"]["betas"][:, mask][masked_best_cam_id, torch.arange(masked_best_cam_id.shape[0])]
+ merged_collector[gt_tid]["camera"] = collector[gt_tid]["camera"]
+ merged_collector[gt_tid]["2DBox_appear_mask"] = collector[gt_tid]["2DBox_appear_mask"]
+ merged_collector[gt_tid]["selected_cam_idx"][mask] = masked_best_cam_id
+
+ # 3. Fill gaps in SMPL data:
+ # - For frames where must_appear_mask is true but matched_mask is false:
+ # * Interpolate SMPL parameters (global_orient, body_pose, betas)
+ # * Interpolate camera info and selected camera index
+ # - This addresses frames where the person is visible but not detected
+ for gt_tid in valid_gt_tids:
+ need_fill = merged_collector[gt_tid]["must_appear_mask"] & (~merged_collector[gt_tid]["matched_mask"])
+ if need_fill.any():
+ final_complete_mask = merged_collector[gt_tid]["must_appear_mask"]
+ mask = merged_collector[gt_tid]["matched_mask"][final_complete_mask]
+
+ feature_masked = merged_collector[gt_tid]["smpl"]["global_orient"][final_complete_mask]
+ feature_full = interpolate_features(mask, feature_masked, is_rot_mat=True)
+ merged_collector[gt_tid]["smpl"]["global_orient"][final_complete_mask] = feature_full
+
+ feature_masked = merged_collector[gt_tid]["smpl"]["body_pose"][final_complete_mask]
+ feature_full = interpolate_features(mask, feature_masked, is_rot_mat=True)
+ merged_collector[gt_tid]["smpl"]["body_pose"][final_complete_mask] = feature_full
+
+ feature_masked = merged_collector[gt_tid]["smpl"]["betas"][final_complete_mask]
+ feature_full = interpolate_features(mask, feature_masked)
+ merged_collector[gt_tid]["smpl"]["betas"][final_complete_mask] = feature_full
+
+ feature_masked = merged_collector[gt_tid]["camera"][:, final_complete_mask]
+ feature_full = interpolate_features(mask, feature_masked)
+ merged_collector[gt_tid]["camera"][:, final_complete_mask] = feature_full
+
+ feature_masked = merged_collector[gt_tid]["selected_cam_idx"][final_complete_mask]
+ feature_full = interpolate_features(mask, feature_masked)
+ merged_collector[gt_tid]["selected_cam_idx"][final_complete_mask] = feature_full.round()
+
+ # 4. Fill breaks in visibility:
+ # - Detect short gaps in must_appear_mask sequence
+ # - For these gaps:
+ # * Set must_appear_mask to true
+ # * Interpolate SMPL parameters, camera info, and selected camera index
+ # - This creates continuity in tracks, filling brief disappearances
+ for gt_tid in valid_gt_tids:
+ need_fill = torch.tensor(
+ detect_breaks_mask(merged_collector[gt_tid]["must_appear_mask"])
+ )
+ if need_fill.any():
+ final_complete_mask = need_fill | merged_collector[gt_tid]["must_appear_mask"]
+ mask = merged_collector[gt_tid]["must_appear_mask"][final_complete_mask]
+
+ feature_masked = merged_collector[gt_tid]["smpl"]["global_orient"][final_complete_mask]
+ feature_full = interpolate_features(mask, feature_masked, is_rot_mat=True)
+ merged_collector[gt_tid]["smpl"]["global_orient"][final_complete_mask] = feature_full
+
+ feature_masked = merged_collector[gt_tid]["smpl"]["body_pose"][final_complete_mask]
+ feature_full = interpolate_features(mask, feature_masked, is_rot_mat=True)
+ merged_collector[gt_tid]["smpl"]["body_pose"][final_complete_mask] = feature_full
+
+ feature_masked = merged_collector[gt_tid]["smpl"]["betas"][final_complete_mask]
+ feature_full = interpolate_features(mask, feature_masked)
+ merged_collector[gt_tid]["smpl"]["betas"][final_complete_mask] = feature_full
+
+ feature_masked = merged_collector[gt_tid]["camera"][:, final_complete_mask]
+ feature_full = interpolate_features(mask, feature_masked)
+ merged_collector[gt_tid]["camera"][:, final_complete_mask] = feature_full
+
+ feature_masked = merged_collector[gt_tid]["selected_cam_idx"][final_complete_mask]
+ feature_full = interpolate_features(mask, feature_masked)
+ merged_collector[gt_tid]["selected_cam_idx"][final_complete_mask] = feature_full.round()
+
+ merged_collector[gt_tid]["must_appear_mask"] = final_complete_mask
+
+ # update valid_mask
+ for gt_tid in valid_gt_tids:
+ merged_collector[gt_tid]["valid_mask"] = merged_collector[gt_tid]["must_appear_mask"].clone()
+
+ # Save merged data in a format compatible with 4D-Humans visualization
+ if save_temp:
+ pkl_dict = {}
+ for cam_id in camera_list:
+ for fi in range(num_f):
+ fi_info = {
+ "time": fi,
+ "shot": 0,
+ "tid": [],
+ "tracked_time": [],
+ "mask": None,
+ "bbox": None,
+ "smpl": [],
+ "camera": [],
+ }
+ for gt_tid in valid_gt_tids:
+ if merged_collector[gt_tid]["2DBox_appear_mask"][cam_id, fi]:
+ fi_info["tid"].append(gt_tid)
+ fi_info["tracked_time"].append(0)
+ fi_info["smpl"].append({
+ "global_orient": merged_collector[gt_tid]["smpl"]["global_orient"][fi].numpy(),
+ "body_pose": merged_collector[gt_tid]["smpl"]["body_pose"][fi].numpy(),
+ "betas": merged_collector[gt_tid]["smpl"]["betas"][fi].numpy(),
+ })
+ fi_info["camera"].append(merged_collector[gt_tid]["camera"][cam_id, fi].numpy())
+ pkl_dict[fi] = fi_info
+ joblib.dump(
+ pkl_dict, os.path.join(temp_dir, f"{cam_id}_merged.pkl")
+ )
+
+ # pop out 2DBox_appear_mask and selected_cam_idx and camera
+ for gt_tid in valid_gt_tids:
+ merged_collector[gt_tid].pop("2DBox_appear_mask")
+ merged_collector[gt_tid].pop("camera")
+
+ # check if nan or inf in smpl, for debug use
+ for gt_tid in valid_gt_tids:
+ for k, v in merged_collector[gt_tid]["smpl"].items():
+ if v.isnan().any() or v.isinf().any():
+ raise ValueError(f"gt_tid {gt_tid} has nan or inf in {k}")
+
+ return merged_collector
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/waymo/waymo_download.py b/worldbench/third_party/drivestudio/datasets/waymo/waymo_download.py
new file mode 100644
index 0000000..0fb2bb6
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/waymo/waymo_download.py
@@ -0,0 +1,87 @@
+import argparse
+import os
+import subprocess
+from concurrent.futures import ThreadPoolExecutor
+from typing import List
+
+
+def download_file(filename, target_dir, source):
+ result = subprocess.run(
+ [
+ "gsutil",
+ "cp",
+ "-n",
+ f"{source}/{filename}.tfrecord",
+ target_dir,
+ ],
+ capture_output=True, # To capture stderr and stdout for detailed error information
+ text=True,
+ )
+
+ # Check the return code of the gsutil command
+ if result.returncode != 0:
+ raise Exception(
+ result.stderr
+ ) # Raise an exception with the error message from the gsutil command
+
+
+def download_files(
+ file_names: List[str],
+ target_dir: str,
+ source: str = "gs://waymo_open_dataset_scene_flow/train",
+) -> None:
+ """
+ Downloads a list of files from a given source to a target directory using multiple threads.
+
+ Args:
+ file_names (List[str]): A list of file names to download.
+ target_dir (str): The target directory to save the downloaded files.
+ source (str, optional): The source directory to download the files from. Defaults to "gs://waymo_open_dataset_scene_flow/train".
+ """
+ # Get the total number of file_names
+ total_files = len(file_names)
+
+ # Use ThreadPoolExecutor to manage concurrent downloads
+ with ThreadPoolExecutor(max_workers=10) as executor:
+ futures = [
+ executor.submit(download_file, filename, target_dir, source)
+ for filename in file_names
+ ]
+
+ for counter, future in enumerate(futures, start=1):
+ # Wait for the download to complete and handle any exceptions
+ try:
+ # inspects the result of the future and raises an exception if one occurred during execution
+ future.result()
+ print(f"[{counter}/{total_files}] Downloaded successfully!")
+ except Exception as e:
+ print(f"[{counter}/{total_files}] Failed to download. Error: {e}")
+
+
+if __name__ == "__main__":
+ print("note: `gcloud auth login` is required before running this script")
+ print("Downloading Waymo dataset from Google Cloud Storage...")
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--target_dir",
+ type=str,
+ default="data/waymo/raw",
+ help="Path to the target directory",
+ )
+ parser.add_argument(
+ "--scene_ids", type=int, nargs="+", help="scene ids to download"
+ )
+ parser.add_argument(
+ "--split_file", type=str, default=None, help="split file in data/waymo_splits"
+ )
+ args = parser.parse_args()
+ os.makedirs(args.target_dir, exist_ok=True)
+ total_list = open("data/waymo_train_list.txt", "r").readlines()
+ if args.split_file is None:
+ file_names = [total_list[i].strip() for i in args.scene_ids]
+ else:
+ # parse the split file
+ split_file = open(args.split_file, "r").readlines()[1:]
+ scene_ids = [int(line.strip().split(",")[0]) for line in split_file]
+ file_names = [total_list[i].strip() for i in scene_ids]
+ download_files(file_names, args.target_dir)
diff --git a/worldbench/third_party/drivestudio/datasets/waymo/waymo_human_utils.py b/worldbench/third_party/drivestudio/datasets/waymo/waymo_human_utils.py
new file mode 100644
index 0000000..c64cb06
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/waymo/waymo_human_utils.py
@@ -0,0 +1,221 @@
+from typing import List
+import os
+import cv2
+import json
+import logging
+import numpy as np
+from tqdm import tqdm
+
+from utils.geometry import (
+ get_corners,
+ project_camera_points_to_image
+)
+from .waymo_sourceloader import (
+ SMPLNODE_CLASSES,
+ OPENCV2DATASET,
+ AVAILABLE_CAM_LIST,
+)
+
+logger = logging.getLogger()
+
+CAMERA_LIST = AVAILABLE_CAM_LIST
+
+# TODO(ziyu): Refactor human utils scripts for all datasets into a single, unified module
+# Implement dataset-specific loaders for intrinsic and extrinsic parameters
+# This will improve code reusability and make it easier to add support for new datasets
+def project_human_boxes(
+ scene_dir: str,
+ camera_list: List[int],
+ save_temp=True,
+ verbose=False,
+ narrow_width_ratio=0.2,
+ fps=12
+):
+ """Project human boxes to 2D image space and save the results to pkl files
+
+ Args:
+ scene_dir: str, path to the scene directory
+ camera_list: List[int], a list of camera ids to be processed
+ save_temp: bool, whether to save the intermediate results
+ verbose: bool, whether to visualize the projected boxes
+ narrow_width_ratio: sometimes the projected boxes are too wide
+ we can narrow them with this ratio to get more accurate results
+ fps: int, FPS for the visualization video
+
+ Returns:
+ collector_all: dict, a dictionary containing the projected boxes for each camera
+ """
+ # check if the necessary directories exist
+ images_dir = f'{scene_dir}/images'
+ poses_dir = f'{scene_dir}/ego_pose'
+ extrinsics_dir = f'{scene_dir}/extrinsics'
+ intrinsics_dir = f'{scene_dir}/intrinsics'
+ instances_dir = f'{scene_dir}/instances'
+ valid_paths = [images_dir, poses_dir, extrinsics_dir, intrinsics_dir, instances_dir]
+ for path in valid_paths:
+ assert os.path.exists(path), \
+ f"Path {path} does not exist, you need to run waymo preprocess to generate the necessary files"
+
+ # create directories for saving the results
+ save_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks'
+ if not os.path.exists(save_dir):
+ os.makedirs(save_dir)
+
+ if verbose:
+ # create directories for saving the intermediate visualization results
+ video_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks/vis'
+ if not os.path.exists(video_dir):
+ os.makedirs(video_dir)
+ per_human_img_dir = f'{scene_dir}/humanpose/temp/Pedes_GTTracks/vis/images'
+ if not os.path.exists(per_human_img_dir):
+ os.makedirs(per_human_img_dir)
+
+ # load instances and frame infos
+ frame_infos = json.load(
+ open(f'{instances_dir}/frame_instances.json')
+ )
+ instances_meta = json.load(
+ open(f'{instances_dir}/instances_info.json')
+ )
+
+ collector_all = {}
+ # iterate over each camera
+ for cam_id in camera_list:
+ # check if already processed
+ pkl_path = os.path.join(save_dir, f"{cam_id}.pkl")
+ if os.path.exists(pkl_path):
+ collector_all[cam_id] = json.load(open(pkl_path))
+ logger.info(f"Results for camera {cam_id} already exists at {pkl_path}")
+ continue
+
+ if verbose:
+ per_cam_vis_dir = os.path.join(per_human_img_dir, f"{cam_id}")
+ if not os.path.exists(per_cam_vis_dir):
+ os.makedirs(per_cam_vis_dir)
+
+ collector = {}
+ frames = []
+ for frame_id, frame_ins_list in frame_infos.items():
+ frame_id = int(frame_id)
+
+ # define empty instance collector for each frame
+ frame_collector = {
+ "gt_bbox": [],
+ "extra_data": {
+ "gt_track_id": [],
+ "gt_class": [],
+ }
+ }
+
+ # load extrinsic
+ cam_to_ego = np.loadtxt(os.path.join(extrinsics_dir, f"{cam_id}.txt"))
+ cam_to_ego = cam_to_ego @ OPENCV2DATASET
+ ego_to_world = np.loadtxt(os.path.join(poses_dir, f"{str(frame_id).zfill(3)}.txt"))
+ cam2world = ego_to_world @ cam_to_ego
+
+ # load intrinsic
+ Ks = np.loadtxt(os.path.join(intrinsics_dir, f"{cam_id}.txt"))
+ fx, fy, cx, cy = Ks[0], Ks[1], Ks[2], Ks[3]
+ # k1, k2, p1, p2, k3 = Ks[4], Ks[5], Ks[6], Ks[7], Ks[8]
+ intrinsic = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float32)
+
+ # load image
+ ori_image = cv2.imread(
+ os.path.join(images_dir, f"{str(frame_id).zfill(3)}_{cam_id}.jpg")
+ )
+ image = ori_image.copy()
+ H, W = image.shape[:2]
+
+ # iterate over each instance in the frame
+ if len(frame_ins_list) > 0:
+ # if there are pedestrians in this frame, project the boxes to the image
+ image_plotted = image.copy()
+ for instance_id in frame_ins_list:
+ ins = instances_meta[str(instance_id)]
+
+ if ins["class_name"] not in SMPLNODE_CLASSES:
+ continue
+
+ ins_anno = ins["frame_annotations"]
+ index = ins_anno['frame_idx'].index(frame_id)
+ obj_to_world = np.array(ins_anno['obj_to_world'][index])
+ l, w, h = ins_anno['box_size'][index]
+
+ # get box corners in object space
+ corners = get_corners(l, w, h)
+ # transform box corners to world space
+ corners_world = obj_to_world[:3, :3] @ corners + obj_to_world[:3, 3:4]
+ # transform box corners to image space
+ world2cam = np.linalg.inv(cam2world)
+ corners_cam = world2cam[:3, :3] @ corners_world + world2cam[:3, 3:4]
+ cam_points, depth = project_camera_points_to_image(corners_cam.T, intrinsic)
+
+ x_min, y_min = np.min(cam_points, axis=0)
+ x_max, y_max = np.max(cam_points, axis=0)
+ # clip left and right with this ratio
+ if narrow_width_ratio > 0.:
+ length = x_max - x_min
+ x_min += length * narrow_width_ratio
+ x_max -= length * narrow_width_ratio
+
+ # clip the box to the image
+ original_area = (x_max - x_min) * (y_max - y_min)
+ x_min, x_max = np.clip(x_min, 0, W), np.clip(x_max, 0, W)
+ y_min, y_max = np.clip(y_min, 0, H), np.clip(y_max, 0, H)
+ new_area = (x_max - x_min) * (y_max - y_min)
+
+ # filter out boxes that are too small or too large
+ behind = depth.max() < 0
+ too_small = new_area < W * H * (0.03)**2
+ too_large = new_area > W * H / 1.1
+ too_far = np.linalg.norm(obj_to_world[:3, 3] - cam2world[:3, 3]) > 40
+ clip_large = new_area / original_area < 1/3
+ if too_small or too_large or clip_large or behind or too_far:
+ continue
+
+ # gt box on image
+ gt_box = [x_min, y_min, x_max - x_min, y_max - y_min]
+
+ # save the projected box to the collector
+ frame_collector["gt_bbox"].append(gt_box)
+ frame_collector["extra_data"]["gt_track_id"].append(instance_id)
+ frame_collector["extra_data"]["gt_class"].append([0])
+
+ if verbose:
+ # visualize the projected boxes of ONE instance
+ raw_image = cv2.rectangle(
+ ori_image.copy(), (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 2
+ )
+ raw_image_path = os.path.join(per_cam_vis_dir, f"{frame_id}_{instance_id}.jpg")
+ cv2.imwrite(raw_image_path, raw_image)
+
+ # add this instance to the image
+ image_plotted = cv2.rectangle(
+ image, (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 2
+ )
+
+ if verbose:
+ frames.append(image_plotted)
+ else:
+ # if no instance in this frame, just save the original image
+ if verbose:
+ frames.append(ori_image)
+
+ collector[frame_id] = frame_collector
+
+ if verbose:
+ height, width = frames[0].shape[:2]
+ output_path = os.path.join(video_dir, f"cam_{cam_id}.mp4")
+ fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+ out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+ for frame in tqdm(frames, desc=f"Writing video for camera {cam_id}"):
+ out.write(frame)
+ out.release()
+
+ if save_temp:
+ # save collector to pkl
+ json.dump(collector, open(pkl_path, "w"))
+
+ collector_all[cam_id] = collector
+
+ return collector_all
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/waymo/waymo_preprocess.py b/worldbench/third_party/drivestudio/datasets/waymo/waymo_preprocess.py
new file mode 100644
index 0000000..71e5ede
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/waymo/waymo_preprocess.py
@@ -0,0 +1,558 @@
+# Acknowledgement:
+# 1. https://github.com/open-mmlab/mmdetection3d/blob/main/tools/dataset_converters/waymo_converter.py
+# 2. https://github.com/leolyj/DCA-SRSFE/blob/main/data_preprocess/Waymo/generate_flow.py
+try:
+ from waymo_open_dataset import dataset_pb2
+except ImportError:
+ raise ImportError(
+ 'Please run "pip install waymo-open-dataset-tf-2-6-0" '
+ ">1.4.5 to install the official devkit first."
+ )
+
+import json
+import os
+
+import numpy as np
+import tensorflow as tf
+from PIL import Image
+from tqdm import tqdm
+from waymo_open_dataset import label_pb2
+from waymo_open_dataset.protos import camera_segmentation_pb2 as cs_pb2
+from waymo_open_dataset.utils import box_utils
+from waymo_open_dataset.utils.frame_utils import parse_range_image_and_camera_projection
+
+from datasets.tools.multiprocess_utils import track_parallel_progress
+from .waymo_utils import (
+ parse_range_image_flow_and_camera_projection,
+ convert_range_image_to_point_cloud_flow,
+ project_vehicle_to_image,
+ get_ground_np
+)
+
+MOVEABLE_OBJECTS_IDS = [
+ cs_pb2.CameraSegmentation.TYPE_CAR,
+ cs_pb2.CameraSegmentation.TYPE_TRUCK,
+ cs_pb2.CameraSegmentation.TYPE_BUS,
+ cs_pb2.CameraSegmentation.TYPE_OTHER_LARGE_VEHICLE,
+ cs_pb2.CameraSegmentation.TYPE_BICYCLE,
+ cs_pb2.CameraSegmentation.TYPE_MOTORCYCLE,
+ cs_pb2.CameraSegmentation.TYPE_TRAILER,
+ cs_pb2.CameraSegmentation.TYPE_PEDESTRIAN,
+ cs_pb2.CameraSegmentation.TYPE_CYCLIST,
+ cs_pb2.CameraSegmentation.TYPE_MOTORCYCLIST,
+ cs_pb2.CameraSegmentation.TYPE_BIRD,
+ cs_pb2.CameraSegmentation.TYPE_GROUND_ANIMAL,
+ cs_pb2.CameraSegmentation.TYPE_PEDESTRIAN_OBJECT,
+]
+
+WAYMO_CLASSES = ['unknown', 'Vehicle', 'Pedestrian', 'Sign', 'Cyclist']
+# TODO(ziyu): consider all dynamic classes
+WAYMO_DYNAMIC_CLASSES = ['Vehicle', 'Pedestrian', 'Cyclist']
+WAYMO_HUMAN_CLASSES = ['Pedestrian', 'Cyclist']
+WAYMO_VEHICLE_CLASSES = ['Vehicle']
+
+class WaymoProcessor(object):
+ """Process Waymo dataset.
+
+ Args:
+ load_dir (str): Directory to load waymo raw data.
+ save_dir (str): Directory to save data in KITTI format.
+ prefix (str): Prefix of filename.
+ workers (int, optional): Number of workers for the parallel process.
+ Defaults to 64.
+ Defaults to False.
+ save_cam_sync_labels (bool, optional): Whether to save cam sync labels.
+ Defaults to True.
+ """
+
+ def __init__(
+ self,
+ load_dir,
+ save_dir,
+ prefix,
+ process_keys=[
+ "images",
+ "lidar",
+ "calib",
+ "pose",
+ "dynamic_masks",
+ "objects"
+ ],
+ process_id_list=None,
+ workers=64,
+ ):
+ self.filter_no_label_zone_points = True
+
+ # Only data collected in specific locations will be converted
+ # If set None, this filter is disabled
+ # Available options: location_sf (main dataset)
+ self.selected_waymo_locations = None
+ self.save_track_id = False
+ self.process_id_list = process_id_list
+ self.process_keys = process_keys
+ print("will process keys: ", self.process_keys)
+
+ # turn on eager execution for older tensorflow versions
+ if int(tf.__version__.split(".")[0]) < 2:
+ tf.enable_eager_execution()
+
+ # keep the order defined by the official protocol
+ self.cam_list = [
+ "_FRONT",
+ "_FRONT_LEFT",
+ "_FRONT_RIGHT",
+ "_SIDE_LEFT",
+ "_SIDE_RIGHT",
+ ]
+ self.lidar_list = ["TOP", "FRONT", "SIDE_LEFT", "SIDE_RIGHT", "REAR"]
+
+ self.load_dir = load_dir
+ self.save_dir = f"{save_dir}/{prefix}"
+ self.workers = int(workers)
+ # a list of tfrecord pathnames
+ training_files = open("data/waymo_train_list.txt").read().splitlines()
+ self.tfrecord_pathnames = [
+ f"{self.load_dir}/{f}.tfrecord" for f in training_files
+ ]
+ # self.tfrecord_pathnames = sorted(glob(join(self.load_dir, "*.tfrecord")))
+ self.create_folder()
+
+ def convert(self):
+ """Convert action."""
+ print("Start converting ...")
+ if self.process_id_list is None:
+ id_list = range(len(self))
+ else:
+ id_list = self.process_id_list
+ track_parallel_progress(self.convert_one, id_list, self.workers)
+ print("\nFinished ...")
+
+ def convert_one(self, file_idx):
+ """Convert action for single file.
+
+ Args:
+ file_idx (int): Index of the file to be converted.
+ """
+ pathname = self.tfrecord_pathnames[file_idx]
+ dataset = tf.data.TFRecordDataset(pathname, compression_type="")
+ num_frames = sum(1 for _ in dataset)
+ for frame_idx, data in enumerate(
+ tqdm(dataset, desc=f"File {file_idx}", total=num_frames, dynamic_ncols=True)
+ ):
+ frame = dataset_pb2.Frame()
+ frame.ParseFromString(bytearray(data.numpy()))
+ if (
+ self.selected_waymo_locations is not None
+ and frame.context.stats.location not in self.selected_waymo_locations
+ ):
+ continue
+ if "images" in self.process_keys:
+ self.save_image(frame, file_idx, frame_idx)
+ if "calib" in self.process_keys:
+ self.save_calib(frame, file_idx, frame_idx)
+ if "lidar" in self.process_keys:
+ self.save_lidar(frame, file_idx, frame_idx)
+ if "pose" in self.process_keys:
+ self.save_pose(frame, file_idx, frame_idx)
+ if "dynamic_masks" in self.process_keys:
+ self.save_dynamic_mask(frame, file_idx, frame_idx, class_valid='all')
+ self.save_dynamic_mask(frame, file_idx, frame_idx, class_valid='human')
+ self.save_dynamic_mask(frame, file_idx, frame_idx, class_valid='vehicle')
+ if frame_idx == 0:
+ self.save_interested_labels(frame, file_idx)
+ if "objects" in self.process_keys:
+ instances_info, frame_instances = self.save_objects(dataset)
+
+ # Save instances info and frame instances
+ object_info_dir = f"{self.save_dir}/{str(file_idx).zfill(3)}/instances"
+ with open(f"{object_info_dir}/instances_info.json", "w") as fp:
+ json.dump(instances_info, fp, indent=4)
+ with open(f"{object_info_dir}/frame_instances.json", "w") as fp:
+ json.dump(frame_instances, fp, indent=4)
+
+ def __len__(self):
+ """Length of the filename list."""
+ return len(self.tfrecord_pathnames)
+
+ def save_interested_labels(self, frame, file_idx):
+ """
+ Saves the interested labels of a given frame to a JSON file.
+
+ Args:
+ frame: A `Frame` object containing the labels to be saved.
+ file_idx: An integer representing the index of the file to be saved.
+
+ Returns:
+ None
+ """
+ frame_data = {
+ "time_of_day": frame.context.stats.time_of_day,
+ "location": frame.context.stats.location,
+ "weather": frame.context.stats.weather,
+ }
+ object_type_name = lambda x: label_pb2.Label.Type.Name(x)
+ object_counts = {
+ object_type_name(x.type): x.count
+ for x in frame.context.stats.camera_object_counts
+ }
+ frame_data.update(object_counts)
+ # write as json
+ with open(
+ f"{self.save_dir}/{str(file_idx).zfill(3)}/frame_info.json",
+ "w",
+ ) as fp:
+ json.dump(frame_data, fp)
+
+ def save_image(self, frame, file_idx, frame_idx):
+ """Parse and save the images in jpg format.
+
+ Args:
+ frame (:obj:`Frame`): Open dataset frame proto.
+ file_idx (int): Current file index.
+ frame_idx (int): Current frame index.
+ """
+ for img in frame.images:
+ img_path = (
+ f"{self.save_dir}/{str(file_idx).zfill(3)}/images/"
+ + f"{str(frame_idx).zfill(3)}_{str(img.name - 1)}.jpg"
+ )
+ with open(img_path, "wb") as fp:
+ fp.write(img.image)
+
+ def save_calib(self, frame, file_idx, frame_idx):
+ """Parse and save the calibration data.
+
+ Args:
+ frame (:obj:`Frame`): Open dataset frame proto.
+ file_idx (int): Current file index.
+ frame_idx (int): Current frame index.
+ """
+ # waymo front camera to kitti reference camera
+ extrinsics = []
+ intrinsics = []
+ for camera in frame.context.camera_calibrations:
+ # extrinsic parameters
+ extrinsic = np.array(camera.extrinsic.transform).reshape(4, 4)
+ intrinsic = list(camera.intrinsic)
+ extrinsics.append(extrinsic)
+ intrinsics.append(intrinsic)
+ # all camera ids are saved as id-1 in the result because
+ # camera 0 is unknown in the proto
+ for i in range(5):
+ np.savetxt(
+ f"{self.save_dir}/{str(file_idx).zfill(3)}/extrinsics/"
+ + f"{str(i)}.txt",
+ extrinsics[i],
+ )
+ np.savetxt(
+ f"{self.save_dir}/{str(file_idx).zfill(3)}/intrinsics/"
+ + f"{str(i)}.txt",
+ intrinsics[i],
+ )
+
+ def save_lidar(self, frame, file_idx, frame_idx):
+ """Parse and save the lidar data in psd format.
+
+ Args:
+ frame (:obj:`Frame`): Open dataset frame proto.
+ file_idx (int): Current file index.
+ frame_idx (int): Current frame index.
+ """
+ (
+ range_images,
+ camera_projections,
+ seg_labels,
+ range_image_top_pose,
+ ) = parse_range_image_and_camera_projection(frame)
+
+ # https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/protos/segmentation.proto
+ if range_image_top_pose is None:
+ # the camera only split doesn't contain lidar points.
+ return
+
+ # collect first return only
+ range_images_flow, _, _ = parse_range_image_flow_and_camera_projection(frame)
+ (
+ origins,
+ points,
+ flows,
+ cp_points,
+ intensity,
+ elongation,
+ laser_ids,
+ ) = convert_range_image_to_point_cloud_flow(
+ frame,
+ range_images,
+ range_images_flow,
+ camera_projections,
+ range_image_top_pose,
+ ri_index=0,
+ )
+ origins = np.concatenate(origins, axis=0)
+ points = np.concatenate(points, axis=0)
+ ground_label = get_ground_np(points)
+ intensity = np.concatenate(intensity, axis=0)
+ elongation = np.concatenate(elongation, axis=0)
+ laser_ids = np.concatenate(laser_ids, axis=0)
+
+ # -1: no-flow-label, the point has no flow information.
+ # 0: unlabeled or "background,", i.e., the point is not contained in a
+ # bounding box.
+ # 1: vehicle, i.e., the point corresponds to a vehicle label box.
+ # 2: pedestrian, i.e., the point corresponds to a pedestrian label box.
+ # 3: sign, i.e., the point corresponds to a sign label box.
+ # 4: cyclist, i.e., the point corresponds to a cyclist label box.
+ flows = np.concatenate(flows, axis=0)
+
+ point_cloud = np.column_stack(
+ (
+ origins,
+ points,
+ flows,
+ ground_label,
+ intensity,
+ elongation,
+ laser_ids,
+ )
+ )
+ pc_path = (
+ f"{self.save_dir}/"
+ + f"{str(file_idx).zfill(3)}/lidar/{str(frame_idx).zfill(3)}.bin"
+ )
+ point_cloud.astype(np.float32).tofile(pc_path)
+
+ def save_pose(self, frame, file_idx, frame_idx):
+ """Parse and save the pose data.
+
+ Note that SDC's own pose is not included in the regular training
+ of KITTI dataset. KITTI raw dataset contains ego motion files
+ but are not often used. Pose is important for algorithms that
+ take advantage of the temporal information.
+
+ Args:
+ frame (:obj:`Frame`): Open dataset frame proto.
+ file_idx (int): Current file index.
+ frame_idx (int): Current frame index.
+ """
+ pose = np.array(frame.pose.transform).reshape(4, 4)
+ np.savetxt(
+ f"{self.save_dir}/{str(file_idx).zfill(3)}/ego_pose/"
+ + f"{str(frame_idx).zfill(3)}.txt",
+ pose,
+ )
+
+ def save_dynamic_mask(self, frame, file_idx, frame_idx, class_valid='all'):
+ assert class_valid in ['all', 'human', 'vehicle'], "Invalid class valid"
+ if class_valid == 'all':
+ VALID_CLASSES = WAYMO_DYNAMIC_CLASSES
+ elif class_valid == 'human':
+ VALID_CLASSES = WAYMO_HUMAN_CLASSES
+ elif class_valid == 'vehicle':
+ VALID_CLASSES = WAYMO_VEHICLE_CLASSES
+ mask_dir = f"{self.save_dir}/{str(file_idx).zfill(3)}/dynamic_masks/{class_valid}"
+ if not os.path.exists(mask_dir):
+ os.makedirs(mask_dir)
+
+ """Parse and save the segmentation data.
+
+ Args:
+ frame (:obj:`Frame`): Open dataset frame proto.
+ file_idx (int): Current file index.
+ frame_idx (int): Current frame index.
+ """
+ for img in frame.images:
+ # dynamic_mask
+ img_path = (
+ f"{self.save_dir}/{str(file_idx).zfill(3)}/images/"
+ + f"{str(frame_idx).zfill(3)}_{str(img.name - 1)}.jpg"
+ )
+ img_shape = np.array(Image.open(img_path))
+ dynamic_mask = np.zeros_like(img_shape, dtype=np.float32)[..., 0]
+
+ filter_available = any(
+ [label.num_top_lidar_points_in_box > 0 for label in frame.laser_labels]
+ )
+ calibration = next(
+ cc for cc in frame.context.camera_calibrations if cc.name == img.name
+ )
+ for label in frame.laser_labels:
+ # camera_synced_box is not available for the data with flow.
+ # box = label.camera_synced_box
+
+ class_name = WAYMO_CLASSES[label.type]
+ if class_name not in VALID_CLASSES:
+ continue
+
+ box = label.box
+ meta = label.metadata
+ speed = np.linalg.norm([meta.speed_x, meta.speed_y])
+ if not box.ByteSize():
+ continue # Filter out labels that do not have a camera_synced_box.
+ if (filter_available and not label.num_top_lidar_points_in_box) or (
+ not filter_available and not label.num_lidar_points_in_box
+ ):
+ continue # Filter out likely occluded objects.
+
+ # Retrieve upright 3D box corners.
+ box_coords = np.array(
+ [
+ [
+ box.center_x,
+ box.center_y,
+ box.center_z,
+ box.length,
+ box.width,
+ box.height,
+ box.heading,
+ ]
+ ]
+ )
+ corners = box_utils.get_upright_3d_box_corners(box_coords)[
+ 0
+ ].numpy() # [8, 3]
+
+ # Project box corners from vehicle coordinates onto the image.
+ projected_corners = project_vehicle_to_image(
+ frame.pose, calibration, corners
+ )
+ u, v, ok = projected_corners.transpose()
+ ok = ok.astype(bool)
+
+ # Skip object if any corner projection failed. Note that this is very
+ # strict and can lead to exclusion of some partially visible objects.
+ if not all(ok):
+ continue
+ u = u[ok]
+ v = v[ok]
+
+ # Clip box to image bounds.
+ u = np.clip(u, 0, calibration.width)
+ v = np.clip(v, 0, calibration.height)
+
+ if u.max() - u.min() == 0 or v.max() - v.min() == 0:
+ continue
+
+ # Draw projected 2D box onto the image.
+ xy = (u.min(), v.min())
+ width = u.max() - u.min()
+ height = v.max() - v.min()
+ # max pooling
+ dynamic_mask[
+ int(xy[1]) : int(xy[1] + height),
+ int(xy[0]) : int(xy[0] + width),
+ ] = np.maximum(
+ dynamic_mask[
+ int(xy[1]) : int(xy[1] + height),
+ int(xy[0]) : int(xy[0] + width),
+ ],
+ speed,
+ )
+ # thresholding, use 1.0 m/s to determine whether the pixel is moving
+ dynamic_mask = np.clip((dynamic_mask > 1.0) * 255, 0, 255).astype(np.uint8)
+ dynamic_mask = Image.fromarray(dynamic_mask, "L")
+ dynamic_mask_path = os.path.join(mask_dir, f"{str(frame_idx).zfill(3)}_{str(img.name - 1)}.png")
+ dynamic_mask.save(dynamic_mask_path)
+
+ def save_objects(self, dataset):
+ """Parse and save the ground truth bounding boxes."""
+ instances_info, frame_instances = {}, {}
+
+ for frame_idx, data in enumerate(dataset):
+ frame = dataset_pb2.Frame()
+ frame.ParseFromString(bytearray(data.numpy()))
+
+ frame_instances[frame_idx] = []
+ for l in frame.laser_labels:
+ frame_pose = np.array(frame.pose.transform).reshape(4, 4)
+
+ str_id = str(l.id)
+ if WAYMO_CLASSES[l.type] not in WAYMO_DYNAMIC_CLASSES:
+ continue
+
+ frame_instances[frame_idx].append(str_id)
+
+ if str_id not in instances_info:
+ instances_info[str_id] = dict(
+ id=l.id,
+ # class_ind=l.type,
+ class_name=WAYMO_CLASSES[l.type],
+ frame_annotations={
+ "frame_idx": [],
+ "obj_to_world": [],
+ "box_size": [],
+ }
+ )
+
+ # https://github.com/waymo-research/waymo-open-dataset/blob/master/waymo_open_dataset/label.proto
+ box = l.box
+
+ # Box coordinates in vehicle frame.
+ tx, ty, tz = box.center_x, box.center_y, box.center_z
+
+ # The heading of the bounding box (in radians). The heading is the angle
+ # required to rotate +x to the surface normal of the box front face. It is
+ # normalized to [-pi, pi).
+ c = np.math.cos(box.heading)
+ s = np.math.sin(box.heading)
+
+ # [object to vehicle]
+ # https://github.com/gdlg/simple-waymo-open-dataset-reader/blob/d488196b3ded6574c32fad391467863b948dfd8e/simple_waymo_open_dataset_reader/utils.py#L32
+ o2v = np.array([
+ [ c, -s, 0, tx],
+ [ s, c, 0, ty],
+ [ 0, 0, 1, tz],
+ [ 0, 0, 0, 1]])
+
+ # [object to ENU world]
+ pose = frame_pose @ o2v # o2w = v2w @ o2v
+
+ # difficulty = l.detection_difficulty_level
+
+ # tracking_difficulty = l.tracking_difficulty_level
+
+ # Dimensions of the box. length: dim x. width: dim y. height: dim z.
+ # length: dim_x: along heading; dim_y: verticle to heading; dim_z: verticle up
+ dimension = [box.length, box.width, box.height]
+
+ instances_info[str_id]['frame_annotations']['frame_idx'].append(frame_idx)
+ instances_info[str_id]['frame_annotations']['obj_to_world'].append(pose.tolist())
+ instances_info[str_id]['frame_annotations']['box_size'].append(dimension)
+
+ # Correct ID mapping
+ id_map = {}
+ for i, (k, v) in enumerate(instances_info.items()):
+ id_map[v["id"]] = i
+
+ # Update keys in instances_info
+ new_instances_info = {}
+ for k, v in instances_info.items():
+ new_instances_info[id_map[v["id"]]] = v
+
+ # Update keys in frame_instances
+ new_frame_instances = {}
+ for k, v in frame_instances.items():
+ new_frame_instances[k] = [id_map[i] for i in v]
+
+ return new_instances_info, new_frame_instances
+
+ def create_folder(self):
+ """Create folder for data preprocessing."""
+ if self.process_id_list is None:
+ id_list = range(len(self))
+ else:
+ id_list = self.process_id_list
+ for i in id_list:
+ if "images" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/images", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/sky_masks", exist_ok=True)
+ if "calib" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/extrinsics", exist_ok=True)
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/intrinsics", exist_ok=True)
+ if "pose" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/ego_pose", exist_ok=True)
+ if "lidar" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/lidar", exist_ok=True)
+ if "dynamic_masks" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/dynamic_masks", exist_ok=True)
+ if "objects" in self.process_keys:
+ os.makedirs(f"{self.save_dir}/{str(i).zfill(3)}/instances", exist_ok=True)
diff --git a/worldbench/third_party/drivestudio/datasets/waymo/waymo_sourceloader.py b/worldbench/third_party/drivestudio/datasets/waymo/waymo_sourceloader.py
new file mode 100644
index 0000000..f083577
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/waymo/waymo_sourceloader.py
@@ -0,0 +1,533 @@
+from typing import Dict
+import logging
+import os
+import json
+import joblib
+import numpy as np
+from tqdm import trange, tqdm
+from omegaconf import OmegaConf
+
+import torch
+from torch import Tensor
+
+from pytorch3d.transforms import matrix_to_quaternion
+from datasets.base.scene_dataset import ModelType
+from datasets.base.lidar_source import SceneLidarSource
+from datasets.base.pixel_source import ScenePixelSource, CameraData
+
+logger = logging.getLogger()
+
+# define each class's node type
+OBJECT_CLASS_NODE_MAPPING = {
+ "Vehicle": ModelType.RigidNodes,
+ "Pedestrian": ModelType.SMPLNodes,
+ "Cyclist": ModelType.DeformableNodes
+}
+SMPLNODE_CLASSES = ["Pedestrian"]
+
+# OpenCV to Dataset coordinate transformation
+# opencv coordinate system: x right, y down, z front
+# waymo coordinate system: x front, y left, z up
+OPENCV2DATASET = np.array(
+ [[0, 0, 1, 0], [-1, 0, 0, 0], [0, -1, 0, 0], [0, 0, 0, 1]]
+)
+
+# Waymo Camera List:
+# 0: front_camera
+# 1: front_left_camera
+# 2: front_right_camera
+# 3: left_camera
+# 4: right_camera
+AVAILABLE_CAM_LIST = [0, 1, 2, 3, 4]
+
+class WaymoCameraData(CameraData):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ def load_calibrations(self):
+ """
+ Load the camera intrinsics, extrinsics, timestamps, etc.
+ Compute the camera-to-world matrices, ego-to-world matrices, etc.
+ """
+ # load camera intrinsics
+ # 1d Array of [f_u, f_v, c_u, c_v, k{1, 2}, p{1, 2}, k{3}].
+ # ====!! we did not use distortion parameters for simplicity !!====
+ # to be improved!!
+ intrinsic = np.loadtxt(
+ os.path.join(self.data_path, "intrinsics", f"{self.cam_id}.txt")
+ )
+ fx, fy, cx, cy = intrinsic[0], intrinsic[1], intrinsic[2], intrinsic[3]
+ k1, k2, p1, p2, k3 = intrinsic[4], intrinsic[5], intrinsic[6], intrinsic[7], intrinsic[8]
+ # scale intrinsics w.r.t. load size
+ fx, fy = (
+ fx * self.load_size[1] / self.original_size[1],
+ fy * self.load_size[0] / self.original_size[0],
+ )
+ cx, cy = (
+ cx * self.load_size[1] / self.original_size[1],
+ cy * self.load_size[0] / self.original_size[0],
+ )
+ _intrinsics = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
+ _distortions = np.array([k1, k2, p1, p2, k3])
+
+ # load camera extrinsics
+ cam_to_ego = np.loadtxt(
+ os.path.join(self.data_path, "extrinsics", f"{self.cam_id}.txt")
+ )
+ # because we use opencv coordinate system to generate camera rays,
+ # we need a transformation matrix to covnert rays from opencv coordinate
+ # system to waymo coordinate system.
+ # opencv coordinate system: x right, y down, z front
+ # waymo coordinate system: x front, y left, z up
+ cam_to_ego = cam_to_ego @ OPENCV2DATASET
+
+ # compute per-image poses and intrinsics
+ cam_to_worlds, ego_to_worlds = [], []
+ intrinsics, distortions = [], []
+
+ # we tranform the camera poses w.r.t. the first timestep to make the translation vector of
+ # the first ego pose as the origin of the world coordinate system.
+ ego_to_world_start = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{self.start_timestep:03d}.txt")
+ )
+ for t in range(self.start_timestep, self.end_timestep):
+ ego_to_world_current = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{t:03d}.txt")
+ )
+ # compute ego_to_world transformation
+ ego_to_world = np.linalg.inv(ego_to_world_start) @ ego_to_world_current
+ ego_to_worlds.append(ego_to_world)
+ # transformation:
+ # (opencv_cam -> waymo_cam -> waymo_ego_vehicle) -> current_world
+ cam2world = ego_to_world @ cam_to_ego
+ cam_to_worlds.append(cam2world)
+ intrinsics.append(_intrinsics)
+ distortions.append(_distortions)
+
+ self.intrinsics = torch.from_numpy(np.stack(intrinsics, axis=0)).float()
+ self.distortions = torch.from_numpy(np.stack(distortions, axis=0)).float()
+ self.cam_to_worlds = torch.from_numpy(np.stack(cam_to_worlds, axis=0)).float()
+
+ @classmethod
+ def get_camera2worlds(cls, data_path: str, cam_id: str, start_timestep: int, end_timestep: int) -> torch.Tensor:
+ """
+ Returns camera-to-world matrices for the specified camera and time range.
+
+ Args:
+ data_path (str): Path to the dataset.
+ cam_id (str): Camera ID.
+ start_timestep (int): Start timestep.
+ end_timestep (int): End timestep.
+
+ Returns:
+ torch.Tensor: Camera-to-world matrices of shape (num_frames, 4, 4).
+ """
+ # Load camera extrinsics
+ cam_to_ego = np.loadtxt(os.path.join(data_path, "extrinsics", f"{cam_id}.txt"))
+ cam_to_ego = cam_to_ego @ OPENCV2DATASET
+
+ # Load ego poses and compute camera-to-world matrices
+ cam_to_worlds = []
+ ego_to_world_start = np.loadtxt(os.path.join(data_path, "ego_pose", f"{start_timestep:03d}.txt"))
+
+ for t in range(start_timestep, end_timestep):
+ ego_to_world_current = np.loadtxt(os.path.join(data_path, "ego_pose", f"{t:03d}.txt"))
+ ego_to_world = np.linalg.inv(ego_to_world_start) @ ego_to_world_current
+ cam2world = ego_to_world @ cam_to_ego
+ cam_to_worlds.append(cam2world)
+
+ return torch.from_numpy(np.stack(cam_to_worlds, axis=0)).float()
+
+class WaymoPixelSource(ScenePixelSource):
+ def __init__(
+ self,
+ dataset_name: str,
+ pixel_data_config: OmegaConf,
+ data_path: str,
+ start_timestep: int,
+ end_timestep: int,
+ device: torch.device = torch.device("cpu"),
+ ):
+ super().__init__(dataset_name, pixel_data_config, device=device)
+ self.data_path = data_path
+ self.start_timestep = start_timestep
+ self.end_timestep = end_timestep
+ self.load_data()
+
+ def load_cameras(self):
+ self._timesteps = torch.arange(self.start_timestep, self.end_timestep)
+ self.register_normalized_timestamps()
+
+ for idx, cam_id in enumerate(self.camera_list):
+ logger.info(f"Loading camera {cam_id}")
+ camera = WaymoCameraData(
+ dataset_name=self.dataset_name,
+ data_path=self.data_path,
+ cam_id=cam_id,
+ start_timestep=self.start_timestep,
+ end_timestep=self.end_timestep,
+ load_dynamic_mask=self.data_cfg.load_dynamic_mask,
+ load_sky_mask=self.data_cfg.load_sky_mask,
+ downscale_when_loading=self.data_cfg.downscale_when_loading[idx],
+ undistort=self.data_cfg.undistort,
+ buffer_downscale=self.buffer_downscale,
+ device=self.device,
+ )
+ camera.load_time(self.normalized_time)
+ unique_img_idx = torch.arange(len(camera), device=self.device) * len(self.camera_list) + idx
+ camera.set_unique_ids(
+ unique_cam_idx = idx,
+ unique_img_idx = unique_img_idx
+ )
+ logger.info(f"Camera {camera.cam_name} loaded.")
+ self.camera_data[cam_id] = camera
+
+ def load_objects(self):
+ """
+ get ground truth bounding boxes of the dynamic objects
+
+ instances_info = {
+ "0": # simplified instance id
+ {
+ "id": str,
+ "class_name": str,
+ "frame_annotations": {
+ "frame_idx": List,
+ "obj_to_world": List,
+ "box_size": List,
+ },
+ ...
+ }
+ frame_instances = {
+ "0": # frame idx
+ List[int] # list of simplified instance ids
+ ...
+ }
+ """
+ instances_info_path = os.path.join(self.data_path, "instances", "instances_info.json")
+ frame_instances_path = os.path.join(self.data_path, "instances", "frame_instances.json")
+ with open(instances_info_path, "r") as f:
+ instances_info = json.load(f)
+ with open(frame_instances_path, "r") as f:
+ frame_instances = json.load(f)
+ # get pose of each instance at each frame
+ # shape (num_frames, num_instances, 4, 4)
+ num_instances = len(instances_info)
+ num_full_frames = len(frame_instances)
+ instances_pose = np.zeros((num_full_frames, num_instances, 4, 4))
+ instances_size = np.zeros((num_full_frames, num_instances, 3))
+ instances_true_id = np.arange(num_instances)
+ instances_model_types = np.ones(num_instances) * -1
+
+ ego_to_world_start = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{self.start_timestep:03d}.txt")
+ )
+ for k, v in instances_info.items():
+ instances_model_types[int(k)] = OBJECT_CLASS_NODE_MAPPING[v["class_name"]]
+ for frame_idx, obj_to_world, box_size in zip(v["frame_annotations"]["frame_idx"], v["frame_annotations"]["obj_to_world"], v["frame_annotations"]["box_size"]):
+ # the first ego pose as the origin of the world coordinate system.
+ obj_to_world = np.array(obj_to_world).reshape(4, 4)
+ obj_to_world = np.linalg.inv(ego_to_world_start) @ obj_to_world
+ instances_pose[frame_idx, int(k)] = np.array(obj_to_world)
+ instances_size[frame_idx, int(k)] = np.array(box_size)
+
+ # get frame valid instances
+ # shape (num_frames, num_instances)
+ per_frame_instance_mask = np.zeros((num_full_frames, num_instances))
+ for frame_idx, valid_instances in frame_instances.items():
+ per_frame_instance_mask[int(frame_idx), valid_instances] = 1
+
+ # select the frames that are in the range of start_timestep and end_timestep
+ instances_pose = torch.from_numpy(instances_pose[self.start_timestep:self.end_timestep]).float()
+ instances_size = torch.from_numpy(instances_size[self.start_timestep:self.end_timestep]).float()
+ instances_true_id = torch.from_numpy(instances_true_id).long()
+ instances_model_types = torch.from_numpy(instances_model_types).long()
+ per_frame_instance_mask = torch.from_numpy(per_frame_instance_mask[self.start_timestep:self.end_timestep]).bool()
+
+ # filter out the instances that are not visible in selected frames
+ ins_frame_cnt = per_frame_instance_mask.sum(dim=0)
+ instances_pose = instances_pose[:, ins_frame_cnt > 0]
+ instances_size = instances_size[:, ins_frame_cnt > 0]
+ instances_true_id = instances_true_id[ins_frame_cnt > 0]
+ instances_model_types = instances_model_types[ins_frame_cnt > 0]
+ per_frame_instance_mask = per_frame_instance_mask[:, ins_frame_cnt > 0]
+
+ # assign to the class
+ # (num_frames, num_instances, 4, 4)
+ self.instances_pose = instances_pose
+ # (num_instances, 3)
+ self.instances_size = instances_size.sum(0) / per_frame_instance_mask.sum(0).unsqueeze(-1)
+ # (num_frames, num_instances)
+ self.per_frame_instance_mask = per_frame_instance_mask
+ # (num_instances)
+ self.instances_true_id = instances_true_id
+ # (num_instances)
+ self.instances_model_types = instances_model_types
+
+ if self.data_cfg.load_smpl:
+ # Collect camera-to-world matrices for all available cameras
+ cam_to_worlds = {}
+ for cam_id in AVAILABLE_CAM_LIST:
+ cam_to_worlds[cam_id] = WaymoCameraData.get_camera2worlds(
+ self.data_path,
+ str(cam_id),
+ self.start_timestep,
+ self.end_timestep
+ )
+
+ # load SMPL parameters
+ smpl_dict = joblib.load(os.path.join(self.data_path, "humanpose", "smpl.pkl"))
+ frame_num = self.end_timestep - self.start_timestep
+
+ smpl_human_all = {}
+ for fi in tqdm(range(self.start_timestep, self.end_timestep), desc="Loading SMPL"):
+ for instance_id, ins_smpl in smpl_dict.items():
+ if instance_id not in smpl_human_all:
+ smpl_human_all[instance_id] = {
+ "smpl_quats": torch.zeros((frame_num, 24, 4), dtype=torch.float32),
+ "smpl_trans": torch.zeros((frame_num, 3), dtype=torch.float32),
+ "smpl_betas": torch.zeros((frame_num, 10), dtype=torch.float32),
+ "frame_valid": torch.zeros((frame_num), dtype=torch.bool)
+ }
+ smpl_human_all[instance_id]["smpl_quats"][:, :, 0] = 1.0
+ if ins_smpl["valid_mask"][fi]:
+ betas = ins_smpl["smpl"]["betas"][fi]
+ smpl_human_all[instance_id]["smpl_betas"][fi - self.start_timestep] = betas
+
+ body_pose = ins_smpl["smpl"]["body_pose"][fi]
+ smpl_orient = ins_smpl["smpl"]["global_orient"][fi]
+ cam_depend = ins_smpl["selected_cam_idx"][fi].item()
+
+ c2w = cam_to_worlds[cam_depend][fi - self.start_timestep]
+ world_orient = c2w[:3, :3].to(smpl_orient.device) @ smpl_orient.squeeze()
+ smpl_quats = matrix_to_quaternion(
+ torch.cat([world_orient[None, ...], body_pose], dim=0)
+ )
+
+ ii = instances_info[str(instance_id)]['frame_annotations']["frame_idx"].index(fi)
+ o2w = np.array(
+ instances_info[str(instance_id)]['frame_annotations']["obj_to_world"][ii]
+ )
+ o2w = torch.from_numpy(
+ np.linalg.inv(ego_to_world_start) @ o2w
+ )
+ # box_size = instances_info[str(instance_id)]['frame_annotations']["box_size"][ii]
+
+ smpl_human_all[instance_id]["smpl_quats"][fi - self.start_timestep] = smpl_quats
+ smpl_human_all[instance_id]["smpl_trans"][fi - self.start_timestep] = o2w[:3, 3]
+ smpl_human_all[instance_id]["frame_valid"][fi - self.start_timestep] = True
+
+ self.smpl_human_all = smpl_human_all
+
+class WaymoLiDARSource(SceneLidarSource):
+ def __init__(
+ self,
+ lidar_data_config: OmegaConf,
+ data_path: str,
+ start_timestep: int,
+ end_timestep: int,
+ device: torch.device = torch.device("cpu"),
+ ):
+ super().__init__(lidar_data_config, device=device)
+ self.data_path = data_path
+ self.start_timestep = start_timestep
+ self.end_timestep = end_timestep
+ self.create_all_filelist()
+ self.load_data()
+
+ def create_all_filelist(self):
+ """
+ Create a list of all the files in the dataset.
+ e.g., a list of all the lidar scans in the dataset.
+ """
+ lidar_filepaths = []
+ for t in range(self.start_timestep, self.end_timestep):
+ lidar_filepaths.append(
+ os.path.join(self.data_path, "lidar", f"{t:03d}.bin")
+ )
+ self.lidar_filepaths = np.array(lidar_filepaths)
+
+ def load_calibrations(self):
+ """
+ Load the calibration files of the dataset.
+ e.g., lidar to world transformation matrices.
+ """
+ # Note that in the Waymo Open Dataset, the lidar coordinate system is the same
+ # as the vehicle coordinate system
+ lidar_to_worlds = []
+
+ # we tranform the poses w.r.t. the first timestep to make the origin of the
+ # first ego pose as the origin of the world coordinate system.
+ ego_to_world_start = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{self.start_timestep:03d}.txt")
+ )
+ for t in range(self.start_timestep, self.end_timestep):
+ ego_to_world_current = np.loadtxt(
+ os.path.join(self.data_path, "ego_pose", f"{t:03d}.txt")
+ )
+ # compute ego_to_world transformation
+ lidar_to_world = np.linalg.inv(ego_to_world_start) @ ego_to_world_current
+ lidar_to_worlds.append(lidar_to_world)
+
+ self.lidar_to_worlds = torch.from_numpy(
+ np.stack(lidar_to_worlds, axis=0)
+ ).float()
+
+ def load_lidar(self):
+ """
+ Load the lidar data of the dataset from the filelist.
+ """
+ origins, directions, ranges, laser_ids = [], [], [], []
+ # flow/ground info are used for evaluation only
+ flows, flow_classes, grounds = [], [], []
+ # in waymo, we simplify timestamps as the time indices
+ timesteps = []
+
+ accumulated_num_original_rays = 0
+ accumulated_num_rays = 0
+ for t in trange(
+ 0, len(self.lidar_filepaths), desc="Loading lidar", dynamic_ncols=True
+ ):
+ # each lidar_info contains an Nx14 array
+ # from left to right:
+ # origins: 3d, points: 3d, flows: 3d, flow_class: 1d,
+ # ground_labels: 1d, intensities: 1d, elongations: 1d, laser_ids: 1d
+ lidar_info = np.memmap(
+ self.lidar_filepaths[t],
+ dtype=np.float32,
+ mode="r",
+ ).reshape(-1, 14)
+ original_length = len(lidar_info)
+ accumulated_num_original_rays += original_length
+
+ # select lidar points based on the laser id
+ if self.data_cfg.only_use_top_lidar:
+ # laser_ids: 0: TOP, 1: FRONT, 2: SIDE_LEFT, 3: SIDE_RIGHT, 4: REAR
+ lidar_info = lidar_info[lidar_info[:, 13] == 0]
+
+ lidar_origins = torch.from_numpy(lidar_info[:, :3]).float()
+ lidar_points = torch.from_numpy(lidar_info[:, 3:6]).float()
+ lidar_ids = torch.from_numpy(lidar_info[:, 13]).float()
+ lidar_flows = torch.from_numpy(lidar_info[:, 6:9]).float()
+ lidar_flow_classes = torch.from_numpy(lidar_info[:, 9]).long()
+ ground_labels = torch.from_numpy(lidar_info[:, 10]).long()
+ # we don't collect intensities and elongations for now
+
+ # select lidar points based on a truncated ego-forward-directional range
+ # this is to make sure most of the lidar points are within the range of the camera
+ valid_mask = torch.ones_like(lidar_origins[:, 0]).bool()
+ if self.data_cfg.truncated_max_range is not None:
+ valid_mask = lidar_points[:, 0] < self.data_cfg.truncated_max_range
+ if self.data_cfg.truncated_min_range is not None:
+ valid_mask = valid_mask & (
+ lidar_points[:, 0] > self.data_cfg.truncated_min_range
+ )
+ lidar_origins = lidar_origins[valid_mask]
+ lidar_points = lidar_points[valid_mask]
+ lidar_ids = lidar_ids[valid_mask]
+ lidar_flows = lidar_flows[valid_mask]
+ lidar_flow_classes = lidar_flow_classes[valid_mask]
+ ground_labels = ground_labels[valid_mask]
+ # transform lidar points from lidar coordinate system to world coordinate system
+ lidar_origins = (
+ self.lidar_to_worlds[t][:3, :3] @ lidar_origins.T
+ + self.lidar_to_worlds[t][:3, 3:4]
+ ).T
+ lidar_points = (
+ self.lidar_to_worlds[t][:3, :3] @ lidar_points.T
+ + self.lidar_to_worlds[t][:3, 3:4]
+ ).T
+ # scene flows are in the lidar coordinate system, so we need to rotate them
+ lidar_flows = (self.lidar_to_worlds[t][:3, :3] @ lidar_flows.T).T
+ # compute lidar directions
+ lidar_directions = lidar_points - lidar_origins
+ lidar_ranges = torch.norm(lidar_directions, dim=-1, keepdim=True)
+ lidar_directions = lidar_directions / lidar_ranges
+ # we use time indices as the timestamp for waymo dataset
+ lidar_timestamp = torch.ones_like(lidar_ranges).squeeze(-1) * t
+ accumulated_num_rays += len(lidar_ranges)
+
+ origins.append(lidar_origins)
+ directions.append(lidar_directions)
+ ranges.append(lidar_ranges)
+ laser_ids.append(lidar_ids)
+ flows.append(lidar_flows)
+ flow_classes.append(lidar_flow_classes)
+ grounds.append(ground_labels)
+ # we use time indices as the timestamp for waymo dataset
+ timesteps.append(lidar_timestamp)
+
+ logger.info(
+ f"Number of lidar rays: {accumulated_num_rays} "
+ f"({accumulated_num_rays / accumulated_num_original_rays * 100:.2f}% of "
+ f"{accumulated_num_original_rays} original rays)"
+ )
+ logger.info("Filter condition:")
+ logger.info(f" only_use_top_lidar: {self.data_cfg.only_use_top_lidar}")
+ logger.info(f" truncated_max_range: {self.data_cfg.truncated_max_range}")
+ logger.info(f" truncated_min_range: {self.data_cfg.truncated_min_range}")
+
+ self.origins = torch.cat(origins, dim=0)
+ self.directions = torch.cat(directions, dim=0)
+ self.ranges = torch.cat(ranges, dim=0)
+ self.laser_ids = torch.cat(laser_ids, dim=0)
+ self.visible_masks = torch.zeros_like(self.ranges).squeeze().bool()
+ self.colors = torch.ones_like(self.directions)
+ # becasue the flows here are velocities (m/s), and the fps of the lidar is 10,
+ # we need to divide the velocities by 10 to get the displacements/flows
+ # between two consecutive lidar scans
+ self.flows = torch.cat(flows, dim=0) / 10.0
+ self.flow_classes = torch.cat(flow_classes, dim=0)
+ self.grounds = torch.cat(grounds, dim=0).bool()
+
+ # the underscore here is important.
+ self._timesteps = torch.cat(timesteps, dim=0)
+ self.register_normalized_timestamps()
+
+ def to(self, device: torch.device):
+ super().to(device)
+ self.flows = self.flows.to(device)
+ self.flow_classes = self.flow_classes.to(device)
+ self.grounds = self.grounds.to(self.device)
+
+ def get_lidar_rays(self, time_idx: int) -> Dict[str, Tensor]:
+ """
+ Get the of rays for rendering at the given timestep.
+ Args:
+ time_idx: the index of the lidar scan to render.
+ Returns:
+ a dict of the sampled rays.
+ """
+ origins = self.origins[self.timesteps == time_idx]
+ directions = self.directions[self.timesteps == time_idx]
+ ranges = self.ranges[self.timesteps == time_idx]
+ normalized_time = self.normalized_time[self.timesteps == time_idx]
+ flows = self.flows[self.timesteps == time_idx]
+ return {
+ "lidar_origins": origins,
+ "lidar_viewdirs": directions,
+ "lidar_ranges": ranges,
+ "lidar_normed_time": normalized_time,
+ "lidar_mask": self.timesteps == time_idx,
+ "lidar_flows": flows,
+ }
+
+ def delete_invisible_pts(self) -> None:
+ """
+ Clear the unvisible points.
+ """
+ if self.visible_masks is not None:
+ num_bf = self.origins.shape[0]
+ self.origins = self.origins[self.visible_masks]
+ self.directions = self.directions[self.visible_masks]
+ self.ranges = self.ranges[self.visible_masks]
+ self.flows = self.flows[self.visible_masks]
+ self._timesteps = self._timesteps[self.visible_masks]
+ self._normalized_time = self._normalized_time[self.visible_masks]
+ self.colors = self.colors[self.visible_masks]
+ logger.info(
+ f"[Lidar] {num_bf - self.visible_masks.sum()} out of {num_bf} points are cleared. {self.visible_masks.sum()} points left."
+ )
+ self.visible_masks = None
+ else:
+ logger.info("[Lidar] No unvisible points to clear.")
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/datasets/waymo/waymo_utils.py b/worldbench/third_party/drivestudio/datasets/waymo/waymo_utils.py
new file mode 100644
index 0000000..239f10b
--- /dev/null
+++ b/worldbench/third_party/drivestudio/datasets/waymo/waymo_utils.py
@@ -0,0 +1,443 @@
+# Acknowledgement:
+# 1. https://github.com/open-mmlab/mmdetection3d/blob/main/tools/dataset_converters/waymo_converter.py
+# 2. https://github.com/leolyj/DCA-SRSFE/blob/main/data_preprocess/Waymo/generate_flow.py
+try:
+ from waymo_open_dataset import dataset_pb2
+except ImportError:
+ raise ImportError(
+ 'Please run "pip install waymo-open-dataset-tf-2-6-0" '
+ ">1.4.5 to install the official devkit first."
+ )
+
+import numpy as np
+import tensorflow as tf
+from waymo_open_dataset.utils import range_image_utils, transform_utils
+from waymo_open_dataset.wdl_limited.camera.ops import py_camera_model_ops
+
+def project_vehicle_to_image(vehicle_pose, calibration, points):
+ """Projects from vehicle coordinate system to image with global shutter.
+
+ Arguments:
+ vehicle_pose: Vehicle pose transform from vehicle into world coordinate
+ system.
+ calibration: Camera calibration details (including intrinsics/extrinsics).
+ points: Points to project of shape [N, 3] in vehicle coordinate system.
+
+ Returns:
+ Array of shape [N, 3], with the latter dimension composed of (u, v, ok).
+ """
+ # Transform points from vehicle to world coordinate system (can be
+ # vectorized).
+ pose_matrix = np.array(vehicle_pose.transform).reshape(4, 4)
+ world_points = np.zeros_like(points)
+ for i, point in enumerate(points):
+ cx, cy, cz, _ = np.matmul(pose_matrix, [*point, 1])
+ world_points[i] = (cx, cy, cz)
+
+ # Populate camera image metadata. Velocity and latency stats are filled with
+ # zeroes.
+ extrinsic = tf.reshape(
+ tf.constant(list(calibration.extrinsic.transform), dtype=tf.float32), [4, 4]
+ )
+ intrinsic = tf.constant(list(calibration.intrinsic), dtype=tf.float32)
+ metadata = tf.constant(
+ [
+ calibration.width,
+ calibration.height,
+ dataset_pb2.CameraCalibration.GLOBAL_SHUTTER,
+ ],
+ dtype=tf.int32,
+ )
+ camera_image_metadata = list(vehicle_pose.transform) + [0.0] * 10
+
+ # Perform projection and return projected image coordinates (u, v, ok).
+ return py_camera_model_ops.world_to_image(
+ extrinsic, intrinsic, metadata, camera_image_metadata, world_points
+ ).numpy()
+
+
+def compute_range_image_cartesian(
+ range_image_polar,
+ extrinsic,
+ pixel_pose=None,
+ frame_pose=None,
+ dtype=tf.float32,
+ scope=None,
+):
+ """Computes range image cartesian coordinates from polar ones.
+
+ Args:
+ range_image_polar: [B, H, W, 3] float tensor. Lidar range image in polar
+ coordinate in sensor frame.
+ extrinsic: [B, 4, 4] float tensor. Lidar extrinsic.
+ pixel_pose: [B, H, W, 4, 4] float tensor. If not None, it sets pose for each
+ range image pixel.
+ frame_pose: [B, 4, 4] float tensor. This must be set when pixel_pose is set.
+ It decides the vehicle frame at which the cartesian points are computed.
+ dtype: float type to use internally. This is needed as extrinsic and
+ inclination sometimes have higher resolution than range_image.
+ scope: the name scope.
+
+ Returns:
+ range_image_cartesian: [B, H, W, 3] cartesian coordinates.
+ """
+ range_image_polar_dtype = range_image_polar.dtype
+ range_image_polar = tf.cast(range_image_polar, dtype=dtype)
+ extrinsic = tf.cast(extrinsic, dtype=dtype)
+ if pixel_pose is not None:
+ pixel_pose = tf.cast(pixel_pose, dtype=dtype)
+ if frame_pose is not None:
+ frame_pose = tf.cast(frame_pose, dtype=dtype)
+
+ with tf.compat.v1.name_scope(
+ scope,
+ "ComputeRangeImageCartesian",
+ [range_image_polar, extrinsic, pixel_pose, frame_pose],
+ ):
+ azimuth, inclination, range_image_range = tf.unstack(range_image_polar, axis=-1)
+
+ cos_azimuth = tf.cos(azimuth)
+ sin_azimuth = tf.sin(azimuth)
+ cos_incl = tf.cos(inclination)
+ sin_incl = tf.sin(inclination)
+
+ # [B, H, W].
+ x = cos_azimuth * cos_incl * range_image_range
+ y = sin_azimuth * cos_incl * range_image_range
+ z = sin_incl * range_image_range
+
+ # [B, H, W, 3]
+ range_image_points = tf.stack([x, y, z], -1)
+ range_image_origins = tf.zeros_like(range_image_points)
+ # [B, 3, 3]
+ rotation = extrinsic[..., 0:3, 0:3]
+ # translation [B, 1, 3]
+ translation = tf.expand_dims(tf.expand_dims(extrinsic[..., 0:3, 3], 1), 1)
+
+ # To vehicle frame.
+ # [B, H, W, 3]
+ range_image_points = (
+ tf.einsum("bkr,bijr->bijk", rotation, range_image_points) + translation
+ )
+ range_image_origins = (
+ tf.einsum("bkr,bijr->bijk", rotation, range_image_origins) + translation
+ )
+ if pixel_pose is not None:
+ # To global frame.
+ # [B, H, W, 3, 3]
+ pixel_pose_rotation = pixel_pose[..., 0:3, 0:3]
+ # [B, H, W, 3]
+ pixel_pose_translation = pixel_pose[..., 0:3, 3]
+ # [B, H, W, 3]
+ range_image_points = (
+ tf.einsum("bhwij,bhwj->bhwi", pixel_pose_rotation, range_image_points)
+ + pixel_pose_translation
+ )
+ range_image_origins = (
+ tf.einsum("bhwij,bhwj->bhwi", pixel_pose_rotation, range_image_origins)
+ + pixel_pose_translation
+ )
+
+ if frame_pose is None:
+ raise ValueError("frame_pose must be set when pixel_pose is set.")
+ # To vehicle frame corresponding to the given frame_pose
+ # [B, 4, 4]
+ world_to_vehicle = tf.linalg.inv(frame_pose)
+ world_to_vehicle_rotation = world_to_vehicle[:, 0:3, 0:3]
+ world_to_vehicle_translation = world_to_vehicle[:, 0:3, 3]
+ # [B, H, W, 3]
+ range_image_points = (
+ tf.einsum(
+ "bij,bhwj->bhwi", world_to_vehicle_rotation, range_image_points
+ )
+ + world_to_vehicle_translation[:, tf.newaxis, tf.newaxis, :]
+ )
+ range_image_origins = (
+ tf.einsum(
+ "bij,bhwj->bhwi", world_to_vehicle_rotation, range_image_origins
+ )
+ + world_to_vehicle_translation[:, tf.newaxis, tf.newaxis, :]
+ )
+
+ range_image_points = tf.cast(range_image_points, dtype=range_image_polar_dtype)
+ range_image_origins = tf.cast(
+ range_image_origins, dtype=range_image_polar_dtype
+ )
+ return range_image_points, range_image_origins
+
+
+def extract_point_cloud_from_range_image(
+ range_image,
+ extrinsic,
+ inclination,
+ pixel_pose=None,
+ frame_pose=None,
+ dtype=tf.float32,
+ scope=None,
+):
+ """Extracts point cloud from range image.
+
+ Args:
+ range_image: [B, H, W] tensor. Lidar range images.
+ extrinsic: [B, 4, 4] tensor. Lidar extrinsic.
+ inclination: [B, H] tensor. Inclination for each row of the range image.
+ 0-th entry corresponds to the 0-th row of the range image.
+ pixel_pose: [B, H, W, 4, 4] tensor. If not None, it sets pose for each range
+ image pixel.
+ frame_pose: [B, 4, 4] tensor. This must be set when pixel_pose is set. It
+ decides the vehicle frame at which the cartesian points are computed.
+ dtype: float type to use internally. This is needed as extrinsic and
+ inclination sometimes have higher resolution than range_image.
+ scope: the name scope.
+
+ Returns:
+ range_image_points: [B, H, W, 3] with {x, y, z} as inner dims in vehicle frame.
+ range_image_origins: [B, H, W, 3] with {x, y, z}, the origin of the range image
+ """
+ with tf.compat.v1.name_scope(
+ scope,
+ "ExtractPointCloudFromRangeImage",
+ [range_image, extrinsic, inclination, pixel_pose, frame_pose],
+ ):
+ range_image_polar = range_image_utils.compute_range_image_polar(
+ range_image, extrinsic, inclination, dtype=dtype
+ )
+ (
+ range_image_points_cartesian,
+ range_image_origins_cartesian,
+ ) = compute_range_image_cartesian(
+ range_image_polar,
+ extrinsic,
+ pixel_pose=pixel_pose,
+ frame_pose=frame_pose,
+ dtype=dtype,
+ )
+ return range_image_origins_cartesian, range_image_points_cartesian
+
+
+def parse_range_image_flow_and_camera_projection(frame):
+ range_images = {}
+ camera_projections = {}
+ range_image_top_pose = None
+ for laser in frame.lasers:
+ if (
+ len(laser.ri_return1.range_image_flow_compressed) > 0
+ ): # pylint: disable=g-explicit-length-test
+ range_image_str_tensor = tf.io.decode_compressed(
+ laser.ri_return1.range_image_flow_compressed, "ZLIB"
+ )
+ ri = dataset_pb2.MatrixFloat()
+ ri.ParseFromString(bytearray(range_image_str_tensor.numpy()))
+ range_images[laser.name] = [ri]
+
+ if laser.name == dataset_pb2.LaserName.TOP:
+ range_image_top_pose_str_tensor = tf.io.decode_compressed(
+ laser.ri_return1.range_image_pose_compressed, "ZLIB"
+ )
+ range_image_top_pose = dataset_pb2.MatrixFloat()
+ range_image_top_pose.ParseFromString(
+ bytearray(range_image_top_pose_str_tensor.numpy())
+ )
+
+ camera_projection_str_tensor = tf.io.decode_compressed(
+ laser.ri_return1.camera_projection_compressed, "ZLIB"
+ )
+ cp = dataset_pb2.MatrixInt32()
+ cp.ParseFromString(bytearray(camera_projection_str_tensor.numpy()))
+ camera_projections[laser.name] = [cp]
+ if (
+ len(laser.ri_return2.range_image_flow_compressed) > 0
+ ): # pylint: disable=g-explicit-length-test
+ range_image_str_tensor = tf.io.decode_compressed(
+ laser.ri_return2.range_image_flow_compressed, "ZLIB"
+ )
+ ri = dataset_pb2.MatrixFloat()
+ ri.ParseFromString(bytearray(range_image_str_tensor.numpy()))
+ range_images[laser.name].append(ri)
+
+ camera_projection_str_tensor = tf.io.decode_compressed(
+ laser.ri_return2.camera_projection_compressed, "ZLIB"
+ )
+ cp = dataset_pb2.MatrixInt32()
+ cp.ParseFromString(bytearray(camera_projection_str_tensor.numpy()))
+ camera_projections[laser.name].append(cp)
+ return range_images, camera_projections, range_image_top_pose
+
+
+def convert_range_image_to_point_cloud_flow(
+ frame,
+ range_images,
+ range_images_flow,
+ camera_projections,
+ range_image_top_pose,
+ ri_index=0,
+):
+ """
+ Modified from the codes of Waymo Open Dataset.
+ Convert range images to point cloud.
+ Convert range images flow to scene flow.
+ Args:
+ frame: open dataset frame
+ range_images: A dict of {laser_name, [range_image_first_return, range_image_second_return]}.
+ range_imaages_flow: A dict similar to range_images.
+ camera_projections: A dict of {laser_name,
+ [camera_projection_from_first_return, camera_projection_from_second_return]}.
+ range_image_top_pose: range image pixel pose for top lidar.
+ ri_index: 0 for the first return, 1 for the second return.
+
+ Returns:
+ points: {[N, 3]} list of 3d lidar points of length 5 (number of lidars).
+ points_flow: {[N, 3]} list of scene flow vector of each point.
+ cp_points: {[N, 6]} list of camera projections of length 5 (number of lidars).
+ """
+ calibrations = sorted(frame.context.laser_calibrations, key=lambda c: c.name)
+ origins, points, cp_points = [], [], []
+ points_intensity = []
+ points_elongation = []
+ points_flow = []
+ laser_ids = []
+
+ frame_pose = tf.convert_to_tensor(
+ np.reshape(np.array(frame.pose.transform), [4, 4])
+ )
+ # [H, W, 6]
+ range_image_top_pose_tensor = tf.reshape(
+ tf.convert_to_tensor(range_image_top_pose.data), range_image_top_pose.shape.dims
+ )
+ # [H, W, 3, 3]
+ range_image_top_pose_tensor_rotation = transform_utils.get_rotation_matrix(
+ range_image_top_pose_tensor[..., 0],
+ range_image_top_pose_tensor[..., 1],
+ range_image_top_pose_tensor[..., 2],
+ )
+ range_image_top_pose_tensor_translation = range_image_top_pose_tensor[..., 3:]
+ range_image_top_pose_tensor = transform_utils.get_transform(
+ range_image_top_pose_tensor_rotation, range_image_top_pose_tensor_translation
+ )
+ for c in calibrations:
+ range_image = range_images[c.name][ri_index]
+ range_image_flow = range_images_flow[c.name][ri_index]
+ if len(c.beam_inclinations) == 0: # pylint: disable=g-explicit-length-test
+ beam_inclinations = range_image_utils.compute_inclination(
+ tf.constant([c.beam_inclination_min, c.beam_inclination_max]),
+ height=range_image.shape.dims[0],
+ )
+ else:
+ beam_inclinations = tf.constant(c.beam_inclinations)
+
+ beam_inclinations = tf.reverse(beam_inclinations, axis=[-1])
+ extrinsic = np.reshape(np.array(c.extrinsic.transform), [4, 4])
+
+ range_image_tensor = tf.reshape(
+ tf.convert_to_tensor(range_image.data), range_image.shape.dims
+ )
+ range_image_flow_tensor = tf.reshape(
+ tf.convert_to_tensor(range_image_flow.data), range_image_flow.shape.dims
+ )
+ pixel_pose_local = None
+ frame_pose_local = None
+ if c.name == dataset_pb2.LaserName.TOP:
+ pixel_pose_local = range_image_top_pose_tensor
+ pixel_pose_local = tf.expand_dims(pixel_pose_local, axis=0)
+ frame_pose_local = tf.expand_dims(frame_pose, axis=0)
+ range_image_mask = range_image_tensor[..., 0] > 0
+ range_image_intensity = range_image_tensor[..., 1]
+ range_image_elongation = range_image_tensor[..., 2]
+
+ flow_x = range_image_flow_tensor[..., 0]
+ flow_y = range_image_flow_tensor[..., 1]
+ flow_z = range_image_flow_tensor[..., 2]
+ flow_class = range_image_flow_tensor[..., 3]
+
+ mask_index = tf.where(range_image_mask)
+
+ (origins_cartesian, points_cartesian,) = extract_point_cloud_from_range_image(
+ tf.expand_dims(range_image_tensor[..., 0], axis=0),
+ tf.expand_dims(extrinsic, axis=0),
+ tf.expand_dims(tf.convert_to_tensor(beam_inclinations), axis=0),
+ pixel_pose=pixel_pose_local,
+ frame_pose=frame_pose_local,
+ )
+ origins_cartesian = tf.squeeze(origins_cartesian, axis=0)
+ points_cartesian = tf.squeeze(points_cartesian, axis=0)
+
+ origins_tensor = tf.gather_nd(origins_cartesian, mask_index)
+ points_tensor = tf.gather_nd(points_cartesian, mask_index)
+
+ points_intensity_tensor = tf.gather_nd(range_image_intensity, mask_index)
+ points_elongation_tensor = tf.gather_nd(range_image_elongation, mask_index)
+
+ points_flow_x_tensor = tf.expand_dims(tf.gather_nd(flow_x, mask_index), axis=1)
+ points_flow_y_tensor = tf.expand_dims(tf.gather_nd(flow_y, mask_index), axis=1)
+ points_flow_z_tensor = tf.expand_dims(tf.gather_nd(flow_z, mask_index), axis=1)
+ points_flow_class_tensor = tf.expand_dims(
+ tf.gather_nd(flow_class, mask_index), axis=1
+ )
+
+ origins.append(origins_tensor.numpy())
+ points.append(points_tensor.numpy())
+ points_intensity.append(points_intensity_tensor.numpy())
+ points_elongation.append(points_elongation_tensor.numpy())
+ laser_ids.append(np.full_like(points_intensity_tensor.numpy(), c.name - 1))
+
+ points_flow.append(
+ tf.concat(
+ [
+ points_flow_x_tensor,
+ points_flow_y_tensor,
+ points_flow_z_tensor,
+ points_flow_class_tensor,
+ ],
+ axis=-1,
+ ).numpy()
+ )
+
+ return (
+ origins,
+ points,
+ points_flow,
+ cp_points,
+ points_intensity,
+ points_elongation,
+ laser_ids,
+ )
+
+def get_ground_np(pts):
+ """
+ This function performs ground removal on a point cloud.
+ Modified from https://github.com/tusen-ai/LiDAR_SOT/blob/main/waymo_data/data_preprocessing/ground_removal.py
+
+ Args:
+ pts (numpy.ndarray): The input point cloud.
+
+ Returns:
+ numpy.ndarray: A boolean array indicating whether each point is ground or not.
+ """
+ th_seeds_ = 1.2
+ num_lpr_ = 20
+ n_iter = 10
+ th_dist_ = 0.3
+ pts_sort = pts[pts[:, 2].argsort(), :]
+ lpr = np.mean(pts_sort[:num_lpr_, 2])
+ pts_g = pts_sort[pts_sort[:, 2] < lpr + th_seeds_, :]
+ normal_ = np.zeros(3)
+ for i in range(n_iter):
+ mean = np.mean(pts_g, axis=0)[:3]
+ xx = np.mean((pts_g[:, 0] - mean[0]) * (pts_g[:, 0] - mean[0]))
+ xy = np.mean((pts_g[:, 0] - mean[0]) * (pts_g[:, 1] - mean[1]))
+ xz = np.mean((pts_g[:, 0] - mean[0]) * (pts_g[:, 2] - mean[2]))
+ yy = np.mean((pts_g[:, 1] - mean[1]) * (pts_g[:, 1] - mean[1]))
+ yz = np.mean((pts_g[:, 1] - mean[1]) * (pts_g[:, 2] - mean[2]))
+ zz = np.mean((pts_g[:, 2] - mean[2]) * (pts_g[:, 2] - mean[2]))
+ cov = np.array(
+ [[xx, xy, xz], [xy, yy, yz], [xz, yz, zz]],
+ dtype=np.float32,
+ )
+ U, S, V = np.linalg.svd(cov)
+ normal_ = U[:, 2]
+ d_ = -normal_.dot(mean)
+ th_dist_d_ = th_dist_ - d_
+ result = pts[:, :3] @ normal_[..., np.newaxis]
+ pts_g = pts[result.squeeze(-1) < th_dist_d_]
+ ground_label = result < th_dist_d_
+ return ground_label
diff --git a/worldbench/third_party/drivestudio/models/__init__.py b/worldbench/third_party/drivestudio/models/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/worldbench/third_party/drivestudio/models/gaussians/__init__.py b/worldbench/third_party/drivestudio/models/gaussians/__init__.py
new file mode 100644
index 0000000..c9a972b
--- /dev/null
+++ b/worldbench/third_party/drivestudio/models/gaussians/__init__.py
@@ -0,0 +1,4 @@
+from .vanilla import VanillaGaussians
+from .deformgs import DeformableGaussians
+from .pvg import PeriodicVibrationGaussians
+from .scaffold import ScaffoldGaussians
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/models/gaussians/basics.py b/worldbench/third_party/drivestudio/models/gaussians/basics.py
new file mode 100644
index 0000000..369289d
--- /dev/null
+++ b/worldbench/third_party/drivestudio/models/gaussians/basics.py
@@ -0,0 +1,246 @@
+import math
+import numpy as np
+
+import torch
+from numpy.typing import NDArray
+
+from typing import Dict, List, Optional, Union
+from dataclasses import dataclass, fields
+from sklearn.neighbors import NearestNeighbors
+from pytorch3d.transforms import matrix_to_quaternion
+
+from gsplat.rendering import rasterization
+from gsplat.cuda_legacy._wrapper import num_sh_bases
+from gsplat.cuda_legacy._torch_impl import quat_to_rotmat
+from gsplat.cuda._wrapper import spherical_harmonics
+
+def interpolate_quats(q1, q2, fraction=0.5):
+ q1 = q1 / torch.norm(q1, dim=-1, keepdim=True)
+ q2 = q2 / torch.norm(q2, dim=-1, keepdim=True)
+
+ dot = (q1 * q2).sum(dim=-1)
+ dot = torch.clamp(dot, -1, 1)
+
+ neg_mask = dot < 0
+ q2[neg_mask] = -q2[neg_mask]
+ dot[neg_mask] = -dot[neg_mask]
+
+ similar_mask = dot > 0.9995
+ q_interp_similar = q1 + fraction * (q2 - q1)
+
+ theta_0 = torch.acos(dot)
+ theta = theta_0 * fraction
+
+ sin_theta = torch.sin(theta)
+ sin_theta_0 = torch.sin(theta_0)
+
+ s1 = torch.cos(theta) - dot * sin_theta / sin_theta_0
+ s2 = sin_theta / sin_theta_0
+
+ q_interp = (s1[..., None] * q1) + (s2[..., None] * q2)
+
+ final_q_interp = torch.zeros_like(q1)
+ final_q_interp[similar_mask] = q_interp_similar[similar_mask]
+ final_q_interp[~similar_mask] = q_interp[~similar_mask]
+ return final_q_interp
+
+def random_quat_tensor(N):
+ """
+ Defines a random quaternion tensor of shape (N, 4)
+ """
+ u = torch.rand(N)
+ v = torch.rand(N)
+ w = torch.rand(N)
+ return torch.stack(
+ [
+ torch.sqrt(1 - u) * torch.sin(2 * math.pi * v),
+ torch.sqrt(1 - u) * torch.cos(2 * math.pi * v),
+ torch.sqrt(u) * torch.sin(2 * math.pi * w),
+ torch.sqrt(u) * torch.cos(2 * math.pi * w),
+ ],
+ dim=-1,
+ )
+
+def quat_mult(q1, q2):
+ # NOTE:
+ # Q1 is the quaternion that rotates the vector from the original position to the final position
+ # Q2 is the quaternion that been rotated
+ w1, x1, y1, z1 = q1.T
+ w2, x2, y2, z2 = q2.T
+ w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
+ x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
+ y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
+ z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
+ return torch.stack([w, x, y, z]).T
+
+def RGB2SH(rgb):
+ """
+ Converts from RGB values [0,1] to the 0th spherical harmonic coefficient
+ """
+ C0 = 0.28209479177387814
+ return (rgb - 0.5) / C0
+
+
+def SH2RGB(sh):
+ """
+ Converts from the 0th spherical harmonic coefficient to RGB values [0,1]
+ """
+ C0 = 0.28209479177387814
+ return sh * C0 + 0.5
+
+
+def projection_matrix(znear, zfar, fovx, fovy, device:Union[str,torch.device]="cpu"):
+ """
+ Constructs an OpenGL-style perspective projection matrix.
+ """
+ t = znear * math.tan(0.5 * fovy)
+ b = -t
+ r = znear * math.tan(0.5 * fovx)
+ l = -r
+ n = znear
+ f = zfar
+ return torch.tensor(
+ [
+ [2 * n / (r - l), 0.0, (r + l) / (r - l), 0.0],
+ [0.0, 2 * n / (t - b), (t + b) / (t - b), 0.0],
+ [0.0, 0.0, (f + n) / (f - n), -1.0 * f * n / (f - n)],
+ [0.0, 0.0, 1.0, 0.0],
+ ],
+ device=device,
+ )
+
+@dataclass
+class dataclass_camera:
+ camtoworlds: torch.Tensor
+ camtoworlds_gt: torch.Tensor
+ Ks: torch.Tensor
+ H: int
+ W: int
+
+@dataclass
+class dataclass_gs:
+ _opacities: torch.Tensor
+ _means: torch.Tensor
+ _rgbs: torch.Tensor
+ _scales: torch.Tensor
+ _quats: torch.Tensor
+ detach_keys: List[str]
+ extras: Optional[Dict[str, torch.Tensor]] = None
+ def set_grad_controller(self, detach_keys):
+ self.detach_keys = detach_keys
+ @property
+ def opacities(self):
+ if "activated_opacities" in self.detach_keys:
+ return self._opacities.detach()
+ else:
+ return self._opacities
+ @property
+ def means(self):
+ if "means" in self.detach_keys:
+ return self._means.detach()
+ else:
+ return self._means
+ @property
+ def rgbs(self):
+ if "colors" in self.detach_keys:
+ return self._rgbs.detach()
+ else:
+ return self._rgbs
+ @property
+ def scales(self):
+ if "scales" in self.detach_keys:
+ return self._scales.detach()
+ else:
+ return self._scales
+ @property
+ def quats(self):
+ if "quats" in self.detach_keys:
+ return self._quats.detach()
+ else:
+ return self._quats
+
+def remove_from_optim(optimizer, deleted_mask, param_dict):
+ """removes the deleted_mask from the optimizer provided"""
+ for group_idx, group in enumerate(optimizer.param_groups):
+ name = group["name"]
+ if name in param_dict.keys():
+ old_params = group["params"][0]
+ new_params = param_dict[name]
+ assert len(new_params) == 1
+ param_state = optimizer.state[old_params]
+ del optimizer.state[old_params]
+
+ # Modify the state directly without deleting and reassigning.
+ param_state["exp_avg"] = param_state["exp_avg"][~deleted_mask]
+ param_state["exp_avg_sq"] = param_state["exp_avg_sq"][~deleted_mask]
+
+ # Update the parameter in the optimizer's param group.
+ del optimizer.param_groups[group_idx]["params"][0]
+ del optimizer.param_groups[group_idx]["params"]
+ optimizer.param_groups[group_idx]["params"] = new_params
+ optimizer.state[new_params[0]] = param_state
+
+def dup_in_optim(optimizer, dup_mask, param_dict, n=2):
+ """adds the parameters to the optimizer"""
+ for group_idx, group in enumerate(optimizer.param_groups):
+ name = group["name"]
+ if name in param_dict.keys():
+ old_params = group["params"][0]
+ new_params = param_dict[name]
+ param_state = optimizer.state[old_params]
+ repeat_dims = (n,) + tuple(1 for _ in range(param_state["exp_avg"].dim() - 1))
+ param_state["exp_avg"] = torch.cat(
+ [param_state["exp_avg"], torch.zeros_like(param_state["exp_avg"][dup_mask.squeeze()]).repeat(*repeat_dims)],
+ dim=0,
+ )
+ param_state["exp_avg_sq"] = torch.cat(
+ [
+ param_state["exp_avg_sq"],
+ torch.zeros_like(param_state["exp_avg_sq"][dup_mask.squeeze()]).repeat(*repeat_dims),
+ ],
+ dim=0,
+ )
+ del optimizer.state[old_params]
+ optimizer.state[new_params[0]] = param_state
+ optimizer.param_groups[group_idx]["params"] = new_params
+ del old_params
+
+def k_nearest_sklearn(x: torch.Tensor, k: int):
+ """
+ Find the k nearest neighbors (excluding the query point).
+
+ Returns arrays of shape (num_samples, effective_k) where
+ effective_k = min(k, num_samples - 1). Downstream code should
+ handle the case effective_k < k.
+ """
+ num_samples = x.shape[0]
+ effective_k = min(k, max(num_samples - 1, 0))
+ if effective_k == 0:
+ return np.zeros((num_samples, 0), dtype=np.float32), np.empty((num_samples, 0), dtype=np.float32)
+
+ x_np = x.cpu().numpy()
+ nn_model = NearestNeighbors(
+ n_neighbors=effective_k + 1,
+ algorithm="auto",
+ metric="euclidean",
+ ).fit(x_np)
+
+ distances, indices = nn_model.kneighbors(x_np)
+ distances = distances[:, 1:].astype(np.float32)
+ indices = indices[:, 1:].astype(np.float32)
+ return distances, indices
+
+if __name__ == "__main__":
+ quats_prev_frame = torch.tensor([
+ [ 4.3390e-02, 4.1600e-06, -9.5784e-05, 9.9906e-01],
+ [ 1.1272e-04, 1.0807e-08, 9.5874e-05, -1.0000e+00],
+ [ 1.7490e-04, 1.6769e-08, -9.5874e-05, 1.0000e+00]
+ ], device='cuda:0')
+
+ quats_next_frame = torch.tensor([
+ [ 4.2516e-02, 4.0762e-06, -9.5787e-05, 9.9910e-01],
+ [ 3.8867e-05, 3.7264e-09, -9.5874e-05, 1.0000e+00],
+ [ 1.8267e-04, 1.7513e-08, -9.5874e-05, 1.0000e+00]
+ ], device='cuda:0')
+
+ quats_cur_frame = interpolate_quats(quats_prev_frame, quats_next_frame, 0.5)
diff --git a/worldbench/third_party/drivestudio/models/gaussians/deformgs.py b/worldbench/third_party/drivestudio/models/gaussians/deformgs.py
new file mode 100644
index 0000000..c87ce1e
--- /dev/null
+++ b/worldbench/third_party/drivestudio/models/gaussians/deformgs.py
@@ -0,0 +1,167 @@
+"""
+Filename: deformgs.py
+
+Author: Ziyu Chen (ziyu.sjtu@gmail.com)
+
+Description:
+Unofficial implementation of Defromable-GS based on the work by Ziyi Yang et al.
+Original work by Ziyi Yang et al.
+
+Original paper: https://arxiv.org/abs/2309.13101
+"""
+
+from typing import Union, Tuple, Dict, List
+import logging
+import torch
+import numpy as np
+from torch.nn import Parameter
+
+from models.gaussians.basics import *
+from models.gaussians.vanilla import VanillaGaussians
+from models.modules import DeformNetwork
+
+logger = logging.getLogger()
+
+def contract(
+ x: torch.Tensor,
+ aabb: torch.Tensor,
+ ord: Union[str, int] = None,
+) -> torch.Tensor:
+ """
+ Contract the input tensor to the unit cube using piecewise projective function.
+ """
+ # similar to the one in MeRF paper
+ aabb_min, aabb_max = aabb[0, :], aabb[1, :]
+ x = (x - aabb_min) / (aabb_max - aabb_min) # 0~1
+ x = x * 2 - 1 # aabb is at [-1, 1]
+ mag = torch.linalg.norm(x, ord=ord, dim=-1, keepdim=True)
+ x = torch.where(mag < 1, x, (2 - 1 / mag) * (x / mag))
+ x = x / 4 + 0.5 # [-inf, inf] is at [0, 1]
+ return x
+
+def get_linear_noise_func(
+ lr_init, lr_final, lr_delay_steps=0, lr_delay_mult=1.0, max_steps=1000000
+):
+ def helper(step):
+ if step < 0 or (lr_init == 0.0 and lr_final == 0.0):
+ # Disable this parameter
+ return 0.0
+ if lr_delay_steps > 0:
+ # A kind of reverse cosine decay.
+ delay_rate = lr_delay_mult + (1 - lr_delay_mult) * np.sin(
+ 0.5 * np.pi * np.clip(step / lr_delay_steps, 0, 1)
+ )
+ else:
+ delay_rate = 1.0
+ t = np.clip(step / max_steps, 0, 1)
+ log_lerp = lr_init * (1 - t) + lr_final * t
+ return delay_rate * log_lerp
+
+ return helper
+
+class DeformableGaussians(VanillaGaussians):
+ def __init__(
+ self,
+ **kwargs
+ ):
+ super().__init__(**kwargs)
+ self.deform_network = DeformNetwork(
+ input_ch=3, **self.networks_cfg
+ ).to(self.device)
+ self.smooth_term = get_linear_noise_func(lr_init=0.1, lr_final=1e-15, lr_delay_mult=0.01, max_steps=20000)
+ self.delta_xyz_rescale = self.ctrl_cfg.get("delta_xyz_rescale", True)
+
+ @property
+ def defrom_gs(self):
+ if self.step < self.ctrl_cfg.coarse_train_interval:
+ return False
+ else:
+ return True
+
+ def set_bbox(self, bbox: torch.Tensor):
+ self.bbox = bbox.reshape(2, 3)
+
+ def set_cur_frame(self, frame_id: int):
+ self.cur_frame = frame_id
+ def register_normalized_timestamps(self, normalized_timestamps: int):
+ self.normalized_timestamps = normalized_timestamps
+ self.time_interval = 1 / len(normalized_timestamps)
+
+ def create_from_pcd(self, init_means: torch.Tensor, init_colors: torch.Tensor) -> None:
+ super().create_from_pcd(init_means, init_colors)
+
+ def get_param_groups(self) -> Dict[str, List[Parameter]]:
+ param_groups = self.get_gaussian_param_groups()
+ param_groups[self.class_prefix+"deform_network"] = list(self.deform_network.parameters())
+ return param_groups
+
+ def get_deformation(self, canonical_means) -> Tuple:
+ """
+ get the deformation of the nonrigid instances
+ """
+ t = self.normalized_timestamps[self.cur_frame]
+ t = t.unsqueeze(0).repeat(self.num_points, 1)
+ normed_canonical_means = contract(canonical_means, self.bbox) # range from 0 to 1
+
+ ast_noise = torch.randn(1, 1, device=self.device).expand(self.num_points, -1) * self.time_interval * self.smooth_term(self.step)
+
+ delta_xyz, delta_quat, delta_scale = self.deform_network(normed_canonical_means.data, t + ast_noise)
+ return delta_xyz, delta_quat, delta_scale
+
+ def get_gaussians(self, cam: dataclass_camera) -> Dict:
+ filter_mask = torch.ones_like(self._means[:, 0], dtype=torch.bool)
+ self.filter_mask = filter_mask
+
+ delta_xyz, delta_quat, delta_scale = None, None, None
+ if self.defrom_gs:
+ delta_xyz, delta_quat, delta_scale = self.get_deformation(self._means)
+ if self.delta_xyz_rescale:
+ delta_xyz = delta_xyz * self.scene_scale
+
+ if delta_xyz is not None:
+ world_means = self._means + delta_xyz
+ else:
+ world_means = self._means
+
+ if delta_quat is not None:
+ world_quats = self.get_quats + delta_quat
+ else:
+ world_quats = self.get_quats
+
+ if delta_scale is not None:
+ activated_scales = torch.exp(self._scales + delta_scale)
+ else:
+ activated_scales = torch.exp(self._scales)
+
+ # get colors of gaussians
+ colors = torch.cat((self._features_dc[:, None, :], self._features_rest), dim=1)
+ if self.sh_degree > 0:
+ viewdirs = world_means.detach() - cam.camtoworlds.data[..., :3, 3] # (N, 3)
+ viewdirs = viewdirs / viewdirs.norm(dim=-1, keepdim=True)
+ n = min(self.step // self.ctrl_cfg.sh_degree_interval, self.sh_degree)
+ rgbs = spherical_harmonics(n, viewdirs, colors)
+ rgbs = torch.clamp(rgbs + 0.5, 0.0, 1.0)
+ else:
+ rgbs = torch.sigmoid(colors[:, 0, :])
+
+ activated_opacities = self.get_opacity
+ activated_rotations = self.quat_act(world_quats)
+ actovated_colors = rgbs
+
+ # collect gaussians information
+ gs_dict = dict(
+ _means=world_means[filter_mask],
+ _opacities=activated_opacities[filter_mask],
+ _rgbs=actovated_colors[filter_mask],
+ _scales=activated_scales[filter_mask],
+ _quats=activated_rotations[filter_mask],
+ )
+
+ # check nan and inf in gs_dict
+ for k, v in gs_dict.items():
+ if torch.isnan(v).any():
+ raise ValueError(f"NaN detected in gaussian {k} at step {self.step}")
+ if torch.isinf(v).any():
+ raise ValueError(f"Inf detected in gaussian {k} at step {self.step}")
+
+ return gs_dict
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/models/gaussians/pvg.py b/worldbench/third_party/drivestudio/models/gaussians/pvg.py
new file mode 100644
index 0000000..ffe74f2
--- /dev/null
+++ b/worldbench/third_party/drivestudio/models/gaussians/pvg.py
@@ -0,0 +1,450 @@
+"""
+Filename: pvg.py
+
+Author: Ziyu Chen (ziyu.sjtu@gmail.com)
+
+Description:
+Unofficial implementation of PVG based on the work by Yurui Chen, Chun Gu, Junzhe Jiang, Xiatian Zhu, Li Zhang.
+
+Original paper: https://arxiv.org/abs/2311.18561
+"""
+
+from typing import Dict, List
+import random
+import logging
+import torch
+import torch.distributions.uniform as uniform
+from torch.nn import Parameter
+
+from models.gaussians.basics import *
+from models.gaussians.vanilla import VanillaGaussians
+
+logger = logging.getLogger()
+
+class PeriodicVibrationGaussians(VanillaGaussians):
+ def __init__(
+ self,
+ **kwargs
+ ):
+ super().__init__(**kwargs)
+ self._taus = torch.zeros(1, 1, device=self.device)
+ self._betas = torch.zeros(1, 1, device=self.device)
+ self._velocity = torch.zeros(1, 3, device=self.device)
+
+ self.T = self.ctrl_cfg.cycle_length
+ self.t_grad_accum = None
+
+ def set_cur_frame(self, frame_id: int):
+ self.cur_frame = frame_id
+ def register_normalized_timestamps(self, normalized_timestamps: int):
+ """
+ num_timesteps: the total number of timesteps of both train and test
+ """
+ self.normalized_timestamps = normalized_timestamps
+
+ self.num_timestamps = len(normalized_timestamps)
+ self.normalized_time_interval = 1.0 / (self.num_timestamps - 1)
+ self.train_time_scale = self.ctrl_cfg.time_interval / self.normalized_time_interval
+
+ def create_from_pcd(self, init_means: torch.Tensor, init_colors: torch.Tensor, init_times: torch.Tensor) -> None:
+ super().create_from_pcd(init_means, init_colors)
+
+ # time related parameters
+ self._taus = Parameter((init_times * self.train_time_scale).to(self.device)) # life peak
+ self._velocity = Parameter(torch.zeros(self.num_points, 3).to(self.device)) # vibration direction
+ betas_init = torch.sqrt(torch.ones(self.num_points, 1).to(self.device) * self.ctrl_cfg.betas_init)
+ self._betas = Parameter(torch.log(betas_init)) # life span
+
+ @property
+ def get_scaling_t(self):
+ return torch.exp(self._betas)
+ @property
+ def vibr_dirs_norm(self):
+ return self._velocity.norm(dim=-1)
+ @property
+ def temporal_means(self):
+ a = 1/self.T * torch.pi * 2
+ means = self._means + self._velocity * torch.sin(
+ (self.cur_time - self._taus) * a
+ ) / a
+ if self.in_smooth:
+ return means + self.velocity * self.delta_t
+ else:
+ return means
+ @property
+ def temporal_opacities(self):
+ return self.get_opacity * torch.exp(
+ -0.5 * (self.cur_time - self._taus)**2 / (self.get_scaling_t ** 2)
+ )
+ @property
+ def get_marginal_t(self):
+ return torch.exp(-0.5 * (self._taus - self.cur_time) ** 2 / self.get_scaling_t ** 2)
+ @property
+ def rho(self):
+ """staticness coefficient"""
+ return self.get_scaling_t / self.T
+ @property
+ def velocity(self):
+ return self._velocity * torch.exp(-0.5 * self.rho)
+ @property
+ def gamma(self):
+ """
+ dynamic scale factor for Position-aware point adaptive control
+ refer to PVG Section 3.3
+ """
+ with torch.no_grad():
+ gamma = (self._means - self.scene_origin).norm(dim=-1) * self.scene_scale - 1
+ gamma = torch.where(gamma<=1, 1, gamma) / self.scene_scale
+ return gamma
+
+ def after_train(
+ self,
+ radii: torch.Tensor,
+ xys_grad: torch.Tensor,
+ last_size: int,
+ ) -> None:
+ with torch.no_grad():
+ # keep track of a moving average of grad norms
+ visible_mask = (radii > 0).flatten()
+ full_mask = torch.zeros(self.num_points, device=radii.device, dtype=torch.bool)
+ full_mask[self.filter_mask] = visible_mask
+
+ grads = xys_grad.norm(dim=-1)
+ t_grads = self._taus.grad.clone().abs()[self.filter_mask].squeeze()
+ if self.xys_grad_norm is None:
+ self.xys_grad_norm = torch.zeros(self.num_points, device=grads.device, dtype=grads.dtype)
+ self.xys_grad_norm[self.filter_mask] = grads
+
+ self.t_grad_accum = torch.zeros(self.num_points, device=grads.device, dtype=grads.dtype)
+ self.t_grad_accum[self.filter_mask] = t_grads
+ self.vis_counts = torch.ones_like(self.xys_grad_norm)
+ else:
+ assert self.vis_counts is not None
+ self.vis_counts[full_mask] = self.vis_counts[full_mask] + 1
+ self.xys_grad_norm[full_mask] = grads[visible_mask] + self.xys_grad_norm[full_mask]
+ self.t_grad_accum[full_mask] = t_grads[visible_mask] + self.t_grad_accum[full_mask]
+
+ # update the max screen size, as a ratio of number of pixels
+ if self.max_2Dsize is None:
+ self.max_2Dsize = torch.zeros(self.num_points, device=radii.device, dtype=torch.float32)
+ newradii = radii[visible_mask]
+ self.max_2Dsize[full_mask] = torch.maximum(
+ self.max_2Dsize[full_mask], newradii / float(last_size)
+ )
+
+ def get_gaussian_param_groups(self) -> Dict[str, List[Parameter]]:
+ return {
+ self.class_prefix+"xyz": [self._means],
+ self.class_prefix+"sh_dc": [self._features_dc],
+ self.class_prefix+"sh_rest": [self._features_rest],
+ self.class_prefix+"opacity": [self._opacities],
+ self.class_prefix+"scaling": [self._scales],
+ self.class_prefix+"rotation": [self._quats],
+ self.class_prefix+"velocity": [self._velocity],
+ self.class_prefix+"life_peak": [self._taus],
+ self.class_prefix+"life_span": [self._betas]
+ }
+
+ def refinement_after(self, step, optimizer: torch.optim.Optimizer) -> None:
+ assert step == self.step
+ if self.step <= self.ctrl_cfg.warmup_steps:
+ return
+ with torch.no_grad():
+ # only split/cull if we've seen every image since opacity reset
+ reset_interval = self.ctrl_cfg.reset_alpha_interval
+ do_densification = (
+ self.step < self.ctrl_cfg.stop_split_at
+ and self.step % reset_interval > max(self.num_train_images, self.ctrl_cfg.refine_interval)
+ and self.num_points < self.ctrl_cfg.densify_until_num_points
+ )
+ # split & duplicate
+ print(f"Class {self.class_prefix} current points: {self.num_points} @ step {self.step}")
+ if do_densification:
+ assert self.xys_grad_norm is not None and self.vis_counts is not None and self.max_2Dsize is not None
+
+ avg_grad_norm = self.xys_grad_norm / self.vis_counts
+ high_xyz_grads = (avg_grad_norm > self.ctrl_cfg.densify_grad_thresh).squeeze()
+
+ t_avg_grad = self.t_grad_accum / self.vis_counts
+ high_t_grads = t_avg_grad > self.ctrl_cfg.densify_t_grad_thresh
+ high_grads = high_xyz_grads | high_t_grads
+
+ splits_xyz = (
+ self.get_scaling.max(dim=-1).values > \
+ self.ctrl_cfg.densify_size_thresh * self.scene_scale * self.gamma
+ ).squeeze()
+ splits_t = (torch.max(self.get_scaling_t, dim=1).values > self.ctrl_cfg.densify_t_size_thresh) & high_t_grads
+ splits = splits_xyz | splits_t
+
+ if self.step < self.ctrl_cfg.stop_screen_size_at:
+ splits |= (self.max_2Dsize > self.ctrl_cfg.split_screen_size).squeeze()
+ splits &= high_grads
+ nsamps = self.ctrl_cfg.n_split_samples
+ (
+ split_means,
+ split_feature_dc,
+ split_feature_rest,
+ split_opacities,
+ split_scales,
+ split_quats,
+ split_dirs,
+ split_taus,
+ split_betas,
+ ) = self.split_gaussians(splits, nsamps)
+
+ dups_xyz = (
+ self.get_scaling.max(dim=-1).values <= \
+ self.ctrl_cfg.densify_size_thresh * self.scene_scale * self.gamma
+ ).squeeze()
+ dups_t = (torch.max(self.get_scaling_t, dim=1).values <= self.ctrl_cfg.densify_t_size_thresh) & high_t_grads
+
+ dups = dups_xyz | dups_t
+ dups &= high_grads
+ (
+ dup_means,
+ dup_feature_dc,
+ dup_feature_rest,
+ dup_opacities,
+ dup_scales,
+ dup_quats,
+ dup_dirs,
+ dup_taus,
+ dup_betas,
+ ) = self.dup_gaussians(dups)
+
+ self._means = Parameter(torch.cat([self._means.detach(), split_means, dup_means], dim=0))
+ # self.colors_all = Parameter(torch.cat([self.colors_all.detach(), split_colors, dup_colors], dim=0))
+ self._features_dc = Parameter(torch.cat([self._features_dc.detach(), split_feature_dc, dup_feature_dc], dim=0))
+ self._features_rest = Parameter(torch.cat([self._features_rest.detach(), split_feature_rest, dup_feature_rest], dim=0))
+ self._opacities = Parameter(torch.cat([self._opacities.detach(), split_opacities, dup_opacities], dim=0))
+ self._scales = Parameter(torch.cat([self._scales.detach(), split_scales, dup_scales], dim=0))
+ self._quats = Parameter(torch.cat([self._quats.detach(), split_quats, dup_quats], dim=0))
+ self._velocity = Parameter(torch.cat([self._velocity.detach(), split_dirs, dup_dirs], dim=0))
+ self._taus = Parameter(torch.cat([self._taus.detach(), split_taus, dup_taus], dim=0))
+ self._betas = Parameter(torch.cat([self._betas.detach(), split_betas, dup_betas], dim=0))
+
+ # append zeros to the max_2Dsize tensor
+ self.max_2Dsize = torch.cat(
+ [self.max_2Dsize, torch.zeros_like(split_scales[:, 0]), torch.zeros_like(dup_scales[:, 0])],
+ dim=0,
+ )
+
+ split_idcs = torch.where(splits)[0]
+ param_groups = self.get_gaussian_param_groups()
+ dup_in_optim(optimizer, split_idcs, param_groups, n=nsamps)
+
+ dup_idcs = torch.where(dups)[0]
+ param_groups = self.get_gaussian_param_groups()
+ dup_in_optim(optimizer, dup_idcs, param_groups, 1)
+
+ # cull NOTE: Offset all the opacity reset logic by refine_every so that we don't
+ # save checkpoints right when the opacity is reset (saves every 2k)
+ if self.step % reset_interval > max(self.num_train_images, self.ctrl_cfg.refine_interval):
+ deleted_mask = self.cull_gaussians()
+ param_groups = self.get_gaussian_param_groups()
+ remove_from_optim(optimizer, deleted_mask, param_groups)
+ print(f"Class {self.class_prefix} left points: {self.num_points}")
+
+ # reset opacity
+ if self.step % reset_interval == self.ctrl_cfg.refine_interval:
+ # NOTE: in nerfstudio, reset_value = cull_alpha_thresh * 0.8
+ # we align to original repo of gaussians spalting
+ reset_value = torch.min(self.get_opacity.data,
+ torch.ones_like(self._opacities.data) * self.ctrl_cfg.reset_alpha_value)
+ self._opacities.data = torch.logit(reset_value)
+ # reset the exp of optimizer
+ for group in optimizer.param_groups:
+ if group["name"] == self.class_prefix+"opacity":
+ old_params = group["params"][0]
+ param_state = optimizer.state[old_params]
+ param_state["exp_avg"] = torch.zeros_like(param_state["exp_avg"])
+ param_state["exp_avg_sq"] = torch.zeros_like(param_state["exp_avg_sq"])
+ self.xys_grad_norm = None
+ self.t_grad_accum = None
+ self.vis_counts = None
+ self.max_2Dsize = None
+
+ def cull_gaussians(self):
+ """
+ This function deletes gaussians with under a certain opacity threshold
+ """
+ n_bef = self.num_points
+ # cull transparent ones
+ culls = (self.get_opacity.data < self.ctrl_cfg.cull_alpha_thresh).squeeze()
+ if self.step > self.ctrl_cfg.reset_alpha_interval:
+ # cull huge ones
+ toobigs = (
+ torch.exp(self._scales).max(dim=-1).values >
+ self.ctrl_cfg.cull_scale_thresh * self.scene_scale * self.gamma
+ ).squeeze()
+ culls = culls | toobigs
+ if self.step < self.ctrl_cfg.stop_screen_size_at:
+ # cull big screen space
+ assert self.max_2Dsize is not None
+ culls = culls | (self.max_2Dsize > self.ctrl_cfg.cull_screen_size).squeeze()
+ self._means = Parameter(self._means[~culls].detach())
+ self._scales = Parameter(self._scales[~culls].detach())
+ self._quats = Parameter(self._quats[~culls].detach())
+ self._features_dc = Parameter(self._features_dc[~culls].detach())
+ self._features_rest = Parameter(self._features_rest[~culls].detach())
+ self._opacities = Parameter(self._opacities[~culls].detach())
+ self._velocity = Parameter(self._velocity[~culls].detach())
+ self._taus = Parameter(self._taus[~culls].detach())
+ self._betas = Parameter(self._betas[~culls].detach())
+
+ print(f" Cull: {n_bef - self.num_points}")
+ return culls
+
+ def split_gaussians(self, split_mask, samps):
+ """
+ This function splits gaussians that are too large
+ """
+
+ n_splits = split_mask.sum().item()
+ print(f" Split: {n_splits}")
+ centered_samples = torch.randn((samps * n_splits, 3), device=self.device) # Nx3 of axis-aligned scales
+ scaled_samples = (
+ self.get_scaling[split_mask].repeat(samps, 1) * centered_samples
+ # torch.exp(self._scales[split_mask].repeat(samps, 1)) * centered_samples
+ ) # how these scales are rotated
+ quats = self.quat_act(self._quats[split_mask]) # normalize them first
+ rots = quat_to_rotmat(quats.repeat(samps, 1)) # how these scales are rotated
+ rotated_samples = torch.bmm(rots, scaled_samples[..., None]).squeeze()
+ new_means = rotated_samples + self._means[split_mask].repeat(samps, 1)
+ # step 2, sample new colors
+ # new_colors_all = self.colors_all[split_mask].repeat(samps, 1, 1)
+ new_feature_dc = self._features_dc[split_mask].repeat(samps, 1)
+ new_feature_rest = self._features_rest[split_mask].repeat(samps, 1, 1)
+ # step 3, sample new opacities
+ new_opacities = self._opacities[split_mask].repeat(samps, 1)
+ # step 4, sample new scales
+ size_fac = 1.6
+ new_scales = torch.log(torch.exp(self._scales[split_mask]) / size_fac).repeat(samps, 1)
+ self._scales[split_mask] = torch.log(torch.exp(self._scales[split_mask]) / size_fac)
+ # step 5, sample new quats
+ new_quats = self._quats[split_mask].repeat(samps, 1)
+ # step 6, sample new temporal properties
+ # new_velocity = self._velocity[split_mask].repeat(samps, 1)
+ # new_taus = self._taus[split_mask].repeat(samps, 1)
+ # new_betas = self._betas[split_mask].repeat(samps, 1)
+ stds_t = self.get_scaling_t[split_mask].repeat(samps, 1)
+ means_t = torch.zeros((stds_t.size(0), 1), device="cuda")
+ samples_t = torch.normal(mean=means_t, std=stds_t)
+ new_taus = samples_t+self._taus[split_mask].repeat(samps, 1)
+
+ new_betas = torch.log(self.get_scaling_t[split_mask].repeat(samps, 1) / size_fac)
+ new_velocity = self._velocity[split_mask].repeat(samps, 1)
+ new_means = new_means + self.velocity[split_mask].repeat(samps, 1) * (samples_t)
+
+ not_split_xyz_mask = (
+ self.get_scaling.max(dim=-1).values <= \
+ self.ctrl_cfg.densify_size_thresh * self.scene_scale * self.gamma
+ ).squeeze()[split_mask]
+ new_scales[not_split_xyz_mask.repeat(samps)] = torch.log(
+ self.get_scaling[split_mask].repeat(samps, 1)
+ )[not_split_xyz_mask.repeat(samps)]
+
+ not_split_t_mask = (torch.max(self.get_scaling_t, dim=1).values <= self.ctrl_cfg.densify_t_size_thresh)[split_mask]
+ new_betas[not_split_t_mask.repeat(samps)] = torch.log(
+ self.get_scaling_t[split_mask].repeat(samps, 1)
+ )[not_split_t_mask.repeat(samps)]
+
+ if self.ctrl_cfg.no_time_split:
+ new_betas = torch.log(self.get_scaling_t[split_mask].repeat(samps, 1))
+ return new_means, new_feature_dc, new_feature_rest, new_opacities, new_scales, new_quats, new_velocity, new_taus, new_betas
+
+ def dup_gaussians(self, dup_mask):
+ """
+ This function duplicates gaussians that are too small
+ """
+ n_dups = dup_mask.sum().item()
+ print(f" Dup: {n_dups}")
+ dup_means = self._means[dup_mask]
+ # dup_colors = self.colors_all[dup_mask]
+ dup_feature_dc = self._features_dc[dup_mask]
+ dup_feature_rest = self._features_rest[dup_mask]
+ dup_opacities = self._opacities[dup_mask]
+ dup_scales = self._scales[dup_mask]
+ dup_quats = self._quats[dup_mask]
+ dup_velocity = self._velocity[dup_mask]
+ dup_taus = self._taus[dup_mask]
+ dup_betas = self._betas[dup_mask]
+ return dup_means, dup_feature_dc, dup_feature_rest, dup_opacities, dup_scales, dup_quats, dup_velocity, dup_taus, dup_betas
+
+ def get_gaussians(self, cam: dataclass_camera) -> Dict:
+ # set time and smooth strategy
+ scaled_train_t = self.normalized_timestamps[self.cur_frame] * self.train_time_scale # t2 in paper
+ if self.training and (
+ self.ctrl_cfg.enable_temporal_smoothing and random.random() < self.ctrl_cfg.smooth_probability
+ ):
+ self.in_smooth = True
+ bound = self.normalized_time_interval * self.ctrl_cfg.distribution_span * self.train_time_scale
+ self.cur_time = scaled_train_t + uniform.Uniform(-bound, bound).sample((1,)).item() # t1 in paper
+ self.delta_t = scaled_train_t - self.cur_time # t2 - t1
+ else:
+ self.in_smooth = False
+ self.cur_time = scaled_train_t
+ self.delta_t = 0.0
+
+ filter_mask = (self.get_marginal_t > 0.05).squeeze()
+ self.filter_mask = filter_mask
+
+ means = self.temporal_means
+ activated_opacities = self.temporal_opacities
+ activated_scales = self.get_scaling
+ activated_rotations = self.get_quats
+
+ # get colors of gaussians
+ colors = torch.cat((self._features_dc[:, None, :], self._features_rest), dim=1)
+ if self.sh_degree > 0:
+ viewdirs = means.detach() - cam.camtoworlds.data[..., :3, 3] # (N, 3)
+ viewdirs = viewdirs / viewdirs.norm(dim=-1, keepdim=True)
+ n = min(self.step // self.ctrl_cfg.sh_degree_interval, self.sh_degree)
+ rgbs = spherical_harmonics(n, viewdirs, colors)
+ rgbs = torch.clamp(rgbs + 0.5, 0.0, 1.0)
+ else:
+ rgbs = torch.sigmoid(colors[:, 0, :])
+ actovated_colors = rgbs
+
+ # collect gaussians information
+ gs_dict = dict(
+ _means=means[filter_mask],
+ _opacities=activated_opacities[filter_mask],
+ _rgbs=actovated_colors[filter_mask],
+ _scales=activated_scales[filter_mask],
+ _quats=activated_rotations[filter_mask],
+ )
+
+ # check nan and inf in gs_dict
+ for k, v in gs_dict.items():
+ if torch.isnan(v).any():
+ raise ValueError(f"NaN detected in gaussian {k} at step {self.step}")
+ if torch.isinf(v).any():
+ raise ValueError(f"Inf detected in gaussian {k} at step {self.step}")
+
+ return gs_dict
+
+ def compute_reg_loss(self):
+ loss_dict = super().compute_reg_loss()
+
+ velocity_reg = self.reg_cfg.get("velocity_reg", None)
+ # use per point velocity regularization to replace per image velocity regularization
+ if velocity_reg:
+ if (self.cur_radii > 0).sum():
+ velocity_loss = self.velocity.norm(dim=-1)[self.filter_mask]
+ loss_dict["velocity_reg"] = velocity_loss[self.cur_radii > 0].mean() * velocity_reg.w
+ return loss_dict
+
+ def load_state_dict(self, state_dict: Dict, **kwargs) -> str:
+ N = state_dict["_means"].shape[0]
+ self._means = Parameter(torch.zeros((N,) + self._means.shape[1:], device=self.device))
+ self._scales = Parameter(torch.zeros((N,) + self._scales.shape[1:], device=self.device))
+ self._quats = Parameter(torch.zeros((N,) + self._quats.shape[1:], device=self.device))
+ self._features_dc = Parameter(torch.zeros((N,) + self._features_dc.shape[1:], device=self.device))
+ self._features_rest = Parameter(torch.zeros((N,) + self._features_rest.shape[1:], device=self.device))
+ self._opacities = Parameter(torch.zeros((N,) + self._opacities.shape[1:], device=self.device))
+ self._taus = Parameter(torch.zeros((N,) + self._taus.shape[1:], device=self.device))
+ self._betas = Parameter(torch.zeros((N,) + self._betas.shape[1:], device=self.device))
+ self._velocity = Parameter(torch.zeros((N,) + self._velocity.shape[1:], device=self.device))
+ msg = super().load_state_dict(state_dict, **kwargs)
+ return msg
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/models/gaussians/scaffold.py b/worldbench/third_party/drivestudio/models/gaussians/scaffold.py
new file mode 100644
index 0000000..a1ab2ce
--- /dev/null
+++ b/worldbench/third_party/drivestudio/models/gaussians/scaffold.py
@@ -0,0 +1,23 @@
+from typing import Dict, List
+from omegaconf import OmegaConf
+import logging
+
+import torch
+import torch.nn as nn
+from torch.nn import Parameter
+
+from models.gaussians.basics import *
+
+logger = logging.getLogger()
+
+class ScaffoldGaussians(nn.Module):
+ def __init__(self,
+ ctrl_cfg: OmegaConf,
+ reg_cfg: OmegaConf,
+ scene_scale: float = 1.0,
+ device: torch.device = torch.device("cuda")):
+ super().__init__()
+ self.ctrl_cfg = ctrl_cfg
+ self.reg_cfg = reg_cfg
+ self.scene_scale = scene_scale
+ self.device = device
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/models/gaussians/vanilla.py b/worldbench/third_party/drivestudio/models/gaussians/vanilla.py
new file mode 100644
index 0000000..6cccf18
--- /dev/null
+++ b/worldbench/third_party/drivestudio/models/gaussians/vanilla.py
@@ -0,0 +1,476 @@
+"""
+Filename: 3dgs.py
+
+Author: Ziyu Chen (ziyu.sjtu@gmail.com)
+
+Description:
+Unofficial implementation of 3DGS based on the work by Bernhard Kerbl, Georgios Kopanas, Thomas Leimkühler, and George Drettakis.
+This implementation is modified from the nerfstudio GaussianSplattingModel.
+
+- Original work by Bernhard Kerbl, Georgios Kopanas, Thomas Leimkühler, and George Drettakis.
+- Codebase reference: nerfstudio GaussianSplattingModel (https://github.com/nerfstudio-project/nerfstudio/blob/gaussian-splatting/nerfstudio/models/gaussian_splatting.py)
+
+Original paper: https://arxiv.org/abs/2308.04079
+"""
+
+from typing import Dict, List, Tuple
+from omegaconf import OmegaConf
+import logging
+
+import torch
+import torch.nn as nn
+from torch.nn import Parameter
+
+from models.gaussians.basics import *
+
+logger = logging.getLogger()
+
+class VanillaGaussians(nn.Module):
+
+ def __init__(
+ self,
+ class_name: str,
+ ctrl: OmegaConf,
+ reg: OmegaConf = None,
+ networks: OmegaConf = None,
+ scene_scale: float = 30.,
+ scene_origin: torch.Tensor = torch.zeros(3),
+ num_train_images: int = 300,
+ device: torch.device = torch.device("cuda"),
+ **kwargs
+ ):
+ super().__init__()
+ self.class_prefix = class_name + "#"
+ self.ctrl_cfg = ctrl
+ self.reg_cfg = reg
+ self.networks_cfg = networks
+ self.scene_scale = scene_scale
+ self.scene_origin = scene_origin
+ self.num_train_images = num_train_images
+ self.step = 0
+
+ self.device = device
+ self.ball_gaussians=self.ctrl_cfg.get("ball_gaussians", False)
+ self.gaussian_2d = self.ctrl_cfg.get("gaussian_2d", False)
+
+ # for evaluation
+ self.in_test_set = False
+
+ # init models
+ self.xys_grad_norm = None
+ self.max_2Dsize = None
+ self._means = torch.zeros(1, 3, device=self.device)
+ if self.ball_gaussians:
+ self._scales = torch.zeros(1, 1, device=self.device)
+ else:
+ if self.gaussian_2d:
+ self._scales = torch.zeros(1, 2, device=self.device)
+ else:
+ self._scales = torch.zeros(1, 3, device=self.device)
+ self._quats = torch.zeros(1, 4, device=self.device)
+ self._opacities = torch.zeros(1, 1, device=self.device)
+ self._features_dc = torch.zeros(1, 3, device=self.device)
+ self._features_rest = torch.zeros(1, num_sh_bases(self.sh_degree) - 1, 3, device=self.device)
+
+ @property
+ def sh_degree(self):
+ return self.ctrl_cfg.sh_degree
+
+ def create_from_pcd(self, init_means: torch.Tensor, init_colors: torch.Tensor) -> None:
+ self._means = Parameter(init_means)
+
+ distances_np, _ = k_nearest_sklearn(self._means.data, 3)
+ distances = torch.from_numpy(distances_np)
+ target_dim = 1 if self.ball_gaussians else (2 if self.gaussian_2d else 3)
+
+ if distances.shape[1] > 0:
+ avg_dist = distances.mean(dim=-1, keepdim=True).to(self.device)
+ scales_init = avg_dist.expand(-1, target_dim).clone()
+ else:
+ # fallback for extremely sparse point clouds
+ default_scale = 1e-3
+ scales_init = torch.full((self.num_points, target_dim), default_scale, device=self.device)
+
+ self._scales = Parameter(torch.log(scales_init))
+ self._quats = Parameter(random_quat_tensor(self.num_points).to(self.device))
+ dim_sh = num_sh_bases(self.sh_degree)
+
+ fused_color = RGB2SH(init_colors) # float range [0, 1]
+ shs = torch.zeros((fused_color.shape[0], dim_sh, 3)).float().to(self.device)
+ if self.sh_degree > 0:
+ shs[:, 0, :3] = fused_color
+ shs[:, 1:, 3:] = 0.0
+ else:
+ shs[:, 0, :3] = torch.logit(init_colors, eps=1e-10)
+ self._features_dc = Parameter(shs[:, 0, :])
+ self._features_rest = Parameter(shs[:, 1:, :])
+ self._opacities = Parameter(torch.logit(0.1 * torch.ones(self.num_points, 1, device=self.device)))
+
+ @property
+ def colors(self):
+ if self.sh_degree > 0:
+ return SH2RGB(self._features_dc)
+ else:
+ return torch.sigmoid(self._features_dc)
+ @property
+ def shs_0(self):
+ return self._features_dc
+ @property
+ def shs_rest(self):
+ return self._features_rest
+ @property
+ def num_points(self):
+ return self._means.shape[0]
+ @property
+ def get_scaling(self):
+ if self.ball_gaussians:
+ if self.gaussian_2d:
+ scaling = torch.exp(self._scales).repeat(1, 2)
+ scaling = torch.cat([scaling, torch.zeros_like(scaling[..., :1])], dim=-1)
+ return scaling
+ else:
+ return torch.exp(self._scales).repeat(1, 3)
+ else:
+ if self.gaussian_2d:
+ scaling = torch.exp(self._scales)
+ scaling = torch.cat([scaling[..., :2], torch.zeros_like(scaling[..., :1])], dim=-1)
+ return scaling
+ else:
+ return torch.exp(self._scales)
+ @property
+ def get_opacity(self):
+ return torch.sigmoid(self._opacities)
+ @property
+ def get_quats(self):
+ return self.quat_act(self._quats)
+
+ def quat_act(self, x: torch.Tensor) -> torch.Tensor:
+ return x / x.norm(dim=-1, keepdim=True)
+
+ def preprocess_per_train_step(self, step: int):
+ self.step = step
+
+ def postprocess_per_train_step(
+ self,
+ step: int,
+ optimizer: torch.optim.Optimizer,
+ radii: torch.Tensor,
+ xys_grad: torch.Tensor,
+ last_size: int,
+ ) -> None:
+ self.after_train(radii, xys_grad, last_size)
+ if step % self.ctrl_cfg.refine_interval == 0:
+ self.refinement_after(step, optimizer)
+
+ def after_train(
+ self,
+ radii: torch.Tensor,
+ xys_grad: torch.Tensor,
+ last_size: int,
+ ) -> None:
+ with torch.no_grad():
+ # keep track of a moving average of grad norms
+ visible_mask = (radii > 0).flatten()
+ full_mask = torch.zeros(self.num_points, device=radii.device, dtype=torch.bool)
+ full_mask[self.filter_mask] = visible_mask
+
+ grads = xys_grad.norm(dim=-1)
+ if self.xys_grad_norm is None:
+ self.xys_grad_norm = torch.zeros(self.num_points, device=grads.device, dtype=grads.dtype)
+ self.xys_grad_norm[self.filter_mask] = grads
+ self.vis_counts = torch.ones_like(self.xys_grad_norm)
+ else:
+ assert self.vis_counts is not None
+ self.vis_counts[full_mask] = self.vis_counts[full_mask] + 1
+ self.xys_grad_norm[full_mask] = grads[visible_mask] + self.xys_grad_norm[full_mask]
+
+ # update the max screen size, as a ratio of number of pixels
+ if self.max_2Dsize is None:
+ self.max_2Dsize = torch.zeros(self.num_points, device=radii.device, dtype=torch.float32)
+ newradii = radii[visible_mask]
+ self.max_2Dsize[full_mask] = torch.maximum(
+ self.max_2Dsize[full_mask], newradii / float(last_size)
+ )
+
+ def get_gaussian_param_groups(self) -> Dict[str, List[Parameter]]:
+ return {
+ self.class_prefix+"xyz": [self._means],
+ self.class_prefix+"sh_dc": [self._features_dc],
+ self.class_prefix+"sh_rest": [self._features_rest],
+ self.class_prefix+"opacity": [self._opacities],
+ self.class_prefix+"scaling": [self._scales],
+ self.class_prefix+"rotation": [self._quats],
+ }
+
+ def get_param_groups(self) -> Dict[str, List[Parameter]]:
+ return self.get_gaussian_param_groups()
+
+ def refinement_after(self, step, optimizer: torch.optim.Optimizer) -> None:
+ assert step == self.step
+ if self.step <= self.ctrl_cfg.warmup_steps:
+ return
+ with torch.no_grad():
+ # only split/cull if we've seen every image since opacity reset
+ reset_interval = self.ctrl_cfg.reset_alpha_interval
+ do_densification = (
+ self.step < self.ctrl_cfg.stop_split_at
+ and self.step % reset_interval > max(self.num_train_images, self.ctrl_cfg.refine_interval)
+ )
+ # split & duplicate
+ print(f"Class {self.class_prefix} current points: {self.num_points} @ step {self.step}")
+ if do_densification:
+ assert self.xys_grad_norm is not None and self.vis_counts is not None and self.max_2Dsize is not None
+
+ avg_grad_norm = self.xys_grad_norm / self.vis_counts
+ high_grads = (avg_grad_norm > self.ctrl_cfg.densify_grad_thresh).squeeze()
+
+ splits = (
+ self.get_scaling.max(dim=-1).values > \
+ self.ctrl_cfg.densify_size_thresh * self.scene_scale
+ ).squeeze()
+ if self.step < self.ctrl_cfg.stop_screen_size_at:
+ splits |= (self.max_2Dsize > self.ctrl_cfg.split_screen_size).squeeze()
+ splits &= high_grads
+ nsamps = self.ctrl_cfg.n_split_samples
+ (
+ split_means,
+ split_feature_dc,
+ split_feature_rest,
+ split_opacities,
+ split_scales,
+ split_quats,
+ ) = self.split_gaussians(splits, nsamps)
+
+ dups = (
+ self.get_scaling.max(dim=-1).values <= \
+ self.ctrl_cfg.densify_size_thresh * self.scene_scale
+ ).squeeze()
+ dups &= high_grads
+ (
+ dup_means,
+ dup_feature_dc,
+ dup_feature_rest,
+ dup_opacities,
+ dup_scales,
+ dup_quats,
+ ) = self.dup_gaussians(dups)
+
+ self._means = Parameter(torch.cat([self._means.detach(), split_means, dup_means], dim=0))
+ # self.colors_all = Parameter(torch.cat([self.colors_all.detach(), split_colors, dup_colors], dim=0))
+ self._features_dc = Parameter(torch.cat([self._features_dc.detach(), split_feature_dc, dup_feature_dc], dim=0))
+ self._features_rest = Parameter(torch.cat([self._features_rest.detach(), split_feature_rest, dup_feature_rest], dim=0))
+ self._opacities = Parameter(torch.cat([self._opacities.detach(), split_opacities, dup_opacities], dim=0))
+ self._scales = Parameter(torch.cat([self._scales.detach(), split_scales, dup_scales], dim=0))
+ self._quats = Parameter(torch.cat([self._quats.detach(), split_quats, dup_quats], dim=0))
+
+ # append zeros to the max_2Dsize tensor
+ self.max_2Dsize = torch.cat(
+ [self.max_2Dsize, torch.zeros_like(split_scales[:, 0]), torch.zeros_like(dup_scales[:, 0])],
+ dim=0,
+ )
+
+ split_idcs = torch.where(splits)[0]
+ param_groups = self.get_gaussian_param_groups()
+ dup_in_optim(optimizer, split_idcs, param_groups, n=nsamps)
+
+ dup_idcs = torch.where(dups)[0]
+ param_groups = self.get_gaussian_param_groups()
+ dup_in_optim(optimizer, dup_idcs, param_groups, 1)
+
+ # cull NOTE: Offset all the opacity reset logic by refine_every so that we don't
+ # save checkpoints right when the opacity is reset (saves every 2k)
+ if self.step % reset_interval > max(self.num_train_images, self.ctrl_cfg.refine_interval):
+ deleted_mask = self.cull_gaussians()
+ param_groups = self.get_gaussian_param_groups()
+ remove_from_optim(optimizer, deleted_mask, param_groups)
+ print(f"Class {self.class_prefix} left points: {self.num_points}")
+
+ # reset opacity
+ if self.step % reset_interval == self.ctrl_cfg.refine_interval:
+ # NOTE: in nerfstudio, reset_value = cull_alpha_thresh * 0.8
+ # we align to original repo of gaussians spalting
+ reset_value = torch.min(self.get_opacity.data,
+ torch.ones_like(self._opacities.data) * self.ctrl_cfg.reset_alpha_value)
+ self._opacities.data = torch.logit(reset_value)
+ # reset the exp of optimizer
+ for group in optimizer.param_groups:
+ if group["name"] == self.class_prefix+"opacity":
+ old_params = group["params"][0]
+ param_state = optimizer.state[old_params]
+ param_state["exp_avg"] = torch.zeros_like(param_state["exp_avg"])
+ param_state["exp_avg_sq"] = torch.zeros_like(param_state["exp_avg_sq"])
+ self.xys_grad_norm = None
+ self.vis_counts = None
+ self.max_2Dsize = None
+
+ def cull_gaussians(self):
+ """
+ This function deletes gaussians with under a certain opacity threshold
+ """
+ n_bef = self.num_points
+ # cull transparent ones
+ culls = (self.get_opacity.data < self.ctrl_cfg.cull_alpha_thresh).squeeze()
+ if self.step > self.ctrl_cfg.reset_alpha_interval:
+ # cull huge ones
+ toobigs = (
+ torch.exp(self._scales).max(dim=-1).values >
+ self.ctrl_cfg.cull_scale_thresh * self.scene_scale
+ ).squeeze()
+ culls = culls | toobigs
+ if self.step < self.ctrl_cfg.stop_screen_size_at:
+ # cull big screen space
+ assert self.max_2Dsize is not None
+ culls = culls | (self.max_2Dsize > self.ctrl_cfg.cull_screen_size).squeeze()
+ self._means = Parameter(self._means[~culls].detach())
+ self._scales = Parameter(self._scales[~culls].detach())
+ self._quats = Parameter(self._quats[~culls].detach())
+ # self.colors_all = Parameter(self.colors_all[~culls].detach())
+ self._features_dc = Parameter(self._features_dc[~culls].detach())
+ self._features_rest = Parameter(self._features_rest[~culls].detach())
+ self._opacities = Parameter(self._opacities[~culls].detach())
+
+ print(f" Cull: {n_bef - self.num_points}")
+ return culls
+
+ def split_gaussians(self, split_mask: torch.Tensor, samps: int) -> Tuple:
+ """
+ This function splits gaussians that are too large
+ """
+
+ n_splits = split_mask.sum().item()
+ print(f" Split: {n_splits}")
+ centered_samples = torch.randn((samps * n_splits, 3), device=self.device) # Nx3 of axis-aligned scales
+ scaled_samples = (
+ self.get_scaling[split_mask].repeat(samps, 1) * centered_samples
+ # torch.exp(self._scales[split_mask].repeat(samps, 1)) * centered_samples
+ ) # how these scales are rotated
+ quats = self.quat_act(self._quats[split_mask]) # normalize them first
+ rots = quat_to_rotmat(quats.repeat(samps, 1)) # how these scales are rotated
+ rotated_samples = torch.bmm(rots, scaled_samples[..., None]).squeeze()
+ new_means = rotated_samples + self._means[split_mask].repeat(samps, 1)
+ # step 2, sample new colors
+ # new_colors_all = self.colors_all[split_mask].repeat(samps, 1, 1)
+ new_feature_dc = self._features_dc[split_mask].repeat(samps, 1)
+ new_feature_rest = self._features_rest[split_mask].repeat(samps, 1, 1)
+ # step 3, sample new opacities
+ new_opacities = self._opacities[split_mask].repeat(samps, 1)
+ # step 4, sample new scales
+ size_fac = 1.6
+ new_scales = torch.log(torch.exp(self._scales[split_mask]) / size_fac).repeat(samps, 1)
+ self._scales[split_mask] = torch.log(torch.exp(self._scales[split_mask]) / size_fac)
+ # step 5, sample new quats
+ new_quats = self._quats[split_mask].repeat(samps, 1)
+ return new_means, new_feature_dc, new_feature_rest, new_opacities, new_scales, new_quats
+
+ def dup_gaussians(self, dup_mask: torch.Tensor) -> Tuple:
+ """
+ This function duplicates gaussians that are too small
+ """
+ n_dups = dup_mask.sum().item()
+ print(f" Dup: {n_dups}")
+ dup_means = self._means[dup_mask]
+ # dup_colors = self.colors_all[dup_mask]
+ dup_feature_dc = self._features_dc[dup_mask]
+ dup_feature_rest = self._features_rest[dup_mask]
+ dup_opacities = self._opacities[dup_mask]
+ dup_scales = self._scales[dup_mask]
+ dup_quats = self._quats[dup_mask]
+ return dup_means, dup_feature_dc, dup_feature_rest, dup_opacities, dup_scales, dup_quats
+
+ def get_gaussians(self, cam: dataclass_camera) -> Dict:
+ filter_mask = torch.ones_like(self._means[:, 0], dtype=torch.bool)
+ self.filter_mask = filter_mask
+
+ # get colors of gaussians
+ colors = torch.cat((self._features_dc[:, None, :], self._features_rest), dim=1)
+ if self.sh_degree > 0:
+ viewdirs = self._means.detach() - cam.camtoworlds.data[..., :3, 3] # (N, 3)
+ viewdirs = viewdirs / viewdirs.norm(dim=-1, keepdim=True)
+ n = min(self.step // self.ctrl_cfg.sh_degree_interval, self.sh_degree)
+ rgbs = spherical_harmonics(n, viewdirs, colors)
+ rgbs = torch.clamp(rgbs + 0.5, 0.0, 1.0)
+ else:
+ rgbs = torch.sigmoid(colors[:, 0, :])
+
+ activated_opacities = self.get_opacity
+ activated_scales = self.get_scaling
+ activated_rotations = self.get_quats
+ actovated_colors = rgbs
+
+ # collect gaussians information
+ gs_dict = dict(
+ _means=self._means[filter_mask],
+ _opacities=activated_opacities[filter_mask],
+ _rgbs=actovated_colors[filter_mask],
+ _scales=activated_scales[filter_mask],
+ _quats=activated_rotations[filter_mask],
+ )
+
+ # check nan and inf in gs_dict
+ for k, v in gs_dict.items():
+ if torch.isnan(v).any():
+ raise ValueError(f"NaN detected in gaussian {k} at step {self.step}")
+ if torch.isinf(v).any():
+ raise ValueError(f"Inf detected in gaussian {k} at step {self.step}")
+
+ return gs_dict
+
+ def compute_reg_loss(self):
+ loss_dict = {}
+ sharp_shape_reg_cfg = self.reg_cfg.get("sharp_shape_reg", None)
+ if sharp_shape_reg_cfg is not None:
+ w = sharp_shape_reg_cfg.w
+ max_gauss_ratio = sharp_shape_reg_cfg.max_gauss_ratio
+ step_interval = sharp_shape_reg_cfg.step_interval
+ if self.step % step_interval == 0:
+ # scale regularization
+ scale_exp = self.get_scaling
+ scale_reg = torch.maximum(scale_exp.amax(dim=-1) / scale_exp.amin(dim=-1), torch.tensor(max_gauss_ratio)) - max_gauss_ratio
+ scale_reg = scale_reg.mean() * w
+ loss_dict["sharp_shape_reg"] = scale_reg
+
+ flatten_reg = self.reg_cfg.get("flatten", None)
+ if flatten_reg is not None:
+ sclaings = self.get_scaling
+ min_scale, _ = torch.min(sclaings, dim=1)
+ min_scale = torch.clamp(min_scale, 0, 30)
+ flatten_loss = torch.abs(min_scale).mean()
+ loss_dict["flatten"] = flatten_loss * flatten_reg.w
+
+ sparse_reg = self.reg_cfg.get("sparse_reg", None)
+ if sparse_reg:
+ if (self.cur_radii > 0).sum():
+ opacity = torch.sigmoid(self._opacities)
+ opacity = opacity.clamp(1e-6, 1-1e-6)
+ log_opacity = opacity * torch.log(opacity)
+ log_one_minus_opacity = (1-opacity) * torch.log(1 - opacity)
+ sparse_loss = -1 * (log_opacity + log_one_minus_opacity)[self.cur_radii > 0].mean()
+ loss_dict["sparse_reg"] = sparse_loss * sparse_reg.w
+
+ # compute the max of scaling
+ max_s_square_reg = self.reg_cfg.get("max_s_square_reg", None)
+ if max_s_square_reg is not None and not self.ball_gaussians:
+ loss_dict["max_s_square"] = torch.mean((self.get_scaling.max(dim=1).values) ** 2) * max_s_square_reg.w
+ return loss_dict
+
+ def load_state_dict(self, state_dict: Dict, **kwargs) -> str:
+ N = state_dict["_means"].shape[0]
+ self._means = Parameter(torch.zeros((N,) + self._means.shape[1:], device=self.device))
+ self._scales = Parameter(torch.zeros((N,) + self._scales.shape[1:], device=self.device))
+ self._quats = Parameter(torch.zeros((N,) + self._quats.shape[1:], device=self.device))
+ self._features_dc = Parameter(torch.zeros((N,) + self._features_dc.shape[1:], device=self.device))
+ self._features_rest = Parameter(torch.zeros((N,) + self._features_rest.shape[1:], device=self.device))
+ self._opacities = Parameter(torch.zeros((N,) + self._opacities.shape[1:], device=self.device))
+ msg = super().load_state_dict(state_dict, **kwargs)
+ return msg
+
+ def export_gaussians_to_ply(self, alpha_thresh: float) -> Dict:
+ means = self._means
+ direct_color = self.colors
+
+ activated_opacities = self.get_opacity
+ mask = activated_opacities.squeeze() > alpha_thresh
+ return {
+ "positions": means[mask],
+ "colors": direct_color[mask],
+ }
diff --git a/worldbench/third_party/drivestudio/models/human_body.py b/worldbench/third_party/drivestudio/models/human_body.py
new file mode 100644
index 0000000..b4f9a6e
--- /dev/null
+++ b/worldbench/third_party/drivestudio/models/human_body.py
@@ -0,0 +1,362 @@
+# Acknowledgement: https://github.com/JiahuiLei/GART
+from typing import Optional
+import os
+import sys
+import pickle
+import trimesh
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+from pytorch3d.transforms import (
+ matrix_to_quaternion,
+ quaternion_to_matrix,
+ axis_angle_to_matrix
+)
+
+from models.modules import VoxelDeformer
+from third_party.smplx.smplx import SMPLLayer
+from third_party.smplx.smplx.utils import SMPLOutput
+from third_party.smplx.smplx.lbs import vertices2joints, batch_rigid_transform
+
+def blockPrinting(func):
+ def func_wrapper(*args, **kwargs):
+ # block all printing to the console
+ sys.stdout = open(os.devnull, 'w')
+ # call the method in question
+ value = func(*args, **kwargs)
+ # enable all printing to the console
+ sys.stdout = sys.__stdout__
+ # pass the return value of the method back
+ return value
+
+ return func_wrapper
+
+class SMPL(SMPLLayer):
+
+ @blockPrinting
+ def __init__(self, *args, joint_regressor_extra: Optional[str] = None, **kwargs):
+ """
+ Extension of the official SMPL implementation to support more joints.
+ Args:
+ Same as SMPLLayer.
+ joint_regressor_extra (str): Path to extra joint regressor.
+ """
+ super(SMPL, self).__init__(*args, **kwargs)
+ smpl_to_openpose = [24, 12, 17, 19, 21, 16, 18, 20, 0, 2, 5, 8, 1, 4,
+ 7, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34]
+
+ if joint_regressor_extra is not None:
+ self.register_buffer('joint_regressor_extra', torch.tensor(pickle.load(open(joint_regressor_extra, 'rb'), encoding='latin1'), dtype=torch.float32))
+ self.register_buffer('joint_map', torch.tensor(smpl_to_openpose, dtype=torch.long))
+
+ def forward(self, *args, **kwargs) -> SMPLOutput:
+ """
+ Run forward pass. Same as SMPL and also append an extra set of joints if joint_regressor_extra is specified.
+ """
+ smpl_output = super(SMPL, self).forward(*args, **kwargs)
+ joints = smpl_output.joints[:, self.joint_map, :]
+ if hasattr(self, 'joint_regressor_extra'):
+ extra_joints = vertices2joints(self.joint_regressor_extra, smpl_output.vertices)
+ joints = torch.cat([joints, extra_joints], dim=1)
+ smpl_output.joints = joints
+ return smpl_output
+
+def get_predefined_human_rest_pose(pose_type):
+ print(f"Using predefined pose: {pose_type}")
+ body_pose_t = torch.zeros((1, 69))
+ if pose_type.lower() == "da_pose":
+ body_pose_t[:, 2] = torch.pi / 6
+ body_pose_t[:, 5] = -torch.pi / 6
+ elif pose_type.lower() == "a_pose":
+ body_pose_t[:, 2] = 0.2
+ body_pose_t[:, 5] = -0.2
+ body_pose_t[:, 47] = -0.8
+ body_pose_t[:, 50] = 0.8
+ elif pose_type.lower() == "t_pose":
+ pass
+ else:
+ raise ValueError("Unknown cano_pose: {}".format(pose_type))
+ return body_pose_t.reshape(23, 3)
+
+class SMPLTemplate(nn.Module):
+ def __init__(self, smpl_model_path, num_human, init_beta, cano_pose_type, voxel_deformer_res=64, use_voxel_deformer=False, is_resume=False):
+ super().__init__()
+ assert num_human == init_beta.shape[0], "num_human should be the same as the number of beta"
+ self.num_human = num_human
+ self.dim = 24
+ self._template_layer = SMPLLayer(model_path=smpl_model_path)
+
+ init_beta = torch.as_tensor(init_beta, dtype=torch.float32).cpu()
+ self.register_buffer("init_beta", init_beta)
+ self.cano_pose_type = cano_pose_type
+ self.name = "smpl"
+
+ can_pose = get_predefined_human_rest_pose(cano_pose_type)
+ can_pose = axis_angle_to_matrix(torch.cat([torch.zeros(1, 3), can_pose], 0))
+ self.register_buffer("canonical_pose", can_pose)
+
+ init_smpl_output = self._template_layer(
+ betas=init_beta,
+ body_pose=can_pose[None, 1:].repeat(num_human, 1, 1, 1),
+ global_orient=can_pose[None, 0].repeat(num_human, 1, 1, 1),
+ return_full_pose=True,
+ )
+ J_canonical, A0 = init_smpl_output.J, init_smpl_output.A
+ A0_inv = torch.inverse(A0)
+ self.register_buffer("A0_inv", A0_inv)
+ self.register_buffer("J_canonical", J_canonical)
+
+ v_init = init_smpl_output.vertices # [B, 6890, 3]
+ W_init = self._template_layer.lbs_weights # [6890, 24]
+ self.register_buffer("W", W_init[None].repeat(num_human, 1, 1))
+
+ self.use_voxel_deformer = use_voxel_deformer
+ if self.use_voxel_deformer:
+ self.voxel_deformer = VoxelDeformer(
+ vtx=v_init,
+ vtx_features=W_init.unsqueeze(0).repeat(num_human, 1, 1), # [B, 6890, 24]
+ resolution_dhw=[
+ voxel_deformer_res // 4,
+ voxel_deformer_res,
+ voxel_deformer_res,
+ ],
+ is_resume=is_resume, # if is resume, the weight will be randomly initialized to save time
+ )
+
+ # * Important, record first joint position, because the global orientation is rotating using this joint position as center, so we can compute the action on later As
+ j0_t = init_smpl_output.joints[:, 0]
+ self.register_buffer("j0_t", j0_t)
+ return
+
+ def get_init_vf(self):
+ init_smpl_output = self._template_layer(
+ betas=self.init_beta,
+ body_pose=self.canonical_pose[None, 1:].repeat(self.num_human, 1, 1, 1),
+ global_orient=self.canonical_pose[None, 0].repeat(self.num_human, 1, 1, 1),
+ return_full_pose=True,
+ )
+ v_init = init_smpl_output.vertices # 1,6890,3
+ v_init = v_init
+ faces = self._template_layer.faces_tensor
+ return v_init, faces
+
+ def get_rot_action(self, axis_angle):
+ # apply this action to canonical additional bones
+ # axis_angle: B,3
+ assert axis_angle.ndim == 2 and axis_angle.shape[-1] == 3
+ B = len(axis_angle)
+ R = axis_angle_to_matrix(axis_angle) # B,3,3
+ I = torch.eye(3).to(R)[None].expand(B, -1, -1) # B,3,3
+ t0 = self.j0_t[None].expand(B, -1) # B,3
+ T = torch.eye(4).to(R)[None].expand(B, -1, -1) # B,4,4
+ T[:, :3, :3] = R
+ T[:, :3, 3] = torch.einsum("bij, bj -> bi", I - R, t0)
+ return T # B,4,4
+
+ def forward(self, masked_theta=None, xyz_canonical=None, instances_mask=None):
+ # skinning
+ if masked_theta is None:
+ A = None
+ else:
+ assert (
+ masked_theta.ndim == 3 and masked_theta.shape[-1] == 4
+ ), "pose should have shape Bx24x3, in axis-angle format"
+ nB = len(masked_theta)
+ _, A = batch_rigid_transform(
+ quaternion_to_matrix(masked_theta),
+ self.J_canonical[instances_mask],
+ self._template_layer.parents,
+ )
+ A = torch.einsum("bnij, bnjk->bnik", A, self.A0_inv[instances_mask]) # B,24,4,4
+
+ if xyz_canonical is None or not self.use_voxel_deformer:
+ # forward theta only
+ W = self.W
+ else:
+ W = self.voxel_deformer(xyz_canonical) # B,N,24+K
+ W = W[instances_mask]
+ return W, A
+
+ def remove_instance(self, ins_id):
+ ins_mask = torch.ones(self.num_human, dtype=torch.bool)
+ ins_mask[ins_id] = False
+ self.J_canonical = self.J_canonical[ins_mask]
+ self.W = self.W[ins_mask]
+ self.num_human -= 1
+
+ if self.use_voxel_deformer:
+ self.voxel_deformer.lbs_voxel_base = self.voxel_deformer.lbs_voxel_base[ins_mask]
+ self.voxel_deformer.voxel_w_correction = nn.Parameter(self.voxel_deformer.voxel_w_correction[ins_mask])
+ self.voxel_deformer.offset = self.voxel_deformer.offset[ins_mask]
+ self.voxel_deformer.scale = self.voxel_deformer.scale[ins_mask]
+
+ def add_instance(self, ins_id, new_dict):
+ voxel_deformer_dict = new_dict["voxel_deformer"]
+ self.J_canonical = torch.cat([self.J_canonical, new_dict['J_canonical'][None]], dim=0)
+ self.W = torch.cat([self.W, new_dict['W'][None]], dim=0)
+ self.num_human += 1
+
+ if self.use_voxel_deformer:
+ self.voxel_deformer.lbs_voxel_base = torch.cat([self.voxel_deformer.lbs_voxel_base, voxel_deformer_dict['lbs_voxel_base'][None]], dim=0)
+ self.voxel_deformer.voxel_w_correction = nn.Parameter(torch.cat([self.voxel_deformer.voxel_w_correction, voxel_deformer_dict['voxel_w_correction'][None]], dim=0))
+ self.voxel_deformer.offset = torch.cat([self.voxel_deformer.offset, voxel_deformer_dict['offset'][None]], dim=0)
+ self.voxel_deformer.scale = torch.cat([self.voxel_deformer.scale, voxel_deformer_dict['scale'][None]], dim=0)
+
+def init_xyz_on_mesh(v_init, faces, subdivide_num):
+ # * xyz
+ denser_v, denser_f = v_init.detach().cpu().numpy(), faces
+ for i in range(subdivide_num):
+ denser_v, denser_f = trimesh.remesh.subdivide(denser_v, denser_f)
+ body_mesh = trimesh.Trimesh(denser_v, denser_f, process=False)
+ v_init = torch.as_tensor(denser_v, dtype=torch.float32)
+ return v_init, body_mesh
+
+def init_qso_on_mesh(
+ body_mesh,
+ scale_init_factor,
+ thickness_init_factor,
+ max_scale,
+ min_scale,
+ s_inv_act,
+ opacity_base_logit,
+):
+ # * Quaternion
+ # each column is a basis vector
+ # the local frame is z to normal, xy on the disk
+ normal = body_mesh.vertex_normals.copy()
+ v_init = torch.as_tensor(body_mesh.vertices.copy())
+ faces = torch.as_tensor(body_mesh.faces.copy())
+
+ uz = torch.as_tensor(normal, dtype=torch.float32)
+ rand_dir = torch.randn_like(uz)
+ ux = F.normalize(torch.cross(uz, rand_dir, dim=-1), dim=-1)
+ uy = F.normalize(torch.cross(uz, ux, dim=-1), dim=-1)
+ frame = torch.stack([ux, uy, uz], dim=-1) # N,3,3
+ ret_q = matrix_to_quaternion(frame)
+
+ # * Scaling
+ xy = v_init[faces[:, 1]] - v_init[faces[:, 0]]
+ xz = v_init[faces[:, 2]] - v_init[faces[:, 0]]
+ area = torch.norm(torch.cross(xy, xz, dim=-1), dim=-1) / 2
+ vtx_nn_area = torch.zeros_like(v_init[:, 0])
+ for i in range(3):
+ vtx_nn_area.scatter_add_(0, faces[:, i], area / 3.0)
+ radius = torch.sqrt(vtx_nn_area / np.pi)
+ # radius = torch.clamp(radius * scale_init_factor, max=max_scale, min=min_scale)
+ # ! 2023.11.22, small eps
+ radius = torch.clamp(
+ radius * scale_init_factor, max=max_scale - 1e-4, min=min_scale + 1e-4
+ )
+ thickness = radius * thickness_init_factor
+ # ! 2023.11.22, small eps
+ thickness = torch.clamp(thickness, max=max_scale - 1e-4, min=min_scale + 1e-4)
+ radius_logit = s_inv_act(radius)
+ thickness_logit = s_inv_act(thickness)
+ ret_s = torch.stack([radius_logit, radius_logit, thickness_logit], dim=-1)
+
+ ret_o = torch.ones_like(v_init[:, :1]) * opacity_base_logit
+ return ret_q, ret_s, ret_o
+
+def get_on_mesh_init_geo_values(
+ template,
+ opacity_init_logit,
+ on_mesh_subdivide = 0,
+ scale_init_factor = 1.0,
+ thickness_init_factor = 0.5,
+ max_scale = 1.0,
+ min_scale = 0.0,
+ s_inv_act = torch.logit,
+):
+ v, f = template.get_init_vf()
+ x_all, q_all, s_all, o_all = [], [], [], []
+ for i in range(len(v)):
+ x, mesh = init_xyz_on_mesh(v[i], f, on_mesh_subdivide)
+ q, s, o = init_qso_on_mesh(
+ mesh,
+ scale_init_factor,
+ thickness_init_factor,
+ max_scale,
+ min_scale,
+ s_inv_act,
+ opacity_init_logit,
+ )
+
+ x_all.append(x)
+ q_all.append(q)
+ s_all.append(s)
+ o_all.append(o)
+
+ x_all = torch.cat(x_all, dim=0)
+ q_all = torch.cat(q_all, dim=0)
+ s_all = torch.cat(s_all, dim=0)
+ o_all = torch.cat(o_all, dim=0)
+ return x_all, q_all, s_all, o_all
+
+phalp_colors =[
+ [213.0,255.0,0.0,],
+ [255.0,0.0,86.0,],
+ [158.0,0.0,142.0,],
+ [14.0,76.0,161.0,],
+ [255.0,229.0,2.0,],
+ [0.0,95.0,57.0,],
+ [0.0,255.0,0.0,],
+ [149.0,0.0,58.0,],
+ [255.0,147.0,126.0,],
+ [164.0,36.0,0.0,],
+ [0.0,21.0,68.0,],
+ [145.0,208.0,203.0,],
+ [98.0,14.0,0.0,],
+ [107.0,104.0,130.0,],
+ [0.0,0.0,255.0,],
+ [0.0,125.0,181.0,],
+ [106.0,130.0,108.0,],
+ [0.0,174.0,126.0,],
+ [194.0,140.0,159.0,],
+ [190.0,153.0,112.0,],
+ [0.0,143.0,156.0,],
+ [95.0,173.0,78.0,],
+ [255.0,0.0,0.0,],
+ [255.0,0.0,246.0,],
+ [255.0,2.0,157.0,],
+ [104.0,61.0,59.0,],
+ [255.0,116.0,163.0,],
+ [150.0,138.0,232.0,],
+ [152.0,255.0,82.0,],
+ [167.0,87.0,64.0,],
+ [1.0,255.0,254.0,],
+ [255.0,238.0,232.0,],
+ [254.0,137.0,0.0,],
+ [189.0,198.0,255.0,],
+ [1.0,208.0,255.0,],
+ [187.0,136.0,0.0,],
+ [117.0,68.0,177.0,],
+ [165.0,255.0,210.0,],
+ [255.0,166.0,254.0,],
+ [119.0,77.0,0.0,],
+ [122.0,71.0,130.0,],
+ [38.0,52.0,0.0,],
+ [0.0,71.0,84.0,],
+ [67.0,0.0,44.0,],
+ [181.0,0.0,255.0,],
+ [255.0,177.0,103.0,],
+ [255.0,219.0,102.0,],
+ [144.0,251.0,146.0,],
+ [126.0,45.0,210.0,],
+ [189.0,211.0,147.0,],
+ [229.0,111.0,254.0,],
+ [222.0,255.0,116.0,],
+ [0.0,255.0,120.0,],
+ [0.0,155.0,255.0,],
+ [0.0,100.0,1.0,],
+ [0.0,118.0,255.0,],
+ [133.0,169.0,0.0,],
+ [0.0,185.0,23.0,],
+ [120.0,130.0,49.0,],
+ [0.0,255.0,198.0,],
+ [255.0,110.0,65.0,],
+ [232.0,94.0,190.0,],
+ [0.0,0.0,0.0,],
+ [1.0,0.0,103.0,],
+]
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/models/losses.py b/worldbench/third_party/drivestudio/models/losses.py
new file mode 100644
index 0000000..faaac5c
--- /dev/null
+++ b/worldbench/third_party/drivestudio/models/losses.py
@@ -0,0 +1,176 @@
+import numpy as np
+from typing import Literal, Union
+
+import torch
+import torch.nn.functional as F
+from torch import autograd, nn, Tensor
+
+def reduce(
+ loss: Union[torch.Tensor, np.ndarray],
+ mask: Union[torch.Tensor, np.ndarray] = None,
+ reduction: Literal['mean', 'mean_in_mask', 'sum', 'max', 'min', 'none']='mean'):
+
+ if mask is not None:
+ if mask.dim() == loss.dim() - 1:
+ mask = mask.view(*loss.shape[:-1], 1).expand_as(loss)
+ assert loss.dim() == mask.dim(), f"Expects loss.dim={loss.dim()} to be equal to mask.dim()={mask.dim()}"
+
+ if reduction == 'mean':
+ return loss.mean() if mask is None else (loss * mask).mean()
+ elif reduction == 'mean_in_mask':
+ return loss.mean() if mask is None else (loss * mask).sum() / mask.sum().clip(1e-5)
+ elif reduction == 'sum':
+ return loss.sum() if mask is None else (loss * mask).sum()
+ elif reduction == 'max':
+ return loss.max() if mask is None else loss[mask].max()
+ elif reduction == 'min':
+ return loss.min() if mask is None else loss[mask].min()
+ elif reduction == 'none':
+ return loss if mask is None else loss * mask
+ else:
+ raise RuntimeError(f"Invalid reduction={reduction}")
+
+class SafeBCE(autograd.Function):
+ """ Perform clipped BCE without disgarding gradients (preserve clipped gradients)
+ This function is equivalent to torch.clip(x, limit), 1-limit) before BCE,
+ BUT with grad existing on those clipped values.
+
+ NOTE: pytorch original BCELoss implementation is equivalent to limit = np.exp(-100) here.
+ see doc https://pytorch.org/docs/stable/generated/torch.nn.BCELoss.html
+ """
+ @staticmethod
+ def forward(ctx, x, y, limit):
+ assert (torch.where(y!=1, y+1, y)==1).all(), u'target must all be {0,1}'
+ ln_limit = ctx.ln_limit = np.log(limit)
+ # ctx.clip_grad = clip_grad
+
+ # NOTE: for example, torch.log(1-torch.tensor([1.000001])) = nan
+ x = torch.clip(x, 0, 1)
+ y = torch.clip(y, 0, 1)
+ ctx.save_for_backward(x, y)
+ return -torch.where(y==0, torch.log(1-x).clamp_min_(ln_limit), torch.log(x).clamp_min_(ln_limit))
+ # return -(y * torch.log(x).clamp_min_(ln_limit) + (1-y)*torch.log(1-x).clamp_min_(ln_limit))
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ x, y = ctx.saved_tensors
+ ln_limit = ctx.ln_limit
+
+ # NOTE: for y==0, do not clip small x; for y==1, do not clip small (1-x)
+ limit = np.exp(ln_limit)
+ # x = torch.clip(x, eclip, 1-eclip)
+ x = torch.where(y==0, torch.clip(x, 0, 1-limit), torch.clip(x, limit, 1))
+
+ grad_x = grad_y = None
+ if ctx.needs_input_grad[0]:
+ # ttt = torch.where(y==0, 1/(1-x), -1/x) * grad_output * (~(x==y))
+ # with open('grad.txt', 'a') as fp:
+ # fp.write(f"{ttt.min().item():.05f}, {ttt.max().item():.05f}\n")
+ # NOTE: " * (~(x==y))" so that those already match will not generate gradients.
+ grad_x = torch.where(y==0, 1/(1-x), -1/x) * grad_output * (~(x==y))
+ # grad_x = ( (1-y)/(1-x) - y/x ) * grad_output
+ if ctx.needs_input_grad[1]:
+ grad_y = (torch.log(1-x) - torch.log(x)) * grad_output * (~(x==y))
+ #---- x, y, limit
+ return grad_x, grad_y, None
+
+def safe_binary_cross_entropy(input: torch.Tensor, target: torch.Tensor, limit: float = 0.1, reduction="mean") -> torch.Tensor:
+ loss = SafeBCE.apply(input, target, limit)
+ return reduce(loss, None, reduction=reduction)
+
+def binary_cross_entropy(input: torch.Tensor, target: torch.Tensor, reduction="mean") -> torch.Tensor:
+ loss = F.binary_cross_entropy(input, target, reduction="none")
+ return reduce(loss, None, reduction=reduction)
+
+def normalize_depth(depth: Tensor, max_depth: float = 80.0):
+ return torch.clamp(depth / max_depth, 0.0, 1.0)
+
+def safe_normalize_depth(depth: Tensor, max_depth: float = 80.0):
+ return torch.clamp(depth / max_depth, 1e-06, 1.0)
+
+class DepthLoss(nn.Module):
+ def __init__(
+ self,
+ loss_type: Literal["l1", "l2", "smooth_l1"] = "l2",
+ normalize: bool = True,
+ use_inverse_depth: bool = False,
+ depth_error_percentile: float = None,
+ upper_bound: float = 80,
+ reduction: Literal["mean_on_hit", "mean_on_hw", "sum", "none"] = "mean_on_hit",
+ ):
+ super().__init__()
+ self.loss_type = loss_type
+ self.normalize = normalize
+ self.use_inverse_depth = use_inverse_depth
+ self.upper_bound = upper_bound
+ self.depth_error_percentile = depth_error_percentile
+ self.reduction = reduction
+
+ def _compute_depth_loss(
+ self,
+ pred_depth: Tensor,
+ gt_depth: Tensor,
+ max_depth: float = 80,
+ hit_mask: Tensor = None,
+ ):
+ pred_depth = pred_depth.squeeze()
+ gt_depth = gt_depth.squeeze()
+ if hit_mask is not None:
+ pred_depth = pred_depth * hit_mask
+ gt_depth = gt_depth * hit_mask
+
+ # cal valid mask to make sure gt_depth is valid
+ valid_mask = (gt_depth > 0.01) & (gt_depth < max_depth) & (pred_depth > 0.0001)
+
+ # normalize depth to (0, 1)
+ if self.normalize:
+ pred_depth = safe_normalize_depth(pred_depth[valid_mask], max_depth=max_depth)
+ gt_depth = safe_normalize_depth(gt_depth[valid_mask], max_depth=max_depth)
+ else:
+ pred_depth = pred_depth[valid_mask]
+ gt_depth = gt_depth[valid_mask]
+
+ # inverse the depth map (0, 1) -> (1, +inf)
+ if self.use_inverse_depth:
+ pred_depth = 1./pred_depth
+ gt_depth = 1./gt_depth
+
+ # cal loss
+ if self.loss_type == "smooth_l1":
+ return F.smooth_l1_loss(pred_depth, gt_depth, reduction="none")
+ elif self.loss_type == "l1":
+ return F.l1_loss(pred_depth, gt_depth, reduction="none")
+ elif self.loss_type == "l2":
+ return F.mse_loss(pred_depth, gt_depth, reduction="none")
+ else:
+ raise NotImplementedError(f"Unknown loss type: {self.loss_type}")
+
+ def __call__(
+ self,
+ pred_depth: Tensor,
+ gt_depth: Tensor,
+ hit_mask: Tensor = None,
+ ):
+ depth_error = self._compute_depth_loss(pred_depth, gt_depth, self.upper_bound, hit_mask)
+ if self.depth_error_percentile is not None:
+ # to avoid outliers. not used for now
+ depth_error = depth_error.flatten()
+ depth_error = depth_error[
+ depth_error.argsort()[
+ : int(len(depth_error) * self.depth_error_percentile)
+ ]
+ ]
+
+ if self.reduction == "sum":
+ depth_error = depth_error.sum()
+ elif self.reduction == "none":
+ depth_error = depth_error
+ elif self.reduction == "mean_on_hit":
+ depth_error = depth_error.mean()
+ elif self.reduction == "mean_on_hw":
+ n = gt_depth.shape[0]*gt_depth.shape[1]
+ depth_error = depth_error.sum() / n
+ else:
+ raise NotImplementedError(f"Unknown reduction method: {self.reduction}")
+
+ return depth_error
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/models/modules.py b/worldbench/third_party/drivestudio/models/modules.py
new file mode 100644
index 0000000..45cb35a
--- /dev/null
+++ b/worldbench/third_party/drivestudio/models/modules.py
@@ -0,0 +1,670 @@
+import torch
+from typing import Optional, Tuple
+import logging
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from pytorch3d.ops import knn_points
+import nvdiffrast.torch as dr
+from utils.geometry import rotation_6d_to_matrix
+
+logger = logging.getLogger()
+
+class XYZ_Encoder(nn.Module):
+ encoder_type = "XYZ_Encoder"
+ """Encode XYZ coordinates or directions to a vector."""
+
+ def __init__(self, n_input_dims):
+ super().__init__()
+ self.n_input_dims = n_input_dims
+
+ @property
+ def n_output_dims(self) -> int:
+ raise NotImplementedError
+
+class SinusoidalEncoder(XYZ_Encoder):
+ encoder_type = "SinusoidalEncoder"
+ """Sinusoidal Positional Encoder used in Nerf."""
+
+ def __init__(
+ self,
+ n_input_dims: int = 3,
+ min_deg: int = 0,
+ max_deg: int = 10,
+ enable_identity: bool = True,
+ ):
+ super().__init__(n_input_dims)
+ self.n_input_dims = n_input_dims
+ self.min_deg = min_deg
+ self.max_deg = max_deg
+ self.enable_identity = enable_identity
+ self.register_buffer(
+ "scales", Tensor([2**i for i in range(min_deg, max_deg + 1)])
+ )
+
+ @property
+ def n_output_dims(self) -> int:
+ return (
+ int(self.enable_identity) + (self.max_deg - self.min_deg + 1) * 2
+ ) * self.n_input_dims
+
+ @torch.no_grad()
+ def forward(self, x: Tensor) -> Tensor:
+ """
+ Args:
+ x: [..., n_input_dims]
+ Returns:
+ encoded: [..., n_output_dims]
+ """
+ if self.max_deg == self.min_deg:
+ return x
+ xb = torch.reshape(
+ (x[..., None, :] * self.scales[:, None]),
+ list(x.shape[:-1])
+ + [(self.max_deg - self.min_deg + 1) * self.n_input_dims],
+ )
+ encoded = torch.sin(torch.cat([xb, xb + 0.5 * torch.pi], dim=-1))
+ if self.enable_identity:
+ encoded = torch.cat([x] + [encoded], dim=-1)
+ return encoded
+
+class MLP(nn.Module):
+ """A simple MLP with skip connections."""
+
+ def __init__(
+ self,
+ in_dims: int,
+ out_dims: int,
+ num_layers: int = 3,
+ hidden_dims: Optional[int] = 256,
+ skip_connections: Optional[Tuple[int]] = [0],
+ ) -> None:
+ super().__init__()
+ self.in_dims = in_dims
+ self.hidden_dims = hidden_dims
+ self.n_output_dims = out_dims
+ self.num_layers = num_layers
+ self.skip_connections = skip_connections
+ layers = []
+ if self.num_layers == 1:
+ layers.append(nn.Linear(in_dims, out_dims))
+ else:
+ for i in range(self.num_layers - 1):
+ if i == 0:
+ layers.append(nn.Linear(in_dims, hidden_dims))
+ elif i in skip_connections:
+ layers.append(nn.Linear(in_dims + hidden_dims, hidden_dims))
+ else:
+ layers.append(nn.Linear(hidden_dims, hidden_dims))
+ layers.append(nn.Linear(hidden_dims, out_dims))
+ self.layers = nn.ModuleList(layers)
+
+ def forward(self, x: Tensor) -> Tensor:
+ input = x
+ for i, layer in enumerate(self.layers):
+ if i in self.skip_connections:
+ x = torch.cat([x, input], -1)
+ x = layer(x)
+ if i < len(self.layers) - 1:
+ x = nn.functional.relu(x)
+ return x
+
+class SkyModel(nn.Module):
+ def __init__(
+ self,
+ class_name: str,
+ n: int,
+ head_mlp_layer_width: int = 64,
+ enable_appearance_embedding: bool = True,
+ appearance_embedding_dim: int = 16,
+ device: torch.device = torch.device("cuda")
+ ):
+ super().__init__()
+ self.class_prefix = class_name + "#"
+ self.device = device
+ self.direction_encoding = SinusoidalEncoder(
+ n_input_dims=3, min_deg=0, max_deg=6
+ )
+ self.direction_encoding.requires_grad_(False)
+
+ self.enable_appearance_embedding = enable_appearance_embedding
+ if self.enable_appearance_embedding:
+ self.appearance_embedding_dim = appearance_embedding_dim
+ self.appearance_embedding = nn.Embedding(n, appearance_embedding_dim, dtype=torch.float32)
+
+ in_dims = self.direction_encoding.n_output_dims + appearance_embedding_dim \
+ if self.enable_appearance_embedding else self.direction_encoding.n_output_dims
+ self.sky_head = MLP(
+ in_dims=in_dims,
+ out_dims=3,
+ num_layers=3,
+ hidden_dims=head_mlp_layer_width,
+ skip_connections=[1],
+ )
+ self.in_test_set = False
+
+ def forward(self, image_infos):
+ directions = image_infos["viewdirs"]
+ self.device = directions.device
+ prefix = directions.shape[:-1]
+
+ dd = self.direction_encoding(directions.reshape(-1, 3)).to(self.device)
+ if self.enable_appearance_embedding:
+ # optionally add appearance embedding
+ if "img_idx" in image_infos and not self.in_test_set:
+ appearance_embedding = self.appearance_embedding(image_infos["img_idx"]).reshape(-1, self.appearance_embedding_dim)
+ else:
+ # use mean appearance embedding
+ appearance_embedding = torch.ones(
+ (*dd.shape[:-1], self.appearance_embedding_dim),
+ device=dd.device,
+ ) * self.appearance_embedding.weight.mean(dim=0)
+ dd = torch.cat([dd, appearance_embedding], dim=-1)
+ rgb_sky = self.sky_head(dd).to(self.device)
+ rgb_sky = F.sigmoid(rgb_sky)
+ return rgb_sky.reshape(prefix + (3,))
+
+ def get_param_groups(self):
+ return {
+ self.class_prefix+"all": self.parameters(),
+ }
+
+class EnvLight(torch.nn.Module):
+
+ def __init__(
+ self,
+ class_name: str,
+ resolution=1024,
+ device: torch.device = torch.device("cuda"),
+ **kwargs
+ ):
+ super().__init__()
+ self.class_prefix = class_name + "#"
+ self.device = device
+ self.to_opengl = torch.tensor([[1, 0, 0], [0, 0, 1], [0, -1, 0]], dtype=torch.float32, device="cuda")
+ self.base = torch.nn.Parameter(
+ 0.5 * torch.ones(6, resolution, resolution, 3, requires_grad=True),
+ )
+
+ def forward(self, image_infos):
+ l = image_infos["viewdirs"]
+
+ l = (l.reshape(-1, 3) @ self.to_opengl.T).reshape(*l.shape)
+ l = l.contiguous()
+ prefix = l.shape[:-1]
+ if len(prefix) != 3: # reshape to [B, H, W, -1]
+ l = l.reshape(1, 1, -1, l.shape[-1])
+
+ light = dr.texture(self.base[None, ...], l, filter_mode='linear', boundary_mode='cube')
+ light = light.view(*prefix, -1)
+
+ return light
+
+ def get_param_groups(self):
+ return {
+ self.class_prefix+"all": self.parameters(),
+ }
+
+class AffineTransform(nn.Module):
+ def __init__(
+ self,
+ class_name: str,
+ n: int,
+ embedding_dim: int = 4,
+ pixel_affine: bool = False,
+ base_mlp_layer_width: int = 64,
+ device: torch.device = torch.device("cuda")
+ ):
+ super().__init__()
+ self.class_prefix = class_name + "#"
+ self.device = device
+ self.embedding_dim = embedding_dim
+ self.pixel_affine = pixel_affine
+ self.embedding = nn.Embedding(n, embedding_dim, dtype=torch.float32)
+
+ input_dim = (embedding_dim + 2)if self.pixel_affine else embedding_dim
+ self.decoder = nn.Sequential(
+ nn.Linear(input_dim, base_mlp_layer_width),
+ nn.ReLU(),
+ nn.Linear(base_mlp_layer_width, 12),
+ )
+ self.in_test_set = False
+
+ self.zero_init()
+
+ def zero_init(self):
+ torch.nn.init.zeros_(self.embedding.weight)
+ for layer in self.decoder:
+ if isinstance(layer, nn.Linear):
+ torch.nn.init.zeros_(layer.weight)
+ torch.nn.init.zeros_(layer.bias)
+
+ def forward(self, image_infos):
+ if "img_idx" in image_infos and not self.in_test_set:
+ embedding = self.embedding(image_infos["img_idx"])
+ else:
+ # use mean appearance embedding
+ embedding = torch.ones(
+ (*image_infos["viewdirs"].shape[:-1], self.embedding_dim),
+ device=image_infos["viewdirs"].device,
+ ) * self.embedding.weight.mean(dim=0)
+ if self.pixel_affine:
+ embedding = torch.cat([embedding, image_infos["pixel_coords"]], dim=-1)
+ affine = self.decoder(embedding)
+ affine = affine.reshape(*embedding.shape[:-1], 3, 4)
+
+ affine[..., :3, :3] = affine[..., :3, :3] + torch.eye(3, device=affine.device).reshape(1, 3, 3)
+ return affine
+
+ def get_param_groups(self):
+ return {
+ self.class_prefix+"all": self.parameters(),
+ }
+
+class CameraOptModule(torch.nn.Module):
+ """Camera pose optimization module."""
+
+ def __init__(
+ self,
+ class_name: str,
+ n: int,
+ device: torch.device = torch.device("cuda"),
+ ):
+ super().__init__()
+ self.class_prefix = class_name + "#"
+ self.device = device
+ # Delta positions (3D) + Delta rotations (6D)
+ self.embeds = torch.nn.Embedding(n, 9)
+ # Identity rotation in 6D representation
+ self.register_buffer("identity", torch.tensor([1.0, 0.0, 0.0, 0.0, 1.0, 0.0]))
+
+ self.zero_init() # important for initialization !!
+
+ def zero_init(self):
+ torch.nn.init.zeros_(self.embeds.weight)
+
+ def random_init(self, std: float):
+ torch.nn.init.normal_(self.embeds.weight, std=std)
+
+ def forward(self, camtoworlds: Tensor, embed_ids: Tensor) -> Tensor:
+ """Adjust camera pose based on deltas.
+
+ Args:
+ camtoworlds: (..., 4, 4)
+ embed_ids: (...,)
+
+ Returns:
+ updated camtoworlds: (..., 4, 4)
+ """
+ assert camtoworlds.shape[:-2] == embed_ids.shape
+ batch_shape = camtoworlds.shape[:-2]
+ pose_deltas = self.embeds(embed_ids) # (..., 9)
+ dx, drot = pose_deltas[..., :3], pose_deltas[..., 3:]
+ rot = rotation_6d_to_matrix(
+ drot + self.identity.expand(*batch_shape, -1)
+ ) # (..., 3, 3)
+ transform = torch.eye(4, device=pose_deltas.device).repeat((*batch_shape, 1, 1))
+ transform[..., :3, :3] = rot
+ transform[..., :3, 3] = dx
+ return torch.matmul(camtoworlds, transform)
+
+ def get_param_groups(self):
+ return {
+ self.class_prefix+"all": self.parameters(),
+ }
+
+def get_embedder(multires, i=1):
+ if i == -1:
+ return nn.Identity(), 3
+
+ embed_kwargs = {
+ 'include_input': True,
+ 'input_dims': i,
+ 'max_freq_log2': multires - 1,
+ 'num_freqs': multires,
+ 'log_sampling': True,
+ 'periodic_fns': [torch.sin, torch.cos],
+ }
+
+ embedder_obj = Embedder(**embed_kwargs)
+ embed = lambda x, eo=embedder_obj: eo.embed(x)
+ return embed, embedder_obj.out_dim
+
+
+class Embedder:
+ def __init__(self, **kwargs):
+ self.kwargs = kwargs
+ self.create_embedding_fn()
+
+ def create_embedding_fn(self):
+ embed_fns = []
+ d = self.kwargs['input_dims']
+ out_dim = 0
+ if self.kwargs['include_input']:
+ embed_fns.append(lambda x: x)
+ out_dim += d
+
+ max_freq = self.kwargs['max_freq_log2']
+ N_freqs = self.kwargs['num_freqs']
+
+ if self.kwargs['log_sampling']:
+ freq_bands = 2. ** torch.linspace(0., max_freq, steps=N_freqs)
+ else:
+ freq_bands = torch.linspace(2. ** 0., 2. ** max_freq, steps=N_freqs)
+
+ for freq in freq_bands:
+ for p_fn in self.kwargs['periodic_fns']:
+ embed_fns.append(lambda x, p_fn=p_fn, freq=freq: p_fn(x * freq))
+ out_dim += d
+
+ self.embed_fns = embed_fns
+ self.out_dim = out_dim
+
+ def embed(self, inputs):
+ return torch.cat([fn(inputs) for fn in self.embed_fns], -1)
+
+
+class DeformNetwork(nn.Module):
+ def __init__(self, D=8, W=256, input_ch=3, output_ch=59, x_multires=10, t_multires=10):
+ super(DeformNetwork, self).__init__()
+ self.D = D
+ self.W = W
+ self.input_ch = input_ch
+ self.output_ch = output_ch
+ self.x_multires = x_multires
+ self.t_multires = t_multires
+ self.skips = [D // 2]
+
+ self.embed_time_fn, time_input_ch = get_embedder(self.t_multires, 1)
+ self.embed_fn, xyz_input_ch = get_embedder(self.x_multires, 3)
+ self.input_ch = xyz_input_ch + time_input_ch
+
+ self.linear = nn.ModuleList(
+ [nn.Linear(self.input_ch, W)] + [
+ nn.Linear(W, W) if i not in self.skips else nn.Linear(W + self.input_ch, W)
+ for i in range(D - 1)]
+ )
+
+ self.gaussian_warp = nn.Linear(W, 3)
+ self.gaussian_rotation = nn.Linear(W, 4)
+ self.gaussian_scaling = nn.Linear(W, 3)
+
+ def forward(self, x, t):
+ t_emb = self.embed_time_fn(t)
+ x_emb = self.embed_fn(x)
+ h = torch.cat([x_emb, t_emb], dim=-1)
+ for i, l in enumerate(self.linear):
+ h = self.linear[i](h)
+ h = F.relu(h)
+ if i in self.skips:
+ h = torch.cat([x_emb, t_emb, h], -1)
+
+ d_xyz = self.gaussian_warp(h)
+ scaling = self.gaussian_scaling(h)
+ rotation = self.gaussian_rotation(h)
+
+ return d_xyz, rotation, scaling
+
+
+class ConditionalDeformNetwork(nn.Module):
+ def __init__(self, D=8, W=256, input_ch=3, embed_dim=10,
+ x_multires=10, t_multires=10,
+ deform_quat=True, deform_scale=True):
+ super(ConditionalDeformNetwork, self).__init__()
+ self.D = D
+ self.W = W
+ self.input_ch = input_ch
+ self.embed_dim = embed_dim
+ self.deform_quat = deform_quat
+ self.deform_scale = deform_scale
+ self.skips = [D // 2]
+
+ self.embed_time_fn, time_input_ch = get_embedder(t_multires, 1)
+ self.embed_fn, xyz_input_ch = get_embedder(x_multires, 3)
+ self.input_ch = xyz_input_ch + time_input_ch + embed_dim
+
+ self.linear = nn.ModuleList(
+ [nn.Linear(self.input_ch, W)] + [
+ nn.Linear(W, W) if i not in self.skips else nn.Linear(W + self.input_ch, W)
+ for i in range(D - 1)]
+ )
+
+ self.gaussian_warp = nn.Linear(W, 3)
+ if self.deform_quat:
+ self.gaussian_rotation = nn.Linear(W, 4)
+ if self.deform_scale:
+ self.gaussian_scaling = nn.Linear(W, 3)
+
+ def forward(self, x, t, condition):
+ t_emb = self.embed_time_fn(t)
+ x_emb = self.embed_fn(x)
+ h = torch.cat([x_emb, t_emb, condition], dim=-1)
+ for i, l in enumerate(self.linear):
+ h = self.linear[i](h)
+ h = F.relu(h)
+ if i in self.skips:
+ h = torch.cat([x_emb, t_emb, condition, h], -1)
+
+ d_xyz = self.gaussian_warp(h)
+ scaling, rotation = None, None
+ if self.deform_scale:
+ scaling = self.gaussian_scaling(h)
+ if self.deform_quat:
+ rotation = self.gaussian_rotation(h)
+
+ return d_xyz, rotation, scaling
+
+class VoxelDeformer(nn.Module):
+ def __init__(
+ self,
+ vtx,
+ vtx_features,
+ resolution_dhw=[8, 32, 32],
+ short_dim_dhw=0, # 0 is d, corresponding to z
+ long_dim_dhw=1,
+ is_resume=False
+ ) -> None:
+ super().__init__()
+ # vtx B,N,3, vtx_features: B,N,J
+ # d-z h-y w-x; human is facing z; dog is facing x, z is upward, should compress on y
+ B = vtx.shape[0]
+ assert vtx.shape[0] == vtx_features.shape[0], "Batch size mismatch"
+
+ # * Prepare Grid
+ self.resolution_dhw = resolution_dhw
+ device = vtx.device
+ d, h, w = self.resolution_dhw
+
+ self.register_buffer(
+ "ratio",
+ torch.Tensor(
+ [self.resolution_dhw[long_dim_dhw] / self.resolution_dhw[short_dim_dhw]]
+ ).squeeze(),
+ )
+ self.ratio_dim = -1 - short_dim_dhw
+ x_range = (
+ (torch.linspace(-1, 1, steps=w, device=device))
+ .view(1, 1, 1, w)
+ .expand(1, d, h, w)
+ )
+ y_range = (
+ (torch.linspace(-1, 1, steps=h, device=device))
+ .view(1, 1, h, 1)
+ .expand(1, d, h, w)
+ )
+ z_range = (
+ (torch.linspace(-1, 1, steps=d, device=device))
+ .view(1, d, 1, 1)
+ .expand(1, d, h, w)
+ )
+ grid = (
+ torch.cat((x_range, y_range, z_range), dim=0)
+ .reshape(1, 3, -1)
+ .permute(0, 2, 1)
+ )
+ grid = grid.expand(B, -1, -1)
+
+ gt_bbox_min = (vtx.min(dim=1).values).to(device)
+ gt_bbox_max = (vtx.max(dim=1).values).to(device)
+ offset = (gt_bbox_min + gt_bbox_max) * 0.5
+ self.register_buffer(
+ "global_scale", torch.Tensor([1.2]).squeeze()
+ ) # from Fast-SNARF
+ scale = (
+ (gt_bbox_max - gt_bbox_min).max(dim=-1).values / 2 * self.global_scale
+ ).unsqueeze(-1)
+
+ corner = torch.ones_like(offset) * scale
+ corner[:, self.ratio_dim] /= self.ratio
+ min_vert = (offset - corner).reshape(-1, 1, 3)
+ max_vert = (offset + corner).reshape(-1, 1, 3)
+ self.bbox = torch.cat([min_vert, max_vert], dim=1)
+
+ self.register_buffer("scale", scale.unsqueeze(1)) # [B, 1, 1]
+ self.register_buffer("offset", offset.unsqueeze(1)) # [B, 1, 3]
+
+ grid_denorm = self.denormalize(
+ grid
+ ) # grid_denorm is in the same scale as the canonical body
+
+ if not is_resume:
+ weights = (
+ self._query_weights_smpl(
+ grid_denorm,
+ smpl_verts=vtx.detach().clone(),
+ smpl_weights=vtx_features.detach().clone(),
+ )
+ .detach()
+ .clone()
+ )
+ else:
+ # random initialization
+ weights = torch.randn(
+ B, vtx_features.shape[-1], *resolution_dhw
+ ).to(device)
+
+ self.register_buffer("lbs_voxel_base", weights.detach())
+ self.register_buffer("grid_denorm", grid_denorm)
+
+ self.num_bones = vtx_features.shape[-1]
+
+ # # debug
+ # import numpy as np
+ # np.savetxt("./debug/dbg.xyz", grid_denorm[0].detach().cpu())
+ # np.savetxt("./debug/vtx.xyz", vtx[0].detach().cpu())
+ return
+
+ def enable_voxel_correction(self):
+ voxel_w_correction = torch.zeros_like(self.lbs_voxel_base)
+ self.voxel_w_correction = nn.Parameter(voxel_w_correction)
+
+ def enable_additional_correction(self, additional_channels, std=1e-4):
+ additional_correction = (
+ torch.ones(
+ self.lbs_voxel_base.shape[0],
+ additional_channels,
+ *self.lbs_voxel_base.shape[2:]
+ )
+ * std
+ )
+ self.additional_correction = nn.Parameter(additional_correction)
+
+ @property
+ def get_voxel_weight(self):
+ w = self.lbs_voxel_base
+ if hasattr(self, "voxel_w_correction"):
+ w = w + self.voxel_w_correction
+ if hasattr(self, "additional_correction"):
+ w = torch.cat([w, self.additional_correction], dim=1)
+ return w
+
+ def get_tv(self, name="dc"):
+ if name == "dc":
+ if not hasattr(self, "voxel_w_correction"):
+ return torch.zeros(1).squeeze().to(self.lbs_voxel_base.device)
+ d = self.voxel_w_correction
+ elif name == "rest":
+ if not hasattr(self, "additional_correction"):
+ return torch.zeros(1).squeeze().to(self.lbs_voxel_base.device)
+ d = self.additional_correction
+ tv_x = torch.abs(d[:, :, 1:, :, :] - d[:, :, :-1, :, :]).mean()
+ tv_y = torch.abs(d[:, :, :, 1:, :] - d[:, :, :, :-1, :]).mean()
+ tv_z = torch.abs(d[:, :, :, :, 1:] - d[:, :, :, :, :-1]).mean()
+ return (tv_x + tv_y + tv_z) / 3.0
+ # tv_x = torch.abs(d[:, :, 1:, :, :] - d[:, :, :-1, :, :]).sum()
+ # tv_y = torch.abs(d[:, :, :, 1:, :] - d[:, :, :, :-1, :]).sum()
+ # tv_z = torch.abs(d[:, :, :, :, 1:] - d[:, :, :, :, :-1]).sum()
+ # return tv_x + tv_y + tv_z
+
+ def get_mag(self, name="dc"):
+ if name == "dc":
+ if not hasattr(self, "voxel_w_correction"):
+ return torch.zeros(1).squeeze().to(self.lbs_voxel_base.device)
+ d = self.voxel_w_correction
+ elif name == "rest":
+ if not hasattr(self, "additional_correction"):
+ return torch.zeros(1).squeeze().to(self.lbs_voxel_base.device)
+ d = self.additional_correction
+ return torch.norm(d, dim=1).mean()
+
+ def forward(self, xc, mode="bilinear"):
+ shape = xc.shape # ..., 3
+ # xc = xc.reshape(1, -1, 3)
+ w = F.grid_sample(
+ self.get_voxel_weight,
+ self.normalize(xc)[:, :, None, None],
+ align_corners=True,
+ mode=mode,
+ padding_mode="border",
+ )
+ w = w.squeeze(3, 4).permute(0, 2, 1)
+ w = w.reshape(*shape[:-1], -1)
+ # * the w may have more channels
+ return w
+
+ def normalize(self, x):
+ x_normalized = x.clone()
+ x_normalized -= self.offset
+ x_normalized /= self.scale
+ x_normalized[..., self.ratio_dim] *= self.ratio
+ return x_normalized
+
+ def denormalize(self, x):
+ x_denormalized = x.clone()
+ x_denormalized[..., self.ratio_dim] /= self.ratio
+ x_denormalized *= self.scale
+ x_denormalized += self.offset
+ return x_denormalized
+
+ def _query_weights_smpl(self, x, smpl_verts, smpl_weights):
+ # adapted from https://github.com/jby1993/SelfReconCode/blob/main/model/Deformer.py
+ dist, idx, _ = knn_points(x, smpl_verts.detach(), K=30) # [B, N, 30]
+ dist = dist.sqrt().clamp_(0.0001, 1.0)
+ expanded_smpl_weights = smpl_weights.unsqueeze(2).expand(-1, -1, idx.shape[2], -1) # [B, N, 30, J]
+ weights = expanded_smpl_weights.gather(1, idx.unsqueeze(-1).expand(-1, -1, -1, expanded_smpl_weights.shape[-1])) # [B, N, 30, J]
+
+ ws = 1.0 / dist
+ ws = ws / ws.sum(-1, keepdim=True)
+ weights = (ws[..., None] * weights).sum(-2)
+
+ b = x.shape[0]
+ c = smpl_weights.shape[-1]
+ d, h, w = self.resolution_dhw
+ weights = weights.permute(0, 2, 1).reshape(b, c, d, h, w)
+ for _ in range(30):
+ mean = (
+ weights[:, :, 2:, 1:-1, 1:-1]
+ + weights[:, :, :-2, 1:-1, 1:-1]
+ + weights[:, :, 1:-1, 2:, 1:-1]
+ + weights[:, :, 1:-1, :-2, 1:-1]
+ + weights[:, :, 1:-1, 1:-1, 2:]
+ + weights[:, :, 1:-1, 1:-1, :-2]
+ ) / 6.0
+ weights[:, :, 1:-1, 1:-1, 1:-1] = (
+ weights[:, :, 1:-1, 1:-1, 1:-1] - mean
+ ) * 0.7 + mean
+ sums = weights.sum(1, keepdim=True)
+ weights = weights / sums
+ return weights.detach()
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/models/nodes/__init__.py b/worldbench/third_party/drivestudio/models/nodes/__init__.py
new file mode 100644
index 0000000..440c6e8
--- /dev/null
+++ b/worldbench/third_party/drivestudio/models/nodes/__init__.py
@@ -0,0 +1,3 @@
+from .rigid import RigidNodes
+from .deformable import DeformableNodes
+from .smpl import SMPLNodes
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/models/nodes/deformable.py b/worldbench/third_party/drivestudio/models/nodes/deformable.py
new file mode 100644
index 0000000..775387f
--- /dev/null
+++ b/worldbench/third_party/drivestudio/models/nodes/deformable.py
@@ -0,0 +1,211 @@
+from typing import Dict, List, Tuple
+import logging
+
+import torch
+from torch.nn import Parameter
+
+from models.modules import ConditionalDeformNetwork
+from models.gaussians.basics import *
+from models.nodes.rigid import RigidNodes
+
+logger = logging.getLogger()
+
+class DeformableNodes(RigidNodes):
+ def __init__(
+ self,
+ **kwargs
+ ):
+ super().__init__(**kwargs)
+ self.instances_embedding = torch.zeros(1, self.networks_cfg.embed_dim, device=self.device)
+ self.deform_network = ConditionalDeformNetwork(
+ input_ch=3, **self.networks_cfg
+ ).to(self.device)
+
+ def create_from_pcd(self, instance_pts_dict: Dict[str, torch.Tensor]) -> None:
+ super().create_from_pcd(instance_pts_dict=instance_pts_dict)
+ init_embedding = torch.rand(self.num_instances, self.networks_cfg.embed_dim, device=self.device)
+ self.instances_embedding = Parameter(init_embedding) # overrided the previous one
+
+ def get_param_groups(self) -> Dict[str, List[Parameter]]:
+ param_groups = self.get_gaussian_param_groups()
+ param_groups[self.class_prefix+"embedding"] = [self.instances_embedding]
+ param_groups[self.class_prefix+"deform_network"] = list(self.deform_network.parameters())
+ return param_groups
+
+ def get_deformation(self, local_means) -> Tuple:
+ """
+ get the deformation of the nonrigid instances
+ """
+ assert local_means.shape[0] == self.point_ids.shape[0], \
+ "its a bug here, we need to pass the mask for points_ids"
+ nonrigid_embed = self.instances_embedding[self.point_ids[..., 0]]
+ ins_height = self.instances_size[self.point_ids[..., 0]][..., 2]
+ x = local_means.data / ins_height[:, None] * 2
+ t = self.normalized_timestamps[self.cur_frame]
+ t = t.unsqueeze(0).repeat(self.point_ids.shape[0], 1)
+ delta_xyz, delta_quat, delta_scale = self.deform_network(x, t, nonrigid_embed)
+ return delta_xyz, delta_quat, delta_scale
+
+ def get_gaussians(self, cam: dataclass_camera) -> Dict[str, torch.Tensor]:
+ filter_mask = torch.ones_like(self._means[:, 0], dtype=torch.bool)
+ self.filter_mask = filter_mask
+
+ delta_xyz, delta_quat, delta_scale = None, None, None
+ if self.ctrl_cfg.use_deformgs_for_nonrigid and self.step > self.ctrl_cfg.use_deformgs_after:
+ delta_xyz, delta_quat, delta_scale = self.get_deformation(local_means=self._means)
+
+ if delta_xyz is not None:
+ if self.ctrl_cfg.stop_optimizing_canonical_xyz:
+ means = self._means.data + delta_xyz
+ else:
+ means = self._means + delta_xyz
+ world_means = self.transform_means(means)
+ else:
+ world_means = self.transform_means(self._means)
+
+ if delta_quat is not None:
+ quats = self.get_quats + delta_quat
+ world_quats = self.transform_quats(quats)
+ else:
+ world_quats = self.transform_quats(self._quats)
+
+ if delta_scale is not None:
+ activated_scales = self.get_scaling + delta_scale
+ else:
+ activated_scales = self.get_scaling
+
+ # get colors of gaussians
+ colors = torch.cat((self._features_dc[:, None, :], self._features_rest), dim=1)
+ if self.sh_degree > 0:
+ viewdirs = world_means.detach() - cam.camtoworlds.data[..., :3, 3] # (N, 3)
+ viewdirs = viewdirs / viewdirs.norm(dim=-1, keepdim=True)
+ n = min(self.step // self.ctrl_cfg.sh_degree_interval, self.sh_degree)
+ rgbs = spherical_harmonics(n, viewdirs, colors)
+ rgbs = torch.clamp(rgbs + 0.5, 0.0, 1.0)
+ else:
+ rgbs = torch.sigmoid(colors[:, 0, :])
+
+ valid_mask = self.get_pts_valid_mask()
+
+ activated_opacities = self.get_opacity * valid_mask.float().unsqueeze(-1)
+ activated_rotations = self.quat_act(world_quats)
+ actovated_colors = rgbs
+
+ # collect gaussians information
+ gs_dict = dict(
+ _means=world_means[filter_mask],
+ _opacities=activated_opacities[filter_mask],
+ _rgbs=actovated_colors[filter_mask],
+ _scales=activated_scales[filter_mask],
+ _quats=activated_rotations[filter_mask],
+ )
+
+ # check nan in gs_dict
+ for k, v in gs_dict.items():
+ if torch.isnan(v).any():
+ raise ValueError(f"NaN detected in gaussian {k} at step {self.step}")
+ if torch.isinf(v).any():
+ raise ValueError(f"Inf detected in gaussian {k} at step {self.step}")
+
+ self._gs_cache = {
+ "_scales": activated_scales[filter_mask],
+ "local_xyz_deformed": means[filter_mask] if delta_xyz is not None else None,
+ }
+ return gs_dict
+
+ def compute_reg_loss(self):
+ loss_dict = super().compute_reg_loss()
+ out_of_bound_losscfg = self.reg_cfg.get("out_of_bound_loss", None)
+ if out_of_bound_losscfg is not None:
+ w = out_of_bound_losscfg.w
+ local_xyz_deformed = self._gs_cache["local_xyz_deformed"]
+ if w > 0 and local_xyz_deformed is not None:
+ local_xyz_deformed = self._gs_cache["local_xyz_deformed"]
+ per_pts_size = self.instances_size[self.point_ids[..., 0]]
+ loss_dict["out_of_bound_loss"] = torch.relu(local_xyz_deformed.abs() - per_pts_size / 2).mean() * w
+ return loss_dict
+
+ def deform_gaussian_points(
+ self, gaussian_dict: Dict[str, torch.Tensor], cur_normalized_time: float,
+ ) -> Dict[str, torch.Tensor]:
+ """
+ deform the points
+ """
+ means = gaussian_dict["means"]
+ nonrigid_embed = self.instances_embedding[gaussian_dict["ids"].squeeze()]
+ cur_normalized_time = torch.tensor(cur_normalized_time, dtype=torch.float32, device=self.device).unsqueeze(0).repeat(means.shape[0], 1)
+ delta_xyz, delta_quat, delta_scale = self.deform_network(means, cur_normalized_time, nonrigid_embed)
+ gaussian_dict["means"] = means + delta_xyz
+ if delta_scale is not None:
+ gaussian_dict["scales"] = gaussian_dict["scales"] + delta_scale
+ if delta_quat is not None:
+ gaussian_dict["quats"] = gaussian_dict["quats"] + delta_quat
+ return gaussian_dict
+
+ def load_state_dict(self, dict: Dict, **kwargs) -> str:
+ instances_num = dict["instances_fv"].shape[1]
+ self.instances_embedding = Parameter(
+ torch.rand([instances_num, self.networks_cfg.embed_dim], device=self.device)
+ )
+ # NOTE: keep it, maybe used in the future
+ # del self.instances_embedding
+ # self.instances_embedding = nn.Embedding(
+ # instances_num, self.networks_cfg.embed_dim
+ # ).to(self.device)
+ msg = super().load_state_dict(dict, **kwargs)
+ return msg
+
+ def collect_gaussians_from_ids(self, ids: List[int]) -> Dict:
+ gaussian_dict = super().collect_gaussians_from_ids(ids)
+ # collect embeddings
+ for id in ids:
+ instance_embedding = self.instances_embedding[id]
+ gaussian_dict[id]["embedding"] = instance_embedding
+ return gaussian_dict
+
+ def replace_instances(self, replace_dict: Dict[int, int]) -> None:
+ """
+ replace instances from the model
+
+ Args:
+ replace_dict: {
+ ins_id(to be replaced): ins_id(replace with)
+ ...
+ }
+ """
+ new_gaussians_dict = self.collect_gaussians_from_ids(replace_dict.values())
+ for ins_id, new_id in replace_dict.items():
+ self.remove_instances([ins_id])
+ new_gaussian = new_gaussians_dict[new_id]
+ self._means = Parameter(torch.cat([self._means, new_gaussian["_means"]], dim=0))
+ self._scales = Parameter(torch.cat([self._scales, new_gaussian["_scales"]], dim=0))
+ self._quats = Parameter(torch.cat([self._quats, new_gaussian["_quats"]], dim=0))
+ self._features_dc = Parameter(torch.cat([self._features_dc, new_gaussian["_features_dc"]], dim=0))
+ self._features_rest = Parameter(torch.cat([self._features_rest, new_gaussian["_features_rest"]], dim=0))
+ self._opacities = Parameter(torch.cat([self._opacities, new_gaussian["_opacities"]], dim=0))
+ # keeps original point ids
+ self.point_ids = torch.cat([self.point_ids, torch.full_like(new_gaussian["point_ids"], ins_id)], dim=0)
+ # replace embeddings
+ # NOTE: modify data in nn.Parameter directly
+ self.instances_embedding.data[ins_id] = new_gaussian["embedding"]
+
+ def export_gaussians_to_ply(
+ self, alpha_thresh: float, instance_id: List[int] = None, specific_frame: int = 0,
+ ) -> Dict[str, torch.Tensor]:
+ self.cur_frame = specific_frame
+ pts_mask = self.point_ids[..., 0] == instance_id
+
+ if self.ctrl_cfg.use_deformgs_for_nonrigid and self.step > self.ctrl_cfg.use_deformgs_after:
+ delta_xyz, _, _ = self.get_deformation(local_means=self._means)
+ means = self._means + delta_xyz
+ else:
+ means = self._means
+ means = means[pts_mask]
+ direct_color = self.colors[pts_mask]
+
+ activated_opacities = self.get_opacity[pts_mask]
+ mask = activated_opacities.squeeze() > alpha_thresh
+ return {
+ "positions": means[mask],
+ "colors": direct_color[mask],
+ }
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/models/nodes/rigid.py b/worldbench/third_party/drivestudio/models/nodes/rigid.py
new file mode 100644
index 0000000..7b3fd03
--- /dev/null
+++ b/worldbench/third_party/drivestudio/models/nodes/rigid.py
@@ -0,0 +1,614 @@
+from typing import Dict, List, Tuple
+import logging
+import random
+
+import torch
+from torch.nn import Parameter
+
+from models.modules import ConditionalDeformNetwork
+from models.gaussians.basics import *
+from models.gaussians.vanilla import VanillaGaussians
+
+logger = logging.getLogger()
+
+class RigidNodes(VanillaGaussians):
+ def __init__(
+ self,
+ **kwargs
+ ):
+ super().__init__(**kwargs)
+
+ @property
+ def num_instances(self):
+ return self.instances_fv.shape[1]
+ @property
+ def num_frames(self):
+ return self.instances_fv.shape[0]
+
+ def get_pts_valid_mask(self):
+ """
+ get the mask for valid points
+ """
+ return self.instances_fv[self.cur_frame][self.point_ids[..., 0]]
+
+ def set_cur_frame(self, frame_id: int):
+ self.cur_frame = frame_id
+ def register_normalized_timestamps(self, normalized_timestamps: int):
+ self.normalized_timestamps = normalized_timestamps
+
+ def create_from_pcd(self, instance_pts_dict: Dict[str, torch.Tensor]) -> None:
+ """
+ instance_pts_dict: {
+ id in dataset: {
+ "class_name": str,
+ "pts": torch.Tensor, (N, 3)
+ "colors": torch.Tensor, (N, 3)
+ "poses": torch.Tensor, (num_frame, 4, 4)
+ "size": torch.Tensor, (3, )
+ "frame_info": torch.Tensor, (num_frame)
+ "num_pts": int,
+ },
+ }
+ """
+ # collect all instances
+ init_means = []
+ init_colors = []
+ instances_pose = []
+ instances_size = []
+ instances_fv = []
+ point_ids = []
+ for id_in_model, (id_in_dataset, v) in enumerate(instance_pts_dict.items()):
+ init_means.append(v["pts"])
+ init_colors.append(v["colors"])
+ instances_pose.append(v["poses"].unsqueeze(1))
+ instances_size.append(v["size"])
+ instances_fv.append(v["frame_info"].unsqueeze(1))
+ point_ids.append(torch.full((v["num_pts"], 1), id_in_model, dtype=torch.long))
+ init_means = torch.cat(init_means, dim=0).to(self.device) # (N, 3)
+ init_colors = torch.cat(init_colors, dim=0).to(self.device) # (N, 3)
+ instances_pose = torch.cat(instances_pose, dim=1).to(self.device) # (num_frame, num_instances, 4, 4)
+ self.instances_size = torch.stack(instances_size).to(self.device) # (num_instances, 3)
+ self.instances_fv = torch.cat(instances_fv, dim=1).to(self.device) # (num_frame, num_instances)
+ self.point_ids = torch.cat(point_ids, dim=0).to(self.device)
+ instances_quats = self.get_instances_quats(instances_pose)
+ instances_trans = instances_pose[..., :3, 3]
+
+ # initialize the means, scales, quats, and colors
+ self._means = Parameter(init_means)
+ distances_np, _ = k_nearest_sklearn(self._means.data, 3)
+ distances = torch.from_numpy(distances_np)
+ if distances.shape[1] > 0:
+ avg_dist = distances.mean(dim=-1, keepdim=True).to(self.device)
+ scales_init = avg_dist.expand(-1, 3).clone()
+ else:
+ default_scale = 1e-3
+ scales_init = torch.full((self.num_points, 3), default_scale, device=self.device)
+ scales_init = scales_init.clamp(0.002, 100)
+ self._scales = Parameter(torch.log(scales_init))
+ self._quats = Parameter(random_quat_tensor(self.num_points).to(self.device))
+ dim_sh = num_sh_bases(self.sh_degree)
+
+ # pose refinement
+ self.instances_quats = Parameter(self.quat_act(instances_quats)) # (num_frame, num_instances, 4)
+ self.instances_trans = Parameter(instances_trans) # (num_frame, num_instances, 3)
+
+ fused_color = RGB2SH(init_colors) # float range [0, 1]
+ shs = torch.zeros((fused_color.shape[0], dim_sh, 3)).float().to(self.device)
+ if self.sh_degree > 0:
+ shs[:, 0, :3] = fused_color
+ shs[:, 1:, 3:] = 0.0
+ else:
+ shs[:, 0, :3] = torch.logit(init_colors, eps=1e-10)
+ self._features_dc = Parameter(shs[:, 0, :])
+ self._features_rest = Parameter(shs[:, 1:, :])
+ self._opacities = Parameter(torch.logit(0.1 * torch.ones(self.num_points, 1, device=self.device)))
+
+ def get_param_groups(self) -> Dict[str, List[Parameter]]:
+ param_groups = self.get_gaussian_param_groups()
+ param_groups[self.class_prefix+"ins_rotation"] = [self.instances_quats]
+ param_groups[self.class_prefix+"ins_translation"] = [self.instances_trans]
+ return param_groups
+
+ def get_instances_quats(self, instances_pose: torch.Tensor) -> torch.Tensor:
+ """
+ Convert the pose to quaternion for all frames and instances
+ """
+ num_frames = instances_pose.shape[0]
+ num_instances = instances_pose.shape[1]
+ quats = torch.zeros(num_frames*num_instances, 4, device=self.device)
+
+ poses = instances_pose[..., :3, :3].view(-1, 3, 3)
+ valid_mask = self.instances_fv.view(-1)
+ _quats = matrix_to_quaternion(poses[valid_mask])
+ _quats = self.quat_act(_quats)
+
+ quats[valid_mask] = _quats
+ quats[~valid_mask, 0] = 1.0
+ return quats.reshape(num_frames, num_instances, 4)
+
+ def refinement_after(self, step: int, optimizer: torch.optim.Optimizer) -> None:
+ assert step == self.step
+ if self.step <= self.ctrl_cfg.warmup_steps:
+ return
+ with torch.no_grad():
+ # only split/cull if we've seen every image since opacity reset
+ reset_interval = self.ctrl_cfg.reset_alpha_interval
+ do_densification = (
+ self.step < self.ctrl_cfg.stop_split_at
+ and self.step % reset_interval > max(self.num_train_images, self.ctrl_cfg.refine_interval)
+ )
+ # split & duplicate
+ print(f"Class {self.class_prefix} current points: {self.num_points} @ step {self.step}")
+ if do_densification:
+ assert self.xys_grad_norm is not None and self.vis_counts is not None and self.max_2Dsize is not None
+
+ avg_grad_norm = self.xys_grad_norm / self.vis_counts
+ # Keep mask 1D; avoid 0-D when num_points==1
+ high_grads = (avg_grad_norm > self.ctrl_cfg.densify_grad_thresh).reshape(-1)
+
+ splits = (
+ self.get_scaling.max(dim=-1).values >
+ self.ctrl_cfg.densify_size_thresh * self.scene_scale
+ ).reshape(-1)
+ if self.step < self.ctrl_cfg.stop_screen_size_at:
+ splits = splits | (self.max_2Dsize > self.ctrl_cfg.split_screen_size).reshape(-1)
+ splits = splits & high_grads
+ nsamps = self.ctrl_cfg.n_split_samples
+ (
+ split_means,
+ split_feature_dc,
+ split_feature_rest,
+ split_opacities,
+ split_scales,
+ split_quats,
+ split_ids,
+ ) = self.split_gaussians(splits, nsamps)
+
+ dups = (
+ self.get_scaling.max(dim=-1).values <=
+ self.ctrl_cfg.densify_size_thresh * self.scene_scale
+ ).reshape(-1)
+ dups = dups & high_grads
+ (
+ dup_means,
+ dup_feature_dc,
+ dup_feature_rest,
+ dup_opacities,
+ dup_scales,
+ dup_quats,
+ dup_ids,
+ ) = self.dup_gaussians(dups)
+
+ self._means = Parameter(torch.cat([self._means.detach(), split_means, dup_means], dim=0))
+ # self.colors_all = Parameter(torch.cat([self.colors_all.detach(), split_colors, dup_colors], dim=0))
+ self._features_dc = Parameter(torch.cat([self._features_dc.detach(), split_feature_dc, dup_feature_dc], dim=0))
+ self._features_rest = Parameter(torch.cat([self._features_rest.detach(), split_feature_rest, dup_feature_rest], dim=0))
+ self._opacities = Parameter(torch.cat([self._opacities.detach(), split_opacities, dup_opacities], dim=0))
+ self._scales = Parameter(torch.cat([self._scales.detach(), split_scales, dup_scales], dim=0))
+ self._quats = Parameter(torch.cat([self._quats.detach(), split_quats, dup_quats], dim=0))
+ self.point_ids = torch.cat([self.point_ids, split_ids, dup_ids], dim=0)
+
+ # append zeros to the max_2Dsize tensor
+ self.max_2Dsize = torch.cat(
+ [self.max_2Dsize, torch.zeros_like(split_scales[:, 0]), torch.zeros_like(dup_scales[:, 0])],
+ dim=0,
+ )
+
+ split_idcs = torch.where(splits)[0]
+ param_groups = self.get_gaussian_param_groups()
+ dup_in_optim(optimizer, split_idcs, param_groups, n=nsamps)
+
+ dup_idcs = torch.where(dups)[0]
+ param_groups = self.get_gaussian_param_groups()
+ dup_in_optim(optimizer, dup_idcs, param_groups, 1)
+
+ # cull NOTE: Offset all the opacity reset logic by refine_every so that we don't
+ # save checkpoints right when the opacity is reset (saves every 2k)
+ if self.step % reset_interval > max(self.num_train_images, self.ctrl_cfg.refine_interval):
+ deleted_mask = self.cull_gaussians()
+ param_groups = self.get_gaussian_param_groups()
+ remove_from_optim(optimizer, deleted_mask, param_groups)
+ print(f"Class {self.class_prefix} left points: {self.num_points}")
+
+ # reset opacity
+ if self.step % reset_interval == self.ctrl_cfg.refine_interval:
+ # NOTE: in nerfstudio, reset_value = cull_alpha_thresh * 0.8
+ # we align to original repo of gaussians spalting
+ reset_value = torch.min(self.get_opacity.data,
+ torch.ones_like(self._opacities.data) * self.ctrl_cfg.reset_alpha_value)
+ self._opacities.data = torch.logit(reset_value)
+ # reset the exp of optimizer
+ for group in optimizer.param_groups:
+ if group["name"] == self.class_prefix+"opacity":
+ old_params = group["params"][0]
+ param_state = optimizer.state[old_params]
+ param_state["exp_avg"] = torch.zeros_like(param_state["exp_avg"])
+ param_state["exp_avg_sq"] = torch.zeros_like(param_state["exp_avg_sq"])
+ self.xys_grad_norm = None
+ self.vis_counts = None
+ self.max_2Dsize = None
+
+ def cull_gaussians(self):
+ """
+ This function deletes gaussians with under a certain opacity threshold
+ """
+ n_bef = self.num_points
+ # cull transparent ones
+ culls = (self.get_opacity.data < self.ctrl_cfg.cull_alpha_thresh).squeeze()
+ if self.ctrl_cfg.cull_out_of_bound:
+ culls = culls | self.get_out_of_bound_mask()
+ if self.step > self.ctrl_cfg.reset_alpha_interval:
+ # cull huge ones
+ toobigs = (
+ torch.exp(self._scales).max(dim=-1).values >
+ self.ctrl_cfg.cull_scale_thresh * self.scene_scale
+ ).squeeze()
+ culls = culls | toobigs
+ if self.step < self.ctrl_cfg.stop_screen_size_at:
+ # cull big screen space
+ assert self.max_2Dsize is not None
+ culls = culls | (self.max_2Dsize > self.ctrl_cfg.cull_screen_size).squeeze()
+ self._means = Parameter(self._means[~culls].detach())
+ self._scales = Parameter(self._scales[~culls].detach())
+ self._quats = Parameter(self._quats[~culls].detach())
+ # self.colors_all = Parameter(self.colors_all[~culls].detach())
+ self._features_dc = Parameter(self._features_dc[~culls].detach())
+ self._features_rest = Parameter(self._features_rest[~culls].detach())
+ self._opacities = Parameter(self._opacities[~culls].detach())
+ self.point_ids = self.point_ids[~culls]
+
+ print(f" Cull: {n_bef - self.num_points}")
+ return culls
+
+ def split_gaussians(self, split_mask: torch.Tensor, samps: int = 2) -> Tuple:
+ """
+ This function splits gaussians that are too large
+ """
+ # Normalize mask shape to 1D to avoid 0-D scalars when N == 1
+ if split_mask.dim() != 1:
+ split_mask = split_mask.reshape(-1)
+
+ n_splits = int(split_mask.sum().item())
+ print(f" Split: {n_splits}")
+ # Early return on zero-split to avoid downstream ops on empty tensors
+ if n_splits == 0:
+ device = self.device
+ return (
+ torch.empty((0, self._means.shape[-1]), device=device, dtype=self._means.dtype),
+ torch.empty((0, self._features_dc.shape[-1]), device=device, dtype=self._features_dc.dtype),
+ torch.empty((0, self._features_rest.shape[1], self._features_rest.shape[-1]), device=device, dtype=self._features_rest.dtype),
+ torch.empty((0, self._opacities.shape[-1]), device=device, dtype=self._opacities.dtype),
+ torch.empty((0, self._scales.shape[-1]), device=device, dtype=self._scales.dtype),
+ torch.empty((0, self._quats.shape[-1]), device=device, dtype=self._quats.dtype),
+ torch.empty((0, self.point_ids.shape[-1]), device=device, dtype=self.point_ids.dtype),
+ )
+
+ centered_samples = torch.randn((samps * n_splits, 3), device=self.device) # Nx3 of axis-aligned scales
+ # Normalize shapes to 2D along the last dim, preserving order
+ scales_selected_raw = self._scales[split_mask]
+ scales_selected = torch.exp(scales_selected_raw.reshape(-1, scales_selected_raw.shape[-1])) # (n_splits, C)
+ scaled_samples = (
+ scales_selected.repeat(samps, 1) * centered_samples
+ ) # how these scales are rotated
+ quats_raw = self._quats[split_mask]
+ quats = self.quat_act(quats_raw.reshape(-1, quats_raw.shape[-1])) # normalize them first, (n_splits, 4)
+ rots = quat_to_rotmat(quats.repeat(samps, 1)) # (samps*n_splits, 3, 3)
+ rotated_samples = torch.bmm(rots, scaled_samples[..., None]).squeeze()
+ means_raw = self._means[split_mask]
+ new_means = rotated_samples + means_raw.reshape(-1, means_raw.shape[-1]).repeat(samps, 1)
+ # step 2, sample new colors
+ # new_colors_all = self.colors_all[split_mask].repeat(samps, 1, 1)
+ new_feature_dc = self._features_dc[split_mask].reshape(-1, self._features_dc.shape[-1]).repeat(samps, 1)
+ new_feature_rest = self._features_rest[split_mask].repeat(samps, 1, 1)
+ # step 3, sample new opacities
+ new_opacities = self._opacities[split_mask].reshape(-1, self._opacities.shape[-1]).repeat(samps, 1)
+ # step 4, sample new scales
+ size_fac = 1.6
+ _sel_scales = self._scales[split_mask]
+ new_scales = torch.log(torch.exp(_sel_scales) / size_fac).reshape(-1, _sel_scales.shape[-1]).repeat(samps, 1)
+ self._scales[split_mask] = torch.log(torch.exp(self._scales[split_mask]) / size_fac)
+ # step 5, sample new quats
+ new_quats = self._quats[split_mask].reshape(-1, quats_raw.shape[-1]).repeat(samps, 1)
+ # step 6, sample new ids
+ new_ids = self.point_ids[split_mask].reshape(-1, self.point_ids.shape[-1]).repeat(samps, 1)
+ return new_means, new_feature_dc, new_feature_rest, new_opacities, new_scales, new_quats, new_ids
+
+ def dup_gaussians(self, dup_mask: torch.Tensor) -> Tuple:
+ """
+ This function duplicates gaussians that are too small
+ """
+ n_dups = dup_mask.sum().item()
+ print(f" Dup: {n_dups}")
+ dup_means = self._means[dup_mask]
+ # dup_colors = self.colors_all[dup_mask]
+ dup_feature_dc = self._features_dc[dup_mask]
+ dup_feature_rest = self._features_rest[dup_mask]
+ dup_opacities = self._opacities[dup_mask]
+ dup_scales = self._scales[dup_mask]
+ dup_quats = self._quats[dup_mask]
+ dup_ids = self.point_ids[dup_mask]
+ return dup_means, dup_feature_dc, dup_feature_rest, dup_opacities, dup_scales, dup_quats, dup_ids
+
+ def get_out_of_bound_mask(self):
+ """
+ This function checks if the gaussians are out of instance boxes
+ """
+ # get the instance boxes
+ per_pts_size = self.instances_size[self.point_ids[..., 0]]
+ instance_pts = self._means
+
+ mask = (instance_pts.abs() > per_pts_size / 2).any(dim=-1)
+ return mask
+
+ def transform_means(self, means: torch.Tensor) -> torch.Tensor:
+ """
+ transform the means of instances to world space
+ according to the pose at the current frame
+ """
+ # Normalize input shapes for robustness when N==1
+ if means.dim() == 1:
+ # Accept (3,) → (1,3)
+ means = means.reshape(1, -1)
+ else:
+ # Ensure (N, 3)
+ means = means.reshape(-1, means.shape[-1])
+ assert means.shape[0] == self.point_ids.shape[0], \
+ "its a bug here, we need to pass the mask for points_ids"
+ if self.in_test_set and (
+ self.cur_frame - 1 > 0 and self.cur_frame + 1 < self.num_frames
+ ):
+ # use the previous and next frame to interpolate the pose
+ _quats_prev_frame = self.instances_quats[self.cur_frame - 1]
+ _quats_next_frame = self.instances_quats[self.cur_frame + 1]
+ _quats_cur_frame = self.instances_quats[self.cur_frame]
+ interpolated_quats = interpolate_quats(_quats_prev_frame, _quats_next_frame)
+
+ inter_valid_mask = self.instances_fv[self.cur_frame - 1] & self.instances_fv[self.cur_frame + 1]
+ quats_cur_frame = torch.where(
+ inter_valid_mask[:, None], interpolated_quats, _quats_cur_frame
+ )
+ else:
+ quats_cur_frame = self.instances_quats[self.cur_frame] # (num_instances, 4)
+ rot_cur_frame = quat_to_rotmat(self.quat_act(quats_cur_frame)) # (num_instances, 3, 3)
+ rot_per_pts = rot_cur_frame[self.point_ids[..., 0]]
+ if rot_per_pts.dim() == 2:
+ # Handle degenerate case when advanced indexing returns (3,3)
+ rot_per_pts = rot_per_pts.unsqueeze(0)
+
+ if self.in_test_set and (
+ self.cur_frame - 1 > 0 and self.cur_frame + 1 < self.num_frames
+ ):
+ _prev_ins_trans = self.instances_trans[self.cur_frame - 1]
+ _next_ins_trans = self.instances_trans[self.cur_frame + 1]
+ _cur_ins_trans = self.instances_trans[self.cur_frame]
+ interpolated_trans = (_prev_ins_trans + _next_ins_trans) * 0.5
+
+ inter_valid_mask = self.instances_fv[self.cur_frame - 1] & self.instances_fv[self.cur_frame + 1]
+ trans_cur_frame = torch.where(
+ inter_valid_mask[:, None], interpolated_trans, _cur_ins_trans
+ )
+ else:
+ trans_cur_frame = self.instances_trans[self.cur_frame] # (num_instances, 3)
+ trans_per_pts = trans_cur_frame[self.point_ids[..., 0]]
+ if trans_per_pts.dim() == 1:
+ trans_per_pts = trans_per_pts.unsqueeze(0)
+
+ # transform the means to world space
+ means = torch.bmm(rot_per_pts, means.unsqueeze(-1)).squeeze(-1) + trans_per_pts
+ return means
+
+ def transform_quats(self, quats: torch.Tensor) -> torch.Tensor:
+ """
+ transform the quats of instances to world space
+ according to the pose at the current frame
+ """
+ # Normalize shapes for robustness when N==1
+ if quats.dim() == 1:
+ quats = quats.reshape(1, -1)
+ else:
+ quats = quats.reshape(-1, quats.shape[-1])
+ assert quats.shape[0] == self.point_ids.shape[0], \
+ "its a bug here, we need to pass the mask for points_ids"
+ global_quats_cur_frame = self.instances_quats[self.cur_frame]
+ global_quats_per_pts = global_quats_cur_frame[self.point_ids[..., 0]]
+ if global_quats_per_pts.dim() == 1:
+ global_quats_per_pts = global_quats_per_pts.unsqueeze(0)
+
+ global_quats_per_pts = self.quat_act(global_quats_per_pts)
+ _quats = self.quat_act(quats)
+ return quat_mult(global_quats_per_pts, _quats)
+
+ def get_gaussians(self, cam: dataclass_camera) -> Dict[str, torch.Tensor]:
+ filter_mask = torch.ones_like(self._means[:, 0], dtype=torch.bool)
+ self.filter_mask = filter_mask
+ # NOTE: hack here, need to consider a gaussian filter for efficient rendering
+
+ world_means = self.transform_means(self._means)
+ world_quats = self.transform_quats(self._quats)
+
+ # get colors of gaussians
+ colors = torch.cat((self._features_dc[:, None, :], self._features_rest), dim=1)
+ if self.sh_degree > 0:
+ viewdirs = world_means.detach() - cam.camtoworlds.data[..., :3, 3] # (N, 3)
+ viewdirs = viewdirs / viewdirs.norm(dim=-1, keepdim=True)
+ n = min(self.step // self.ctrl_cfg.sh_degree_interval, self.sh_degree)
+ rgbs = spherical_harmonics(n, viewdirs, colors)
+ rgbs = torch.clamp(rgbs + 0.5, 0.0, 1.0)
+ else:
+ rgbs = torch.sigmoid(colors[:, 0, :])
+
+ valid_mask = self.get_pts_valid_mask()
+
+ activated_opacities = self.get_opacity * valid_mask.float().unsqueeze(-1)
+ activated_scales = self.get_scaling
+ activated_rotations = self.quat_act(world_quats)
+ actovated_colors = rgbs
+
+ # collect gaussians information
+ gs_dict = dict(
+ _means=world_means[filter_mask],
+ _opacities=activated_opacities[filter_mask],
+ _rgbs=actovated_colors[filter_mask],
+ _scales=activated_scales[filter_mask],
+ _quats=activated_rotations[filter_mask],
+ )
+
+ # check nan and inf in gs_dict
+ for k, v in gs_dict.items():
+ if torch.isnan(v).any():
+ raise ValueError(f"NaN detected in gaussian {k} at step {self.step}")
+ if torch.isinf(v).any():
+ raise ValueError(f"Inf detected in gaussian {k} at step {self.step}")
+
+ self._gs_cache = {
+ "_scales": activated_scales[filter_mask],
+ }
+ return gs_dict
+
+ def get_instance_activated_gs_dict(self, ins_id: int) -> Dict[str, torch.Tensor]:
+ pts_mask = self.point_ids[..., 0] == ins_id
+ if pts_mask.sum() < 100:
+ return None
+ local_means = self._means[pts_mask]
+ activated_opacities = torch.sigmoid(self._opacities[pts_mask])
+ activated_scales = torch.exp(self._scales[pts_mask])
+ activated_local_rotations = self.quat_act(self._quats[pts_mask])
+ gaussian_dict = {
+ "means": local_means,
+ "opacities": activated_opacities,
+ "scales": activated_scales,
+ "quats": activated_local_rotations,
+ "sh_dcs": self._features_dc[pts_mask],
+ "sh_rests": self._features_rest[pts_mask],
+ "ids": self.point_ids[pts_mask],
+ }
+ return gaussian_dict
+
+ def compute_reg_loss(self) -> Dict[str, torch.Tensor]:
+ loss_dict = super().compute_reg_loss()
+ scaling_reg = self.reg_cfg.get("scaling_reg", None)
+ if scaling_reg is not None:
+ w = scaling_reg.w
+ precentile = scaling_reg.precentile
+ stop_after = scaling_reg.stop_after
+ start_after = scaling_reg.start_after
+
+ if self.step < stop_after and self.step > start_after and w > 0:
+ scale_prod = self._gs_cache["_scales"].prod(dim=-1)
+ p = torch.kthvalue(scale_prod, int(scale_prod.shape[0] * precentile)).values
+ # penalize the scales that are too large
+ loss_dict["scaling_percentile_reg"] = torch.relu(scale_prod - p).mean() * w
+
+ # temporal smooth regularization
+ temporal_smooth_reg = self.reg_cfg.get("temporal_smooth_reg", None)
+ if temporal_smooth_reg is not None:
+ instance_mask = self.instances_fv[self.cur_frame]
+ if instance_mask.sum() > 0:
+ trans_cfg = temporal_smooth_reg.get("trans", None)
+ if trans_cfg is not None:
+ fi_interval = random.randint(1, trans_cfg.smooth_range)
+ if self.cur_frame >= fi_interval and self.cur_frame < self.num_frames - fi_interval:
+ valid_mask = (
+ self.instances_fv[self.cur_frame - fi_interval] & \
+ self.instances_fv[self.cur_frame + fi_interval] & \
+ self.instances_fv[self.cur_frame]
+ )
+ if valid_mask.sum() > 0:
+ cur_trans = self.instances_trans[self.cur_frame]
+ pre_trans = self.instances_trans[self.cur_frame - fi_interval].data
+ next_trans = self.instances_trans[self.cur_frame + fi_interval].data
+ loss = (next_trans[valid_mask] + pre_trans[valid_mask] - 2 * cur_trans[valid_mask]).abs().mean()
+ loss_dict["trans_temporal_smooth"] = loss * trans_cfg.w
+ return loss_dict
+
+ def state_dict(self) -> Dict:
+ state_dict = super().state_dict()
+ state_dict.update({
+ "points_ids": self.point_ids,
+ "instances_size": self.instances_size,
+ "instances_fv": self.instances_fv,
+ })
+ return state_dict
+
+ def load_state_dict(self, state_dict: Dict, **kwargs) -> str:
+ self.point_ids = state_dict.pop("points_ids")
+ self.instances_size = state_dict.pop("instances_size")
+ self.instances_fv = state_dict.pop("instances_fv")
+ self.instances_trans = Parameter(
+ torch.zeros(self.num_frames, self.num_instances, 3, device=self.device)
+ )
+ self.instances_quats = Parameter(
+ torch.zeros(self.num_frames, self.num_instances, 4, device=self.device)
+ )
+ msg = super().load_state_dict(state_dict, **kwargs)
+ return msg
+
+ # editting functions
+ def remove_instances(self, remove_id_list: List[int]) -> None:
+ """
+ remove instances from the model
+
+ Args:
+ remove_id_list: list of instance ids to be removed
+ """
+ for ins_ids in remove_id_list:
+ mask = ~(self.point_ids[..., 0] == ins_ids)
+ self._means = Parameter(self._means[mask])
+ self._scales = Parameter(self._scales[mask])
+ self._quats = Parameter(self._quats[mask])
+ self._features_dc = Parameter(self._features_dc[mask])
+ self._features_rest = Parameter(self._features_rest[mask])
+ self._opacities = Parameter(self._opacities[mask])
+ self.point_ids = self.point_ids[mask]
+
+ def collect_gaussians_from_ids(self, ids: List[int]) -> Dict:
+ gaussian_dict = {}
+ for id in ids:
+ if id not in gaussian_dict:
+ instance_raw_dict = {
+ "_means": self._means[self.point_ids[..., 0] == id],
+ "_scales": self._scales[self.point_ids[..., 0] == id],
+ "_quats": self._quats[self.point_ids[..., 0] == id],
+ "_features_dc": self._features_dc[self.point_ids[..., 0] == id],
+ "_features_rest": self._features_rest[self.point_ids[..., 0] == id],
+ "_opacities": self._opacities[self.point_ids[..., 0] == id],
+ "point_ids": self.point_ids[self.point_ids[..., 0] == id],
+ }
+ gaussian_dict[id] = instance_raw_dict
+ return gaussian_dict
+
+ def replace_instances(self, replace_dict: Dict[int, int]) -> None:
+ """
+ replace instances from the model
+
+ Args:
+ replace_dict: {
+ ins_id(to be replaced): ins_id(replace with)
+ ...
+ }
+ """
+ new_gaussians_dict = self.collect_gaussians_from_ids(replace_dict.values())
+ for ins_id, new_id in replace_dict.items():
+ self.remove_instances([ins_id])
+ new_gaussian = new_gaussians_dict[new_id]
+ self._means = Parameter(torch.cat([self._means, new_gaussian["_means"]], dim=0))
+ self._scales = Parameter(torch.cat([self._scales, new_gaussian["_scales"]], dim=0))
+ self._quats = Parameter(torch.cat([self._quats, new_gaussian["_quats"]], dim=0))
+ self._features_dc = Parameter(torch.cat([self._features_dc, new_gaussian["_features_dc"]], dim=0))
+ self._features_rest = Parameter(torch.cat([self._features_rest, new_gaussian["_features_rest"]], dim=0))
+ self._opacities = Parameter(torch.cat([self._opacities, new_gaussian["_opacities"]], dim=0))
+ # keeps original point ids
+ self.point_ids = torch.cat([self.point_ids, torch.full_like(new_gaussian["point_ids"], ins_id)], dim=0)
+
+ def export_gaussians_to_ply(self, alpha_thresh: float, instance_id: List[int] = None) -> Dict[str, torch.Tensor]:
+ pts_mask = self.point_ids[..., 0] == instance_id
+
+ means = self._means[pts_mask]
+ direct_color = self.colors[pts_mask]
+
+ activated_opacities = self.get_opacity[pts_mask]
+ mask = activated_opacities.squeeze() > alpha_thresh
+ return {
+ "positions": means[mask],
+ "colors": direct_color[mask],
+ }
diff --git a/worldbench/third_party/drivestudio/models/nodes/smpl.py b/worldbench/third_party/drivestudio/models/nodes/smpl.py
new file mode 100644
index 0000000..8025f4c
--- /dev/null
+++ b/worldbench/third_party/drivestudio/models/nodes/smpl.py
@@ -0,0 +1,754 @@
+from typing import Dict, List, Tuple
+from omegaconf import OmegaConf
+import random
+import logging
+
+import torch
+import numpy as np
+from torch.nn import Parameter
+
+from models.human_body import phalp_colors, SMPLTemplate, get_on_mesh_init_geo_values, batch_rigid_transform, quaternion_to_matrix
+from models.gaussians.basics import *
+from models.nodes.rigid import RigidNodes
+from models.gaussians.vanilla import VanillaGaussians
+from pytorch3d.ops import knn_points
+
+RGB_tuples = torch.tensor(np.vstack([phalp_colors] * 10), dtype=torch.float32) / 255.0
+logger = logging.getLogger()
+
+class SMPLNodes(RigidNodes):
+ def __init__(
+ self,
+ **kwargs
+ ):
+ self.smpl_points_num = 6890
+ super().__init__(**kwargs)
+
+ self.use_voxel_deformer=self.ctrl_cfg.use_voxel_deformer
+ # overide here, because we use only one dimension for scale
+ if self.ball_gaussians:
+ self._scales = torch.zeros(1, 1, device=self.device)
+
+ @property
+ def num_instances(self):
+ return self.instances_fv.shape[1]
+ @property
+ def num_frames(self):
+ return self.instances_fv.shape[0]
+
+ def create_from_pcd(self, instance_pts_dict: Dict[str, torch.Tensor]) -> None:
+ """
+ instance_pts_dict: {
+ id in dataset: {
+ "class_name": str,
+ "pts": torch.Tensor, (N, 3)
+ "colors": torch.Tensor, (N, 3)
+ "poses": torch.Tensor, (num_frame, 4, 4)
+ "size": torch.Tensor, (3, )
+ "frame_info": torch.Tensor, (num_frame)
+ "num_pts": int,
+ },
+ }
+ """
+ """
+ For version 1:
+ we simplified gaussian properties:
+ - means: (Frame_N, N, 3), Not Optimized
+ - scales: (N, 1), Optimized, we use gaussians with the same scale on xyz
+ - quats: (N, 4), Not Optimized
+ - features_dc: (N, 3), Optimized
+ - features_rest: (N, num_sh_bases, 3), Optimized
+ - opacities: (N, 1), Optimized
+ """
+ # collect all instances
+ smpl_betas, smpl_qauts = [], []
+ instances_quats, instances_trans, instances_size = [], [], []
+ instances_fv, point_ids = [], []
+ # instances_pts, instances_colors = [], []
+ for id_in_model, (id_in_dataset, v) in enumerate(instance_pts_dict.items()):
+ smpl_qauts.append(v["smpl_quats"][:, 1:, :].unsqueeze(1))
+ instances_quats.append(v["smpl_quats"][:, 0, :].unsqueeze(1))
+ instances_trans.append(v["smpl_trans"].unsqueeze(1))
+ instances_fv.append(v["frame_info"].unsqueeze(1))
+ smpl_betas.append(v["smpl_betas"].unsqueeze(0))
+ instances_size.append(v["size"])
+ # instances_pts.append(v["pts"])
+ # instances_colors.append(v["colors"])
+ point_ids.append(torch.full((self.smpl_points_num, 1), id_in_model, dtype=torch.long))
+
+ smpl_qauts = torch.cat(smpl_qauts, dim=1).to(self.device) # (num_frame, num_instances, 23, 4)
+ instances_quats = torch.cat(instances_quats, dim=1).to(self.device) # (num_frame, num_instances, 4)
+ instances_trans = torch.cat(instances_trans, dim=1).to(self.device) # (num_frame, num_instances, 3)
+ instances_fv = torch.cat(instances_fv, dim=1).to(self.device) # (num_frame, num_instances)
+ smpl_betas = torch.cat(smpl_betas, dim=0).to(self.device) # (num_instances, 10)
+ instances_size = torch.stack(instances_size).to(self.device) # (num_instances, 3)
+ point_ids = torch.cat(point_ids, dim=0).to(self.device) # (self.smpl_points_num*num_instances, 1)
+ self.instances_fv = instances_fv # (num_frame, num_instances)
+
+ self.template = SMPLTemplate(
+ smpl_model_path="smpl_models/SMPL_NEUTRAL.pkl",
+ num_human=smpl_betas.shape[0],
+ init_beta=smpl_betas,
+ cano_pose_type="da_pose",
+ use_voxel_deformer=self.use_voxel_deformer
+ )
+ if self.use_voxel_deformer:
+ self.template.voxel_deformer.enable_voxel_correction()
+
+ opacity_init_value = torch.tensor(self.ctrl_cfg.opacity_init_value)
+ x, q, s, o = get_on_mesh_init_geo_values(
+ self.template,
+ opacity_init_logit=torch.logit(opacity_init_value),
+ )
+ if self.ball_gaussians:
+ s = s.mean(-1, keepdim=True)
+ x = x.to(dtype=torch.float32, device=self.device)
+ s = s.to(dtype=torch.float32, device=self.device)
+ q = q.to(dtype=torch.float32, device=self.device)
+ o = o.to(dtype=torch.float32, device=self.device)
+
+ # knn init
+ self.update_knn(x)
+
+ if self.ctrl_cfg.constrain_xyz_offset:
+ self.on_mesh_x = x.clone()
+
+ # NOTE: In the future, we will also use colors of lidars to get the initialization of colors
+ self.template = self.template.to(self.device)
+ for fi in range(self.num_frames):
+ instance_mask = instances_fv[fi]
+ if instance_mask.sum() == 0:
+ continue
+
+ theta = torch.cat(
+ (instances_quats[fi].unsqueeze(1), smpl_qauts[fi]), dim=1
+ )
+ masked_theta = theta[instance_mask]
+ masked_theta = masked_theta / masked_theta.norm(dim=-1, keepdim=True)
+ W, A = self.template(
+ masked_theta = masked_theta,
+ instances_mask = instance_mask
+ )
+ T = torch.einsum("bnj, bjrc -> bnrc", W, A)
+ R = T[:, :, :3, :3] # [N, 3, 3]
+ t = T[:, :, :3, 3] # [N, 3]
+
+ reshaped_means = x.reshape(self.num_instances, self.smpl_points_num, 3)
+ deformed_means = torch.einsum(
+ "bnij,bnj->bni", R, reshaped_means[instance_mask]
+ ) + t # [N, 6890, 3]
+ bbox_min = deformed_means.min(dim=1)[0]
+ bbox_max = deformed_means.max(dim=1)[0]
+ local_shift = (bbox_min + bbox_max) / 2
+ instances_trans[fi, instance_mask] = instances_trans[fi, instance_mask] - local_shift
+
+ self._means = Parameter(x, requires_grad=not self.ctrl_cfg.freeze_x)
+ self._scales = Parameter(s, requires_grad=not self.ctrl_cfg.freeze_s)
+ self._quats = Parameter(q, requires_grad=not self.ctrl_cfg.freeze_q)
+ self._opacities = Parameter(o, requires_grad=not self.ctrl_cfg.freeze_o)
+
+ self.instances_quats = Parameter(instances_quats.unsqueeze(2)) # (num_frame, num_instances, 1, 4)
+ self.instances_trans = Parameter(instances_trans) # (num_frame, num_instances, 3)
+ self.smpl_qauts = Parameter(smpl_qauts) # (num_frame, num_instances, 23, 4)
+ self.instances_size = instances_size # (num_instances, 3)
+ self.point_ids = point_ids # (self.smpl_points_num*num_instances, 1)
+
+ dim_sh = num_sh_bases(self.sh_degree)
+ # NOTE: init_colors actually is for visualization, we use random color here
+ # init_colors = RGB_tuples[self.point_ids.squeeze().cpu()].to(self.device)
+ init_colors = torch.rand((self.num_points, 3), device=self.device)
+ fused_color = RGB2SH(init_colors) # float range [0, 1]
+ shs = torch.zeros((fused_color.shape[0], dim_sh, 3)).float().to(self.device)
+ if self.sh_degree > 0:
+ shs[:, 0, :3] = fused_color
+ shs[:, 1:, 3:] = 0.0
+ else:
+ shs[:, 0, :3] = torch.logit(init_colors, eps=1e-10)
+ self._features_dc = Parameter(shs[:, 0, :], requires_grad=not self.ctrl_cfg.freeze_shs_dc)
+ self._features_rest = Parameter(shs[:, 1:, :], requires_grad=not self.ctrl_cfg.freeze_shs_rest)
+
+ def get_param_groups(self) -> Dict[str, List[Parameter]]:
+ param_groups = self.get_gaussian_param_groups()
+ param_groups[self.class_prefix+"ins_rotation"] = [self.instances_quats]
+ param_groups[self.class_prefix+"ins_translation"] = [self.instances_trans]
+ param_groups[self.class_prefix+"smpl_rotation"] = [self.smpl_qauts]
+ if self.use_voxel_deformer:
+ param_groups[self.class_prefix+"w_dc_vox"] = [self.template.voxel_deformer.voxel_w_correction]
+ # param_groups[self.class_prefix+"w_rest_vox"] = [self.template.voxel_deformer.additional_correction]
+ return param_groups
+
+ @property
+ def num_points(self):
+ return self._means.shape[0]
+
+ def update_knn(self, x: torch.Tensor) -> None:
+ reshaped_x = x.reshape(self.num_instances, self.smpl_points_num, 3)
+ _, nn_ind, _ = knn_points(reshaped_x, reshaped_x, K=self.ctrl_cfg.knn_neighbors, return_nn=False)
+ self.nn_ind = nn_ind
+
+ def postprocess_per_train_step(
+ self,
+ step: int,
+ optimizer: torch.optim.Optimizer,
+ radii: torch.Tensor,
+ xys_grad: torch.Tensor,
+ last_size: int,
+ ) -> None:
+ self.radii = radii
+ self.xys_grad = xys_grad
+ knn_update_interval = self.ctrl_cfg.get("knn_update_interval", 1000000)
+ if self.step % knn_update_interval == 0:
+ self.update_knn(self._means)
+
+ def transform_means(self, means: torch.Tensor) -> torch.Tensor:
+ """
+ transform the means of instances to world space
+ according to the pose at the current frame
+ """
+ assert means.shape[0] == self.point_ids.shape[0], \
+ "its a bug here, we need to pass the mask for points_ids"
+ instance_mask = self.instances_fv[self.cur_frame]
+ if self.in_test_set and (
+ self.cur_frame - 1 > 0 and self.cur_frame + 1 < self.num_frames
+ ):
+ _prev_masked_theta = torch.cat((self.instances_quats[self.cur_frame - 1], self.smpl_qauts[self.cur_frame - 1]), dim=1)[instance_mask]
+ _next_masked_theta = torch.cat((self.instances_quats[self.cur_frame + 1], self.smpl_qauts[self.cur_frame + 1]), dim=1)[instance_mask]
+ _cur_masked_theta = torch.cat((self.instances_quats[self.cur_frame], self.smpl_qauts[self.cur_frame]), dim=1)[instance_mask]
+ interpolated_theta = interpolate_quats(_prev_masked_theta, _next_masked_theta)
+
+ inter_valid_mask = self.instances_fv[self.cur_frame - 1, instance_mask] & self.instances_fv[self.cur_frame + 1, instance_mask]
+ masked_theta = torch.where(
+ inter_valid_mask[:, None, None], interpolated_theta, _cur_masked_theta
+ )
+ else:
+ theta = torch.cat(
+ (self.instances_quats[self.cur_frame], self.smpl_qauts[self.cur_frame]), dim=1
+ )
+ masked_theta = theta[instance_mask]
+ masked_theta = self.quat_act(masked_theta)
+ W, A = self.template(
+ masked_theta = masked_theta,
+ instances_mask = instance_mask,
+ xyz_canonical = means.reshape(self.num_instances, self.smpl_points_num, 3) if self.use_voxel_deformer else None
+ )
+ T = torch.einsum("bnj, bjrc -> bnrc", W, A)
+ R = T[:, :, :3, :3] # [N, 3, 3]
+ t = T[:, :, :3, 3] # [N, 3]
+
+ reshaped_means = means.reshape(self.num_instances, self.smpl_points_num, 3)
+ deformed_means = torch.einsum(
+ "bnij,bnj->bni", R, reshaped_means[instance_mask]
+ ) + t # [N, 6890, 3]
+
+ means_container = torch.zeros_like(reshaped_means)
+ means_container.index_add_(0, instance_mask.nonzero().squeeze(), deformed_means)
+ means_container = means_container.reshape(-1, 3)
+
+ if self.in_test_set and (
+ self.cur_frame - 1 > 0 and self.cur_frame + 1 < self.num_frames
+ ):
+ _prev_ins_trans = self.instances_trans[self.cur_frame - 1]
+ _next_ins_trans = self.instances_trans[self.cur_frame + 1]
+ _cur_ins_trans = self.instances_trans[self.cur_frame]
+ interpolated_trans = (_prev_ins_trans + _next_ins_trans) * 0.5
+
+ inter_valid_mask = self.instances_fv[self.cur_frame - 1] & self.instances_fv[self.cur_frame + 1]
+ trans_cur_frame = torch.where(
+ inter_valid_mask[:, None], interpolated_trans, _cur_ins_trans
+ )
+ else:
+ trans_cur_frame = self.instances_trans[self.cur_frame] # (num_instances, 3)
+ trans_per_pts = trans_cur_frame[self.point_ids[..., 0]]
+
+ # transform the means to world space
+ means_container += trans_per_pts
+ return means_container
+
+ def transform_means_and_quats(self, means: torch.Tensor, quats: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ transform the means and quats of gaussians to world space
+ according to the pose at the current frame
+ """
+ assert means.shape[0] == self.point_ids.shape[0], \
+ "its a bug here, we need to pass the mask for points_ids"
+ instance_mask = self.instances_fv[self.cur_frame]
+ if self.in_test_set and (
+ self.cur_frame - 1 > 0 and self.cur_frame + 1 < self.num_frames
+ ):
+ _prev_masked_theta = torch.cat((self.instances_quats[self.cur_frame - 1], self.smpl_qauts[self.cur_frame - 1]), dim=1)[instance_mask]
+ _next_masked_theta = torch.cat((self.instances_quats[self.cur_frame + 1], self.smpl_qauts[self.cur_frame + 1]), dim=1)[instance_mask]
+ _cur_masked_theta = torch.cat((self.instances_quats[self.cur_frame], self.smpl_qauts[self.cur_frame]), dim=1)[instance_mask]
+ interpolated_theta = interpolate_quats(_prev_masked_theta, _next_masked_theta)
+
+ inter_valid_mask = self.instances_fv[self.cur_frame - 1, instance_mask] & self.instances_fv[self.cur_frame + 1, instance_mask]
+ masked_theta = torch.where(
+ inter_valid_mask[:, None, None], interpolated_theta, _cur_masked_theta
+ )
+ else:
+ theta = torch.cat(
+ (self.instances_quats[self.cur_frame], self.smpl_qauts[self.cur_frame]), dim=1
+ )
+ masked_theta = theta[instance_mask]
+ masked_theta = self.quat_act(masked_theta)
+ W, A = self.template(
+ masked_theta = masked_theta,
+ instances_mask = instance_mask,
+ xyz_canonical = means.reshape(self.num_instances, self.smpl_points_num, 3) if self.use_voxel_deformer else None
+ )
+ T = torch.einsum("bnj, bjrc -> bnrc", W, A)
+ R = T[:, :, :3, :3] # [N, 3, 3]
+ t = T[:, :, :3, 3] # [N, 3]
+
+ reshaped_means = means.reshape(self.num_instances, self.smpl_points_num, 3)
+ deformed_means = torch.einsum(
+ "bnij,bnj->bni", R, reshaped_means[instance_mask]
+ ) + t # [N, 6890, 3]
+
+ means_container = torch.zeros_like(reshaped_means)
+ means_container.index_add_(0, instance_mask.nonzero().squeeze(), deformed_means)
+ means_container = means_container.reshape(-1, 3)
+
+ if self.in_test_set and (
+ self.cur_frame - 1 > 0 and self.cur_frame + 1 < self.num_frames
+ ):
+ _prev_ins_trans = self.instances_trans[self.cur_frame - 1]
+ _next_ins_trans = self.instances_trans[self.cur_frame + 1]
+ _cur_ins_trans = self.instances_trans[self.cur_frame]
+ interpolated_trans = (_prev_ins_trans + _next_ins_trans) * 0.5
+
+ inter_valid_mask = self.instances_fv[self.cur_frame - 1] & self.instances_fv[self.cur_frame + 1]
+ trans_cur_frame = torch.where(
+ inter_valid_mask[:, None], interpolated_trans, _cur_ins_trans
+ )
+ else:
+ trans_cur_frame = self.instances_trans[self.cur_frame] # (num_instances, 3)
+ trans_per_pts = trans_cur_frame[self.point_ids[..., 0]]
+
+ # transform the means to world space
+ means_container += trans_per_pts
+
+ reshaped_quats = quats.reshape(self.num_instances, self.smpl_points_num, 4)
+ R_quats = matrix_to_quaternion(R)
+ deformed_quats = quat_mult(
+ self.quat_act(R_quats),
+ self.quat_act(reshaped_quats[instance_mask])
+ )
+ quats_container = torch.zeros_like(reshaped_quats)
+ quats_container.index_add_(0, instance_mask.nonzero().squeeze(), deformed_quats)
+ # fill other with [1, 0, 0, 0]
+ quats_container.index_add_(0, (~instance_mask).nonzero().squeeze(), torch.tensor([[[1., 0., 0., 0.]]], device=self.device).repeat((~instance_mask).sum(), 6890, 1))
+ quats_container = quats_container.reshape(-1, 4)
+ return means_container, quats_container
+
+ def get_gaussians(self, cam: dataclass_camera) -> Dict[str, torch.Tensor]:
+ filter_mask = torch.ones_like(self._means[:, 0], dtype=torch.bool)
+ self.filter_mask = filter_mask
+ # NOTE: hack here, need to consider a gaussian filter for efficient rendering
+
+ instance_mask = self.instances_fv[self.cur_frame]
+ if instance_mask.sum() == 0:
+ return None
+
+ if self.ball_gaussians:
+ world_means = self.transform_means(self._means)
+ world_quats = self._quats
+ else:
+ world_means, world_quats = self.transform_means_and_quats(self._means, self._quats)
+
+ # get colors of gaussians
+ colors = torch.cat((self._features_dc[:, None, :], self._features_rest), dim=1)
+ if self.sh_degree > 0:
+ viewdirs = world_means.detach() - cam.camtoworlds.data[..., :3, 3] # (N, 3)
+ viewdirs_norm = viewdirs.norm(dim=-1, keepdim=True)
+ viewdirs = viewdirs / viewdirs_norm.clamp_min(1e-6)
+ zero_dir_mask = viewdirs_norm.squeeze(-1) < 1e-6
+ if zero_dir_mask.any():
+ viewdirs[zero_dir_mask] = torch.tensor([0.0, 0.0, 1.0], device=viewdirs.device)
+ n = min(self.step // self.ctrl_cfg.sh_degree_interval, self.sh_degree)
+ rgbs = spherical_harmonics(n, viewdirs, colors)
+ rgbs = torch.clamp(rgbs + 0.5, 0.0, 1.0)
+ else:
+ rgbs = torch.sigmoid(colors[:, 0, :])
+
+ valid_mask = self.get_pts_valid_mask()
+
+ activated_opacities = self.get_opacity * valid_mask.float().unsqueeze(-1)
+ if self.ball_gaussians:
+ activated_scales = torch.exp(self._scales.repeat(1, 3))
+ else:
+ activated_scales = torch.exp(self._scales)
+ activated_rotations = self.quat_act(world_quats)
+ actovated_colors = rgbs
+
+ # collect gaussians information
+ gs_dict = dict(
+ _means=world_means[filter_mask],
+ _opacities=activated_opacities[filter_mask],
+ _rgbs=actovated_colors[filter_mask],
+ _scales=activated_scales[filter_mask],
+ _quats=activated_rotations[filter_mask],
+ )
+ # check nan in gs_dict
+ for k, v in gs_dict.items():
+ if torch.isnan(v).any():
+ raise ValueError(f"NaN detected in gaussian {k} at step {self.step}")
+ if torch.isinf(v).any():
+ raise ValueError(f"Inf detected in gaussian {k} at step {self.step}")
+
+ self._gs_cache = {
+ "_scales": activated_scales[filter_mask],
+ }
+ return gs_dict
+
+ def compute_reg_loss(self):
+ loss_dict = super().compute_reg_loss()
+
+ instance_mask = self.instances_fv[self.cur_frame]
+ if instance_mask.sum() == 0:
+ return loss_dict
+
+ # temporal smooth regularization
+ temporal_smooth_reg = self.reg_cfg.get("temporal_smooth_reg", None)
+ if temporal_smooth_reg is not None:
+ joint_smooth_reg = temporal_smooth_reg.get("joint_smooth", None)
+ if joint_smooth_reg is not None:
+ if self.cur_frame >= 1 and self.cur_frame < self.num_frames - 1:
+ valid_mask = (
+ self.instances_fv[self.cur_frame - 1] & \
+ self.instances_fv[self.cur_frame + 1] & \
+ self.instances_fv[self.cur_frame]
+ )
+ cur_theta = torch.cat(
+ (self.instances_quats[self.cur_frame], self.smpl_qauts[self.cur_frame]), dim=1
+ )[valid_mask]
+ next_theta = torch.cat(
+ (self.instances_quats[self.cur_frame + 1], self.smpl_qauts[self.cur_frame + 1]), dim=1
+ )[valid_mask]
+ prev_theta = torch.cat(
+ (self.instances_quats[self.cur_frame - 1], self.smpl_qauts[self.cur_frame - 1]), dim=1
+ )[valid_mask]
+ thetas = torch.vstack([prev_theta, cur_theta, next_theta])
+ thetas = self.quat_act(thetas)
+ J_transformed, _ = batch_rigid_transform(
+ quaternion_to_matrix(thetas),
+ self.template.J_canonical[valid_mask].repeat(3, 1, 1),
+ self.template._template_layer.parents,
+ )
+
+ cur_trans = self.instances_trans[self.cur_frame, valid_mask]
+ next_trans = self.instances_trans[self.cur_frame + 1, valid_mask]
+ prev_trans = self.instances_trans[self.cur_frame - 1, valid_mask]
+ trans = torch.vstack([prev_trans, cur_trans, next_trans])
+ J_transformed += trans.unsqueeze(-2)
+ J_transformed = J_transformed.reshape(3, -1, 24, 3)
+
+ velocity_prev = (J_transformed[1] - J_transformed[0])
+ velocity_next = (J_transformed[2] - J_transformed[1])
+ # l2 loss
+ loss_dict["smpl_temporal_smooth"] = (velocity_next - velocity_prev).abs().mean() \
+ * joint_smooth_reg.w
+
+ # voxel deformer regularization
+ voxel_deformer_reg = self.reg_cfg.get("voxel_deformer_reg", None)
+ if voxel_deformer_reg is not None and self.use_voxel_deformer:
+ w_std = self.template.voxel_deformer.get_tv("dc")
+ w_rest_std = self.template.voxel_deformer.get_tv("rest")
+ w_norm = self.template.voxel_deformer.get_mag("dc")
+ w_rest_norm = self.template.voxel_deformer.get_mag("rest")
+
+ loss_dict["voxel_deformer_reg"] = \
+ voxel_deformer_reg.lambda_std_w * w_std + \
+ voxel_deformer_reg.lambda_std_w_rest * w_rest_std + \
+ voxel_deformer_reg.lambda_w_norm * w_norm + \
+ voxel_deformer_reg.lambda_w_rest_norm * w_rest_norm
+
+ # knn regularization
+ knn_reg = self.reg_cfg.get("knn_reg", None)
+ if knn_reg is not None:
+ K = self.ctrl_cfg.knn_neighbors
+ instances_mask = self.instances_fv[self.cur_frame]
+ nn_ind = self.nn_ind[instances_mask] # (num_instances, smpl_points_num, knn_neighbors)
+
+ if not self.ctrl_cfg.freeze_shs_dc:
+ valid_shs_dc = self._features_dc.reshape(self.num_instances, self.smpl_points_num, 3)[instances_mask] # (num_instances, smpl_points_num, 3)
+ nn_ind_expanded = nn_ind.unsqueeze(-1).expand(-1, -1, -1, 3)
+ knn_shs_dc = torch.gather(valid_shs_dc.unsqueeze(2).expand(-1, -1, K, -1), 1, nn_ind_expanded) # (num_instances, smpl_points_num, knn_neighbors, 3)
+ shs_dc_std = knn_shs_dc.std(dim=2).mean()
+ loss_dict["knn_reg_dc"] = shs_dc_std * knn_reg.lambda_std_shs_dc
+
+ if not self.ctrl_cfg.freeze_shs_rest and self.sh_degree > 0:
+ dim_sh = num_sh_bases(self.sh_degree)
+ valid_shs_rest = self._features_rest.reshape(self.num_instances, self.smpl_points_num, -1)[instances_mask] # (num_instances, smpl_points_num, (dim_sh-1)*3)
+ nn_ind_expanded = nn_ind.unsqueeze(-1).expand(-1, -1, -1, (dim_sh-1)*3)
+ knn_shs_rest = torch.gather(valid_shs_rest.unsqueeze(2).expand(-1, -1, K, -1), 1, nn_ind_expanded) # (num_instances, smpl_points_num, knn_neighbors, (dim_sh-1)*3)
+ shs_rest_std = knn_shs_rest.std(dim=2).mean()
+ loss_dict["knn_reg_rest"] = shs_rest_std * knn_reg.lambda_std_shs_rest
+
+ if not self.ctrl_cfg.freeze_o:
+ valid_o = self.get_opacity.reshape(self.num_instances, self.smpl_points_num, 1)[instances_mask] # (num_instances, smpl_points_num, 1)
+ nn_ind_expanded = nn_ind.unsqueeze(-1).expand(-1, -1, -1, 1)
+ knn_o = torch.gather(valid_o.unsqueeze(2).expand(-1, -1, K, -1), 1, nn_ind_expanded)
+ o_std = knn_o.std(dim=2).mean()
+ loss_dict["knn_reg_o"] = o_std * knn_reg.lambda_std_o
+
+ if not self.ctrl_cfg.freeze_s:
+ scale_dim = 1 if self.ball_gaussians else 3
+ valid_s = self.get_scaling.reshape(self.num_instances, self.smpl_points_num, scale_dim)[instances_mask] # (num_instances, smpl_points_num, 1)
+ nn_ind_expanded = nn_ind.unsqueeze(-1).expand(-1, -1, -1, scale_dim)
+ knn_s = torch.gather(valid_s.unsqueeze(2).expand(-1, -1, K, -1), 1, nn_ind_expanded)
+ s_std = knn_s.std(dim=2).mean()
+ loss_dict["knn_reg_s"] = s_std * knn_reg.lambda_std_s
+
+ if not self.ctrl_cfg.freeze_q:
+ valid_q = self._quats.reshape(self.num_instances, self.smpl_points_num, 4)[instances_mask] # (num_instances, smpl_points_num, 4)
+ nn_ind_expanded = nn_ind.unsqueeze(-1).expand(-1, -1, -1, 4)
+ knn_q = torch.gather(valid_q.unsqueeze(2).expand(-1, -1, K, -1), 1, nn_ind_expanded)
+ q_std = knn_q.std(dim=2).mean()
+ loss_dict["knn_reg_q"] = q_std * knn_reg.lambda_std_q
+
+ # valid_x = self._means.reshape(self.num_instances, self.smpl_points_num, 3)[instances_mask] # (num_instances, smpl_points_num, 3)
+ # nn_ind_expanded = nn_ind.unsqueeze(-1).expand(-1, -1, -1, 3)
+ # knn_x = torch.gather(valid_x.unsqueeze(2).expand(-1, -1, K, -1), 1, nn_ind_expanded)
+ # x_std = knn_x.std(dim=2).mean()
+ # loss_dict["knn_reg_x"] = x_std * knn_reg.lambda_std_x
+
+ x_offset_reg = self.reg_cfg.get("x_offset", None)
+ if x_offset_reg is not None and self.ctrl_cfg.constrain_xyz_offset and not self.ctrl_cfg.freeze_x:
+ instances_mask = self.instances_fv[self.cur_frame]
+ valid_x = self._means.reshape(self.num_instances, self.smpl_points_num, 3)[instances_mask] # (num_instances, smpl_points_num, 3)
+ valid_x_on_mesh = self.on_mesh_x.reshape(self.num_instances, self.smpl_points_num, 3)[instances_mask]
+ x_offset = (valid_x - valid_x_on_mesh).norm(dim=-1).mean()
+
+ loss_dict["x_offset"] = x_offset * x_offset_reg.w
+
+ return loss_dict
+
+ def state_dict(self) -> Dict:
+ state_dict = VanillaGaussians.state_dict(self)
+ state_dict.update({
+ "points_ids": self.point_ids,
+ "instances_size": self.instances_size,
+ "instances_fv": self.instances_fv,
+ })
+ return state_dict
+
+ def load_state_dict(self, state_dict: Dict, **kwargs) -> str:
+ self.point_ids = state_dict.pop("points_ids")
+ self.instances_size = state_dict.pop("instances_size")
+ self.instances_fv = state_dict.pop("instances_fv")
+ self.instances_trans = Parameter(
+ torch.zeros(self.num_frames, self.num_instances, 3, device=self.device)
+ )
+ self.instances_quats = Parameter(
+ torch.zeros(self.num_frames, self.num_instances, 1, 4, device=self.device)
+ )
+ self.smpl_qauts = Parameter(
+ torch.zeros(self.num_frames, self.num_instances, 23, 4, device=self.device)
+ )
+ self.template = SMPLTemplate(
+ smpl_model_path="smpl_models/SMPL_NEUTRAL.pkl",
+ num_human=self.num_instances,
+ init_beta=torch.zeros(self.num_instances, 10, device=self.device),
+ cano_pose_type="da_pose",
+ use_voxel_deformer=self.use_voxel_deformer,
+ is_resume=True
+ ).to(self.device)
+ if self.use_voxel_deformer:
+ self.template.voxel_deformer.enable_voxel_correction()
+ msg = VanillaGaussians.load_state_dict(self, state_dict, **kwargs)
+ return msg
+
+ def get_instance_activated_gs_dict(self, ins_id: int) -> Dict[str, torch.Tensor]:
+ pts_mask = self.point_ids[..., 0] == ins_id
+ if pts_mask.sum() < 100:
+ return None
+ local_means = self._means[pts_mask]
+ activated_opacities = torch.sigmoid(self._opacities[pts_mask])
+ activated_scales = torch.exp(self._scales[pts_mask].repeat(1, 3) if self.ball_gaussians else self._scales[pts_mask])
+ activated_local_rotations = self.quat_act(self._quats[pts_mask])
+ gaussian_dict = {
+ "means": local_means,
+ "opacities": activated_opacities,
+ "scales": activated_scales,
+ "quats": activated_local_rotations,
+ "sh_dcs": self._features_dc[pts_mask],
+ "sh_rests": self._features_rest[pts_mask],
+ "ids": self.point_ids[pts_mask],
+ }
+ return gaussian_dict
+
+ def deform_gaussian_points(
+ self, gaussian_dict: Dict[str, torch.Tensor], cur_normalized_time: float,
+ ) -> Dict[str, torch.Tensor]:
+ """
+ deform the points
+ """
+ means = gaussian_dict["means"]
+ cur_frame = torch.argmin(
+ torch.abs(self.normalized_timestamps - cur_normalized_time)
+ )
+ ins_id = gaussian_dict["ids"].flatten()[0]
+ if not self.instances_fv[cur_frame, ins_id]:
+ # find the nearest frame that has the instance
+ for i in range(1, self.num_frames):
+ if cur_frame - i >= 0:
+ if self.instances_fv[cur_frame-i, ins_id]:
+ cur_frame = cur_frame - i
+ break
+ if cur_frame + i < self.num_frames:
+ if self.instances_fv[cur_frame+i, ins_id]:
+ cur_frame = cur_frame + i
+ break
+ instance_mask = torch.zeros(self.num_instances, device=self.device)
+ instance_mask[ins_id] = 1
+ instance_mask = instance_mask.bool()
+ masked_theta = torch.cat(
+ [torch.tensor([1.0, 1.0, 0.0, 0.0], device=self.device).unsqueeze(0),
+ self.smpl_qauts[cur_frame, ins_id]], dim=0
+ ).unsqueeze(0)
+
+ W, A = self.template(
+ masked_theta = self.quat_act(masked_theta),
+ instances_mask = instance_mask,
+ xyz_canonical = means.reshape(1, self.smpl_points_num, 3).repeat(self.num_instances, 1, 1) if self.use_voxel_deformer else None
+ )
+ T = torch.einsum("bnj, bjrc -> bnrc", W, A)
+ R = T[..., :3, :3].squeeze() # [N, 3, 3]
+ t = T[..., :3, 3].squeeze() # [N, 3]
+
+ deformed_means = torch.einsum(
+ "nij,nj->ni", R, means
+ ) + t # [6890, 3]
+ gaussian_dict["means"] = deformed_means
+
+ # placeholder for quats rotating: TODO
+ gaussian_dict["quats"] = gaussian_dict["quats"]
+ return gaussian_dict
+
+ def collect_gaussians_from_ids(self, ids: List[int]) -> Dict:
+ gaussian_dict = {}
+ for id in ids:
+ if id not in gaussian_dict:
+ instance_raw_dict = {
+ "_means": self._means[self.point_ids[..., 0] == id],
+ "_scales": self._scales[self.point_ids[..., 0] == id],
+ "_quats": self._quats[self.point_ids[..., 0] == id],
+ "_features_dc": self._features_dc[self.point_ids[..., 0] == id],
+ "_features_rest": self._features_rest[self.point_ids[..., 0] == id],
+ "_opacities": self._opacities[self.point_ids[..., 0] == id],
+ "point_ids": self.point_ids[self.point_ids[..., 0] == id],
+ "instances_fv": self.instances_fv[:, id],
+ "instances_quats": self.instances_quats.data[:, id],
+ "smpl_qauts": self.smpl_qauts.data[:, id],
+ "smpl_template": {
+ "J_canonical": self.template.J_canonical[id],
+ "W": self.template.W[id],
+ "voxel_deformer": {
+ "lbs_voxel_base": self.template.voxel_deformer.lbs_voxel_base[id],
+ "voxel_w_correction": self.template.voxel_deformer.voxel_w_correction[id],
+ "offset": self.template.voxel_deformer.offset[id],
+ "scale": self.template.voxel_deformer.scale[id],
+ }
+ }
+ }
+ gaussian_dict[id] = instance_raw_dict
+ return gaussian_dict
+
+ def remove_instances(self, remove_id_list: List[int]) -> None:
+ """
+ remove instances from the model
+
+ Args:
+ remove_id_list: list of instance ids to be removed
+ """
+ for ins_ids in remove_id_list:
+ mask = ~(self.point_ids[..., 0] == ins_ids)
+ self._means = Parameter(self._means[mask])
+ self._scales = Parameter(self._scales[mask])
+ self._quats = Parameter(self._quats[mask])
+ self._features_dc = Parameter(self._features_dc[mask])
+ self._features_rest = Parameter(self._features_rest[mask])
+ self._opacities = Parameter(self._opacities[mask])
+ self.point_ids = self.point_ids[mask]
+ ins_mask = torch.ones(self.num_instances, device=self.device, dtype=torch.bool)
+ ins_mask[ins_ids] = 0
+ self.instances_fv = self.instances_fv[:, ins_mask]
+ self.instances_quats = Parameter(self.instances_quats[:, ins_mask])
+ self.smpl_qauts = Parameter(self.smpl_qauts[:, ins_mask])
+ # process_smpl_template
+ self.template.remove_instance(ins_ids)
+
+ def replace_instances(self, replace_dict: Dict[int, int]) -> None:
+ """
+ replace instances from the model
+
+ Args:
+ replace_dict: {
+ ins_id(to be replaced): ins_id(replace with)
+ ...
+ }
+ """
+ new_gaussians_dict = self.collect_gaussians_from_ids(replace_dict.values())
+ for ins_id, new_id in replace_dict.items():
+ # self.remove_instances([ins_id])
+ new_gaussian = new_gaussians_dict[new_id]
+ # self._means = Parameter(torch.cat([self._means, new_gaussian["_means"]], dim=0))
+ # self._scales = Parameter(torch.cat([self._scales, new_gaussian["_scales"]], dim=0))
+ # self._quats = Parameter(torch.cat([self._quats, new_gaussian["_quats"]], dim=0))
+ # self._features_dc = Parameter(torch.cat([self._features_dc, new_gaussian["_features_dc"]], dim=0))
+ # self._features_rest = Parameter(torch.cat([self._features_rest, new_gaussian["_features_rest"]], dim=0))
+ # self._opacities = Parameter(torch.cat([self._opacities, new_gaussian["_opacities"]], dim=0))
+ self._features_dc.data[self.point_ids[..., 0] == ins_id] = new_gaussian["_features_dc"]
+ # keeps original point ids
+ # self.point_ids = torch.cat([self.point_ids, torch.full_like(new_gaussian["point_ids"], ins_id)], dim=0)
+ # self.template.add_instance(ins_id, new_gaussian["smpl_template"])
+ # self.instances_fv = torch.cat([self.instances_fv, new_gaussian["instances_fv"].unsqueeze(1)], dim=1)
+ # self.instances_quats = Parameter(torch.cat([self.instances_quats, new_gaussian["instances_quats"].unsqueeze(1)], dim=1))
+ # self.smpl_qauts = Parameter(torch.cat([self.smpl_qauts, new_gaussian["smpl_qauts"].unsqueeze(1)], dim=1))
+
+ def export_gaussians_to_ply(
+ self, alpha_thresh: float, instance_id: List[int] = None, specific_frame: int = 0,
+ ) -> Dict[str, torch.Tensor]:
+ self.cur_frame = specific_frame
+ cur_frame = 10
+ ins_id = instance_id
+
+ means = self._means[self.point_ids[..., 0] == ins_id]
+
+ if not self.instances_fv[cur_frame, ins_id]:
+ # find the nearest frame that has the instance
+ for i in range(1, self.num_frames):
+ if cur_frame - i >= 0:
+ if self.instances_fv[cur_frame-i, ins_id]:
+ cur_frame = cur_frame - i
+ break
+ if cur_frame + i < self.num_frames:
+ if self.instances_fv[cur_frame+i, ins_id]:
+ cur_frame = cur_frame + i
+ break
+ instance_mask = torch.zeros(self.num_instances, device=self.device)
+ instance_mask[ins_id] = 1
+ instance_mask = instance_mask.bool()
+ masked_theta = torch.cat(
+ [torch.tensor([1.0, 1.0, 0.0, 0.0], device=self.device).unsqueeze(0),
+ self.smpl_qauts[cur_frame, ins_id]], dim=0
+ ).unsqueeze(0)
+
+ W, A = self.template(
+ masked_theta = self.quat_act(masked_theta),
+ instances_mask = instance_mask,
+ xyz_canonical = means.reshape(1, self.smpl_points_num, 3).repeat(self.num_instances, 1, 1) if self.use_voxel_deformer else None
+ )
+ T = torch.einsum("bnj, bjrc -> bnrc", W, A)
+ R = T[..., :3, :3].squeeze() # [N, 3, 3]
+ t = T[..., :3, 3].squeeze() # [N, 3]
+
+ deformed_means = torch.einsum(
+ "nij,nj->ni", R, means
+ ) + t # [6890, 3]
+
+ color = self.colors[self.point_ids[..., 0] == ins_id]
+ return {
+ "positions": deformed_means,
+ "colors": color,
+ }
diff --git a/worldbench/third_party/drivestudio/models/trainers/__init__.py b/worldbench/third_party/drivestudio/models/trainers/__init__.py
new file mode 100644
index 0000000..a3d2adc
--- /dev/null
+++ b/worldbench/third_party/drivestudio/models/trainers/__init__.py
@@ -0,0 +1,3 @@
+from .base import BasicTrainer
+from .single import SingleTrainer
+from .scene_graph import MultiTrainer
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/models/trainers/base.py b/worldbench/third_party/drivestudio/models/trainers/base.py
new file mode 100644
index 0000000..48ceaa8
--- /dev/null
+++ b/worldbench/third_party/drivestudio/models/trainers/base.py
@@ -0,0 +1,788 @@
+from typing import Dict, List, Tuple
+from omegaconf import OmegaConf
+import os
+import time
+import logging
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+import kornia
+from enum import IntEnum
+import viser
+import nerfview
+from pytorch_msssim import SSIM
+from torchmetrics.image import PeakSignalNoiseRatio
+from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+
+from models.gaussians.basics import *
+
+logger = logging.getLogger()
+
+class GSModelType(IntEnum):
+ Background = 0
+ RigidNodes = 1
+ SMPLNodes = 2
+ DeformableNodes = 3
+
+def lr_scheduler_fn(
+ cfg: OmegaConf,
+ lr_init: float
+):
+ if cfg.lr_final is None:
+ lr_final = lr_init
+ else:
+ lr_final = cfg.lr_final
+
+ def func(step):
+ step = step - cfg.opt_after
+ if step < 0:
+ return 0.
+
+ if step < cfg.warmup_steps:
+ if cfg.ramp == "cosine":
+ lr = cfg.lr_pre_warmup + (lr_init - cfg.lr_pre_warmup) * np.sin(
+ 0.5 * np.pi * np.clip(step / cfg.warmup_steps, 0, 1)
+ )
+ else:
+ lr = (
+ cfg.lr_pre_warmup
+ + (lr_init - cfg.lr_pre_warmup) * step / cfg.warmup_steps
+ )
+ else:
+ t = np.clip(
+ (step - cfg.warmup_steps) / (cfg.max_steps - cfg.warmup_steps), 0, 1
+ )
+ lr = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t)
+ return lr # divided by lr_init because the multiplier is with the initial learning rate
+
+ return func
+
+class BasicTrainer(nn.Module):
+ def __init__(
+ self,
+ type: str = "basic",
+ optim: OmegaConf = None,
+ losses: OmegaConf = None,
+ render: OmegaConf = None,
+ res_schedule: OmegaConf = None,
+ gaussian_optim_general_cfg: OmegaConf = None,
+ gaussian_ctrl_general_cfg: OmegaConf = None,
+ model_config: OmegaConf = None,
+ num_train_images: int = 0,
+ num_full_images: int = 0,
+ test_set_indices: List[int] = None,
+ scene_aabb: torch.Tensor = None,
+ device=None,
+ ):
+ super().__init__()
+ self._type = type
+ self.optim_general = optim
+ self.losses_dict = losses
+ self.render_cfg = render
+ self.res_schedule = res_schedule
+ self.model_config = model_config
+ self.num_iters = self.optim_general.get("num_iters", 30000)
+ self.gaussian_optim_general_cfg = gaussian_optim_general_cfg
+ self.gaussian_ctrl_general_cfg = gaussian_ctrl_general_cfg
+ self.step = 0
+ self.device = device
+
+ # dataset infos
+ self.num_train_images = num_train_images
+ self.num_full_images = num_full_images
+
+ # init scene scale
+ self._init_scene(scene_aabb=scene_aabb)
+
+ # init models
+ self.models = {}
+ self.misc_classes_keys = [
+ 'Sky', 'Affine', 'CamPose', 'CamPosePerturb'
+ ]
+ self.gaussian_classes = {}
+ self._init_models()
+ self.pts_labels = None # will be overwritten in forward
+ self.render_dynamic_mask = False
+
+ # init losses fn
+ self._init_losses()
+
+ # metrics
+ self.psnr = PeakSignalNoiseRatio(data_range=1.0).to(self.device)
+ self.ssim = SSIM(data_range=1.0, size_average=True, channel=3).to(self.device)
+ self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True).to(self.device)
+ self.step = 0
+
+ # background color
+ self.back_color = torch.zeros(3).to(self.device)
+
+ # for evaluation
+ self.cur_frame = torch.tensor(0, device=self.device)
+ self.test_set_indices = test_set_indices # will be override
+
+ # a simple viewer for background visualization
+ self.viewer = None
+
+ @property
+ def in_test_set(self):
+ return self.cur_frame.item() in self.test_set_indices
+
+ def set_train(self):
+ for model in self.models.values():
+ model.train()
+ self.train()
+
+ def set_eval(self):
+ for model in self.models.values():
+ model.eval()
+ self.eval()
+
+ def _get_downscale_factor(self):
+ if self.training:
+ return 2 ** max((self.res_schedule.downscale_times - self.step // self.res_schedule.double_steps), 0)
+ else:
+ return 1
+
+ def update_gaussian_cfg(self, model_cfg: OmegaConf) -> OmegaConf:
+ class_optim_cfg = model_cfg.get('optim', None)
+ class_ctrl_cfg = model_cfg.get('ctrl', None)
+ new_optim_cfg = self.gaussian_optim_general_cfg.copy()
+ new_ctrl_cfg = self.gaussian_ctrl_general_cfg.copy()
+ if class_optim_cfg is not None:
+ new_optim_cfg.update(class_optim_cfg)
+ if class_ctrl_cfg is not None:
+ new_ctrl_cfg.update(class_ctrl_cfg)
+ model_cfg['optim'] = new_optim_cfg
+ model_cfg['ctrl'] = new_ctrl_cfg
+
+ return model_cfg
+
+ def _init_scene(self, scene_aabb) -> None:
+ self.aabb = scene_aabb.to(self.device)
+ scene_origin = (self.aabb[0] + self.aabb[1]) / 2
+ scene_radius = torch.max(self.aabb[1] - self.aabb[0]) / 2 * 1.1
+ self.scene_radius = scene_radius.item()
+ self.scene_origin = scene_origin
+ logger.info(f"scene origin: {scene_origin}")
+ logger.info(f"scene radius: {scene_radius}")
+
+ def _init_models(self) -> None:
+ raise NotImplementedError("Please implement the _init_models function")
+
+ def initialize_optimizer(self) -> None:
+ # get param groups first
+ self.param_groups = {}
+ for class_name, model in self.models.items():
+ self.param_groups.update(model.get_param_groups())
+
+ groups = []
+ lr_schedulers = {}
+ for params_name, params in self.param_groups.items():
+ class_name = params_name.split("#")[0]
+ component_name = params_name.split("#")[1]
+ class_cfg = self.model_config.get(class_name)
+ class_optim_cfg = class_cfg["optim"]
+
+ raw_optim_cfg = class_optim_cfg.get(component_name, None)
+ lr_scale_factor = raw_optim_cfg.get("scale_factor", 1.0)
+ if isinstance(lr_scale_factor, str) and lr_scale_factor == "scene_radius":
+ # scale the spatial learning rate to scene scale
+ lr_scale_factor = self.scene_radius
+
+ optim_cfg = OmegaConf.create({
+ "lr": raw_optim_cfg.get('lr', 0.0005),
+ "eps": raw_optim_cfg.get('eps', 1.0e-15),
+ "weight_decay": raw_optim_cfg.get('weight_decay', 0),
+ })
+ optim_cfg.lr = optim_cfg.lr * lr_scale_factor
+ assert optim_cfg is not None, f"param group {params_name} not found in config"
+ lr_init = optim_cfg.lr
+ groups.append({
+ 'params': params,
+ 'name': params_name,
+ 'lr': optim_cfg.lr,
+ 'eps': optim_cfg.eps,
+ 'weight_decay': optim_cfg.weight_decay
+ })
+
+ if raw_optim_cfg.get("lr_final", None) is not None:
+ sched_cfg = OmegaConf.create({
+ "opt_after": raw_optim_cfg.get('opt_after', 0),
+ "warmup_steps": raw_optim_cfg.get('warmup_steps', 0),
+ "max_steps": raw_optim_cfg.get('max_steps', self.num_iters),
+ "lr_pre_warmup": raw_optim_cfg.get('lr_pre_warmup', 1.0e-8),
+ "lr_final": raw_optim_cfg.get('lr_final', None),
+ "ramp": raw_optim_cfg.get('ramp', "cosine"),
+ })
+ # scale the learning rate according to the scene scale
+ sched_cfg.lr_pre_warmup = sched_cfg.lr_pre_warmup * lr_scale_factor
+ sched_cfg.lr_final = sched_cfg.lr_final * lr_scale_factor if sched_cfg.lr_final is not None else None
+ # adjust max_steps to account for opt_after
+ sched_cfg.max_steps = sched_cfg.max_steps - sched_cfg.opt_after
+ lr_schedulers[params_name] = lr_scheduler_fn(sched_cfg, lr_init)
+
+ self.optimizer = torch.optim.Adam(groups, lr=0.0, eps=1e-15)
+ self.lr_schedulers = lr_schedulers
+ self.grad_scaler = torch.cuda.amp.GradScaler(enabled=self.optim_general.get("use_grad_scaler", False))
+
+ def _init_losses(self) -> None:
+ sky_opacity_loss_fn = None
+ if "Sky" in self.models:
+ if self.losses_dict.mask.opacity_loss_type == "bce":
+ from models.losses import binary_cross_entropy
+ sky_opacity_loss_fn = lambda pred, gt: binary_cross_entropy(pred, gt, reduction="mean")
+ elif self.losses_dict.mask.opacity_loss_type == "safe_bce":
+ from models.losses import safe_binary_cross_entropy
+ sky_opacity_loss_fn = lambda pred, gt: safe_binary_cross_entropy(pred, gt, limit=0.1, reduction="mean")
+ self.sky_opacity_loss_fn = sky_opacity_loss_fn
+
+ depth_loss_fn = None
+ depth_loss_cfg = self.losses_dict.get("depth", None)
+ if depth_loss_cfg is not None:
+ from models.losses import DepthLoss
+ depth_loss_fn = DepthLoss(
+ loss_type=depth_loss_cfg.loss_type,
+ normalize=depth_loss_cfg.normalize,
+ use_inverse_depth=depth_loss_cfg.inverse_depth,
+ )
+ self.depth_loss_fn = depth_loss_fn
+
+ def optimizer_zero_grad(self) -> None:
+ self.optimizer.zero_grad()
+
+ def optimizer_step(self) -> None:
+ # for params_name, optimizer in self.optimizers.items():
+ # class_name = params_name.split("#")[0]
+ # component_name = params_name.split("#")[1]
+ # max_norm = self.model_config[class_name]["optim"][component_name].get("max_norm", None)
+ # if max_norm is not None:
+ # self.grad_scaler.unscale_(optimizer)
+ # torch.nn.utils.clip_grad_norm_(self.param_groups[params_name], max_norm)
+ # if any(any(p.grad is not None for p in g["params"]) for g in optimizer.param_groups):
+ # self.grad_scaler.step(optimizer)
+ self.optimizer.step()
+
+ def preprocess_per_train_step(self, step: int) -> None:
+ self.step = step
+ for class_name in self.gaussian_classes.keys():
+ self.models[class_name].preprocess_per_train_step(step)
+
+ # viewer
+ if self.viewer is not None:
+ while self.viewer.state.status == "paused":
+ time.sleep(0.01)
+ self.viewer.lock.acquire()
+ self.tic = time.time()
+
+ def postprocess_per_train_step(self, step: int) -> None:
+ radii = self.info["radii"]
+ if self.render_cfg.absgrad:
+ grads = self.info["means2d"].absgrad.clone()
+ else:
+ grads = self.info["means2d"].grad.clone()
+ grads[..., 0] *= self.info["width"] / 2.0 * self.render_cfg.batch_size
+ grads[..., 1] *= self.info["height"] / 2.0 * self.render_cfg.batch_size
+
+ for class_name in self.gaussian_classes.keys():
+ gaussian_mask = self.pts_labels == self.gaussian_classes[class_name]
+
+ self.models[class_name].postprocess_per_train_step(
+ step=step,
+ optimizer=self.optimizer,
+ radii=radii[0, gaussian_mask],
+ xys_grad=grads[0, gaussian_mask],
+ last_size=max(self.info["width"], self.info["height"])
+ )
+
+ # viewer
+ if self.viewer is not None:
+ num_train_rays_per_step = self.render_cfg.batch_size * self.info["width"] * self.info["height"]
+ self.viewer.lock.release()
+ num_train_steps_per_sec = 1.0 / (time.time() - self.tic)
+ num_train_rays_per_sec = (
+ num_train_rays_per_step * num_train_steps_per_sec
+ )
+ # Update the viewer state.
+ self.viewer.state.num_train_rays_per_sec = num_train_rays_per_sec
+ # Update the scene.
+ self.viewer.update(step, num_train_rays_per_step)
+
+ def update_visibility_filter(self) -> None:
+ for class_name in self.gaussian_classes.keys():
+ gaussian_mask = self.pts_labels == self.gaussian_classes[class_name]
+ self.models[class_name].cur_radii = self.info["radii"][0, gaussian_mask]
+
+ def process_camera(
+ self,
+ camera_infos: Dict[str, torch.Tensor],
+ image_ids: torch.Tensor,
+ novel_view: bool = False
+ ) -> dataclass_camera:
+ camtoworlds = camtoworlds_gt = camera_infos["camera_to_world"]
+
+ if "CamPosePerturb" in self.models.keys() and not novel_view:
+ camtoworlds = self.models["CamPosePerturb"](camtoworlds, image_ids)
+
+ if "CamPose" in self.models.keys() and not novel_view:
+ camtoworlds = self.models["CamPose"](camtoworlds, image_ids)
+
+ # collect camera information
+ camera_dict = dataclass_camera(
+ camtoworlds=camtoworlds,
+ camtoworlds_gt=camtoworlds_gt,
+ Ks=camera_infos["intrinsics"],
+ H=camera_infos["height"],
+ W=camera_infos["width"]
+ )
+
+ return camera_dict
+
+ def collect_gaussians(
+ self,
+ cam: dataclass_camera,
+ image_ids: torch.Tensor # leave it here for future use
+ ) -> dataclass_gs:
+ gs_dict = {
+ "_means": [],
+ "_scales": [],
+ "_quats": [],
+ "_rgbs": [],
+ "_opacities": [],
+ "class_labels": [],
+ }
+ for class_name in self.gaussian_classes.keys():
+ gs = self.models[class_name].get_gaussians(cam)
+ if gs is None:
+ continue
+
+ # collect gaussians
+ gs["class_labels"] = torch.full((gs["_means"].shape[0],), self.gaussian_classes[class_name], device=self.device)
+ for k, _ in gs.items():
+ gs_dict[k].append(gs[k])
+
+ for k, v in gs_dict.items():
+ gs_dict[k] = torch.cat(v, dim=0)
+
+ # get the class labels
+ self.pts_labels = gs_dict.pop("class_labels")
+ if self.render_dynamic_mask:
+ self.dynamic_pts_mask = (self.pts_labels != 0).float()
+
+ gaussians = dataclass_gs(
+ _means=gs_dict["_means"],
+ _scales=gs_dict["_scales"],
+ _quats=gs_dict["_quats"],
+ _rgbs=gs_dict["_rgbs"],
+ _opacities=gs_dict["_opacities"],
+ detach_keys=[], # if "means" in detach_keys, then the means will be detached
+ extras=None # to save some extra information (TODO) more flexible way
+ )
+
+ return gaussians
+
+ def render_gaussians(
+ self,
+ gs: dataclass_gs,
+ cam: dataclass_camera,
+ **kwargs,
+ ) -> Dict[str, torch.Tensor]:
+
+ def render_fn(opaticy_mask=None, return_info=False):
+ renders, alphas, info = rasterization(
+ means=gs.means,
+ quats=gs.quats,
+ scales=gs.scales,
+ opacities=gs.opacities.squeeze()*opaticy_mask if opaticy_mask is not None else gs.opacities.squeeze(),
+ colors=gs.rgbs,
+ viewmats=torch.linalg.inv(cam.camtoworlds)[None, ...], # [C, 4, 4]
+ Ks=cam.Ks[None, ...], # [C, 3, 3]
+ width=cam.W,
+ height=cam.H,
+ packed=self.render_cfg.packed,
+ absgrad=self.render_cfg.absgrad,
+ sparse_grad=self.render_cfg.sparse_grad,
+ rasterize_mode="antialiased" if self.render_cfg.antialiased else "classic",
+ **kwargs,
+ )
+ renders = renders[0]
+ alphas = alphas[0].squeeze(-1)
+ assert self.render_cfg.batch_size == 1, "batch size must be 1, will support batch size > 1 in the future"
+
+ assert renders.shape[-1] == 4, f"Must render rgb, depth and alpha"
+ rendered_rgb, rendered_depth = torch.split(renders, [3, 1], dim=-1)
+
+ if not return_info:
+ return torch.clamp(rendered_rgb, max=1.0), rendered_depth, alphas[..., None]
+ else:
+ return torch.clamp(rendered_rgb, max=1.0), rendered_depth, alphas[..., None], info
+
+ # render rgb and opacity
+ rgb, depth, opacity, self.info = render_fn(return_info=True)
+ results = {
+ "rgb_gaussians": rgb,
+ "depth": depth,
+ "opacity": opacity
+ }
+
+ if self.training:
+ self.info["means2d"].retain_grad()
+
+ return results, render_fn
+
+ def affine_transformation(
+ self,
+ rgb_blended: torch.Tensor,
+ image_infos: Dict[str, torch.Tensor]
+ ):
+ if "Affine" in self.models:
+ affine_trs = self.models['Affine'](image_infos)
+ rgb_transformed = (affine_trs[..., :3, :3] @ rgb_blended[..., None] + affine_trs[..., :3, 3:])[..., 0]
+
+ return rgb_transformed
+ else:
+ return rgb_blended
+
+ def forward(
+ self,
+ image_infos: Dict[str, torch.Tensor],
+ camera_infos: Dict[str, torch.Tensor],
+ novel_view: bool = False
+ ) -> Dict[str, torch.Tensor]:
+ """Forward pass of the model
+
+ Args:
+ image_infos (Dict[str, torch.Tensor]): image and pixels information
+ camera_infos (Dict[str, torch.Tensor]): camera information
+ novel_view: whether the view is novel, if True, disable the camera refinement
+
+ Returns:
+ Dict[str, torch.Tensor]: output of the model
+ """
+
+ # for evaluation
+ for model in self.models.values():
+ if hasattr(model, 'in_test_set'):
+ model.in_test_set = self.in_test_set
+
+ # prapare data
+ processed_cam = self.process_camera(
+ camera_infos=camera_infos,
+ image_ids=image_infos["img_idx"].flatten()[0],
+ novel_view=novel_view
+ )
+ gs = self.collect_gaussians(
+ cam=processed_cam,
+ image_ids=image_infos["img_idx"].flatten()[0]
+ )
+
+ # render gaussians
+ outputs, _ = self.render_gaussians(
+ gs=gs,
+ cam=processed_cam,
+ near_plane=self.render_cfg.near_plane,
+ far_plane=self.render_cfg.far_plane,
+ render_mode="RGB+ED",
+ radius_clip=self.render_cfg.get('radius_clip', 0.)
+ )
+
+ # render sky
+ sky_model = self.models['Sky']
+ outputs["rgb_sky"] = sky_model(image_infos)
+ outputs["rgb_sky_blend"] = outputs["rgb_sky"] * (1.0 - outputs["opacity"])
+
+ # affine transformation
+ outputs["rgb"] = self.affine_transformation(
+ outputs["rgb_gaussians"] + outputs["rgb_sky"] * (1.0 - outputs["opacity"]), image_infos
+ )
+
+ return outputs
+
+ def backward(self, loss_dict: Dict[str, torch.Tensor]) -> None:
+ # ----------------- backward ----------------
+ total_loss = sum(loss for loss in loss_dict.values())
+ self.grad_scaler.scale(total_loss).backward()
+ self.optimizer_step()
+
+ scale = self.grad_scaler.get_scale()
+ self.grad_scaler.update()
+
+ # If the gradient scaler is decreased, no optimization step is performed so we should not step the scheduler.
+ if scale <= self.grad_scaler.get_scale():
+ for group in self.optimizer.param_groups:
+ if group["name"] in self.lr_schedulers:
+ new_lr = self.lr_schedulers[group["name"]](self.step)
+ group["lr"] = new_lr
+
+ def compute_losses(
+ self,
+ outputs: Dict[str, torch.Tensor],
+ image_infos: Dict[str, torch.Tensor],
+ cam_infos: Dict[str, torch.Tensor]
+ ) -> Dict[str, torch.Tensor]:
+ # calculate loss
+ loss_dict = {}
+
+ if "egocar_masks" in image_infos:
+ # in the case of egocar, we need to mask out the egocar region
+ valid_loss_mask = (1.0 - image_infos["egocar_masks"]).float()
+ else:
+ valid_loss_mask = torch.ones_like(image_infos["sky_masks"])
+
+ gt_rgb = image_infos["pixels"] * valid_loss_mask[..., None]
+ predicted_rgb = outputs["rgb"] * valid_loss_mask[..., None]
+
+ gt_occupied_mask = (1.0 - image_infos["sky_masks"]).float() * valid_loss_mask
+ pred_occupied_mask = outputs["opacity"].squeeze() * valid_loss_mask
+
+ # rgb loss
+ Ll1 = torch.abs(gt_rgb - predicted_rgb).mean()
+ simloss = 1 - self.ssim(gt_rgb.permute(2, 0, 1)[None, ...], predicted_rgb.permute(2, 0, 1)[None, ...])
+ loss_dict.update({
+ "rgb_loss": self.losses_dict.rgb.w * Ll1,
+ "ssim_loss": self.losses_dict.ssim.w * simloss,
+ })
+
+ # mask loss
+ if self.sky_opacity_loss_fn is not None:
+ sky_loss_opacity = self.sky_opacity_loss_fn(pred_occupied_mask, gt_occupied_mask) * self.losses_dict.mask.w
+ loss_dict.update({"sky_loss_opacity": sky_loss_opacity})
+
+ # depth loss
+ if self.depth_loss_fn is not None:
+ gt_depth = image_infos["lidar_depth_map"]
+ lidar_hit_mask = (gt_depth > 0).float() * valid_loss_mask
+ pred_depth = outputs["depth"]
+ depth_loss = self.depth_loss_fn(pred_depth, gt_depth, lidar_hit_mask)
+
+ lidar_w_decay = self.losses_dict.depth.get("lidar_w_decay", -1)
+ if lidar_w_decay > 0:
+ decay_weight = np.exp(-self.step / 8000 * lidar_w_decay)
+ else:
+ decay_weight = 1
+ depth_loss = depth_loss * self.losses_dict.depth.w * decay_weight
+ loss_dict.update({"depth_loss": depth_loss})
+
+ # ----- reg loss -----
+ opacity_entropy_reg = self.losses_dict.get("opacity_entropy", None)
+ if opacity_entropy_reg is not None:
+ pred_opacity = torch.clamp(outputs["opacity"].squeeze(), 1e-6, 1 - 1e-6)
+ loss_dict.update({
+ "opacity_entropy_loss": opacity_entropy_reg.w * (-pred_opacity * torch.log(pred_opacity)).mean()
+ })
+
+ # from pvg: https://github.com/fudan-zvg/PVG/blob/b4162a9135282e0f3c929054f16be1b3fbacd77a/train.py#L161
+ inverse_depth_smoothness_reg = self.losses_dict.get("inverse_depth_smoothness", None)
+ if inverse_depth_smoothness_reg is not None:
+ inverse_depth = 1 / (outputs["depth"] + 1e-5)
+ loss_inv_depth = kornia.losses.inverse_depth_smoothness_loss(
+ inverse_depth[None].repeat(1, 1, 1, 3).permute(0, 3, 1, 2),
+ image_infos["pixels"][None].permute(0, 3, 1, 2)
+ )
+ loss_dict.update({
+ "inverse_depth_smoothness_loss": inverse_depth_smoothness_reg.w * loss_inv_depth
+ })
+
+ # affine reg loss
+ affine_reg = self.losses_dict.get("affine", None)
+ if affine_reg is not None and "Affine" in self.models:
+ affine_trs = self.models['Affine']({"img_idx": image_infos["img_idx"].flatten()[0]})
+ reg_mat = torch.eye(3, device=self.device)
+ reg_shift = torch.zeros(3, device=self.device)
+ loss_affine = torch.abs(affine_trs[..., :3, :3] - reg_mat).mean() + torch.abs(affine_trs[..., :3, 3:] - reg_shift).mean()
+ loss_dict.update({
+ "affine_loss": affine_reg.w * loss_affine
+ })
+
+ # dynamic region loss
+ dynamic_region_weighted_losses = self.losses_dict.get("dynamic_region", None)
+ if dynamic_region_weighted_losses is not None:
+ weight_factor = dynamic_region_weighted_losses.get("w", 1.0)
+ start_from = dynamic_region_weighted_losses.get("start_from", 0)
+ if self.step == start_from:
+ self.render_dynamic_mask = True
+ if self.step > start_from and "Dynamic_opacity" in outputs:
+ dynamic_pred_mask = (outputs["Dynamic_opacity"].data > 0.2).squeeze()
+ dynamic_pred_mask = dynamic_pred_mask & valid_loss_mask.bool()
+
+ if dynamic_pred_mask.sum() > 0:
+ Ll1 = torch.abs(gt_rgb[dynamic_pred_mask] - predicted_rgb[dynamic_pred_mask]).mean()
+ loss_dict.update({
+ "vehicle_region_rgb_loss": weight_factor * Ll1,
+ })
+
+ # compute gaussian reg loss
+ for class_name in self.gaussian_classes.keys():
+ class_reg_loss = self.models[class_name].compute_reg_loss()
+ for k, v in class_reg_loss.items():
+ loss_dict[f"{class_name}_{k}"] = v
+ return loss_dict
+
+ def compute_metrics(
+ self,
+ outputs: Dict[str, torch.Tensor],
+ image_infos: Dict[str, torch.Tensor]
+ ) -> Dict[str, torch.Tensor]:
+ metric_dict = {}
+ psnr = self.psnr(outputs["rgb"], image_infos["pixels"])
+ metric_dict.update({"psnr": psnr})
+ return metric_dict
+
+ def get_gaussian_count(self):
+ num_dict = {}
+ for class_name in self.gaussian_classes.keys():
+ num_dict[class_name] = self.models[class_name].num_points
+ return num_dict
+
+ def state_dict(self, only_model: bool = True):
+ state_dict = super().state_dict()
+ state_dict.update({
+ "models": {k: v.state_dict() for k, v in self.models.items()},
+ "step": self.step,
+ })
+ if not only_model:
+ state_dict.update({
+ "optimizer": {k: v.state_dict() for k, v in self.optimizer.items()},
+ # "lr_schedulers": {k: v.state_dict() for k, v in self.lr_schedulers.items()},
+ # "grad_scaler": self.grad_scaler.state_dict(),
+ })
+ return state_dict
+
+ def load_state_dict(self, state_dict: dict, load_only_model: bool =True, strict: bool = True):
+ step = state_dict.pop("step")
+ self.step = step
+ logger.info(f"Loading checkpoint at step {step}")
+
+ # load optimizer and schedulers
+ if "optimizer" in state_dict:
+ loaded_state_optimizers = state_dict.pop("optimizer")
+ # if "schedulers" in state_dict:
+ # loaded_state_schedulers = state_dict.pop("schedulers")
+ # if "grad_scaler" in state_dict:
+ # loaded_grad_scaler = state_dict.pop("grad_scaler")
+ if not load_only_model:
+ raise NotImplementedError("Now only support loading model, \
+ it seems there is no need to load optimizer and schedulers")
+ for k, v in loaded_state_optimizers.items():
+ self.optimizer[k].load_state_dict(v)
+ for k, v in loaded_state_schedulers.items():
+ self.schedulers[k].load_state_dict(v)
+ self.grad_scaler.load_state_dict(loaded_grad_scaler)
+
+ # load model
+ model_state_dict = state_dict.pop("models")
+ for class_name in self.models.keys():
+ model = self.models[class_name]
+ model.step = step
+ if class_name not in model_state_dict:
+ if class_name in self.gaussian_classes:
+ self.gaussian_classes.pop(class_name)
+ logger.warning(f"Cannot find {class_name} in the checkpoint")
+ continue
+ msg = model.load_state_dict(model_state_dict[class_name], strict=strict)
+ logger.info(f"{class_name}: {msg}")
+ msg = super().load_state_dict(state_dict, strict)
+ logger.info(f"BasicTrainer: {msg}")
+
+ def resume_from_checkpoint(
+ self,
+ ckpt_path: str,
+ load_only_model: bool=True
+ ) -> None:
+ """
+ Load model from checkpoint.
+ """
+ logger.info(f"Loading checkpoint from {ckpt_path}")
+ state_dict = torch.load(ckpt_path)
+ self.load_state_dict(state_dict, load_only_model=load_only_model, strict=True)
+
+ def save_checkpoint(
+ self,
+ log_dir: str,
+ save_only_model: bool=True,
+ is_final: bool=False
+ ) -> None:
+ """
+ Save model to checkpoint.
+ """
+ if is_final:
+ ckpt_path = os.path.join(log_dir, f"checkpoint_final.pth")
+ else:
+ ckpt_path = os.path.join(log_dir, f"checkpoint_{self.step:05d}.pth")
+ torch.save(self.state_dict(only_model=save_only_model), ckpt_path)
+ logger.info(f"Saved a checkpoint to {ckpt_path}")
+
+ def init_viewer(self, port: int = 8080):
+ # a simple viewer for background ONLY visualization
+ self.server = viser.ViserServer(port=port, verbose=False)
+ self.viewer = nerfview.Viewer(
+ server=self.server,
+ render_fn=self._viewer_render_fn,
+ mode="training",
+ )
+
+ @torch.no_grad()
+ def _viewer_render_fn(
+ self, camera_state: nerfview.CameraState, img_wh: Tuple[int, int]
+ ):
+ """Callable function for the viewer."""
+ W, H = img_wh
+ c2w = camera_state.c2w
+ K = camera_state.get_K(img_wh)
+ c2w = torch.from_numpy(c2w).float().to(self.device)
+ K = torch.from_numpy(K).float().to(self.device)
+
+ cam = dataclass_camera(
+ camtoworlds=c2w,
+ camtoworlds_gt=c2w,
+ Ks=K,
+ H=H,
+ W=W
+ )
+
+ gs_dict = {
+ "_means": [],
+ "_scales": [],
+ "_quats": [],
+ "_rgbs": [],
+ "_opacities": [],
+ }
+ for class_name in ["Background"]:
+ gs = self.models[class_name].get_gaussians(cam)
+ if gs is None:
+ continue
+
+ for k, _ in gs.items():
+ gs_dict[k].append(gs[k])
+
+ for k, v in gs_dict.items():
+ gs_dict[k] = torch.cat(v, dim=0)
+
+ gs = dataclass_gs(
+ _means=gs_dict["_means"],
+ _scales=gs_dict["_scales"],
+ _quats=gs_dict["_quats"],
+ _rgbs=gs_dict["_rgbs"],
+ _opacities=gs_dict["_opacities"],
+ detach_keys=[],
+ extras=None
+ )
+
+ render_colors, _, _ = rasterization(
+ means=gs.means,
+ quats=gs.quats,
+ scales=gs.scales,
+ opacities=gs.opacities.squeeze(),
+ colors=gs.rgbs,
+ viewmats=torch.linalg.inv(cam.camtoworlds)[None, ...], # [C, 4, 4]
+ Ks=cam.Ks[None, ...], # [C, 3, 3]
+ width=cam.W,
+ height=cam.H,
+ packed=self.render_cfg.packed,
+ absgrad=self.render_cfg.absgrad,
+ sparse_grad=self.render_cfg.sparse_grad,
+ rasterize_mode="antialiased" if self.render_cfg.antialiased else "classic",
+ radius_clip=4.0, # skip GSs that have small image radius (in pixels)
+ )
+ return render_colors[0].cpu().numpy()
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/models/trainers/scene_graph.py b/worldbench/third_party/drivestudio/models/trainers/scene_graph.py
new file mode 100644
index 0000000..c29df19
--- /dev/null
+++ b/worldbench/third_party/drivestudio/models/trainers/scene_graph.py
@@ -0,0 +1,296 @@
+from typing import Dict
+import torch
+import logging
+
+from datasets.driving_dataset import DrivingDataset
+from models.trainers.base import BasicTrainer, GSModelType
+from utils.misc import import_str
+from utils.geometry import uniform_sample_sphere
+
+logger = logging.getLogger()
+
+class MultiTrainer(BasicTrainer):
+ def __init__(
+ self,
+ num_timesteps: int,
+ **kwargs
+ ):
+ self.num_timesteps = num_timesteps
+ super().__init__(**kwargs)
+ self.render_each_class = True
+
+ def register_normalized_timestamps(self, num_timestamps: int):
+ self.normalized_timestamps = torch.linspace(0, 1, num_timestamps, device=self.device)
+
+ def _init_models(self):
+ # gaussian model classes
+ if "Background" in self.model_config:
+ self.gaussian_classes["Background"] = GSModelType.Background
+ if "RigidNodes" in self.model_config:
+ self.gaussian_classes["RigidNodes"] = GSModelType.RigidNodes
+ if "SMPLNodes" in self.model_config:
+ self.gaussian_classes["SMPLNodes"] = GSModelType.SMPLNodes
+ if "DeformableNodes" in self.model_config:
+ self.gaussian_classes["DeformableNodes"] = GSModelType.DeformableNodes
+
+ for class_name, model_cfg in self.model_config.items():
+ # update model config for gaussian classes
+ if class_name in self.gaussian_classes:
+ model_cfg = self.model_config.pop(class_name)
+ self.model_config[class_name] = self.update_gaussian_cfg(model_cfg)
+
+ if class_name in self.gaussian_classes.keys():
+ model = import_str(model_cfg.type)(
+ **model_cfg,
+ class_name=class_name,
+ scene_scale=self.scene_radius,
+ scene_origin=self.scene_origin,
+ num_train_images=self.num_train_images,
+ device=self.device
+ )
+
+ if class_name in self.misc_classes_keys:
+ model = import_str(model_cfg.type)(
+ class_name=class_name,
+ **model_cfg.get('params', {}),
+ n=self.num_full_images,
+ device=self.device
+ ).to(self.device)
+
+ self.models[class_name] = model
+
+ logger.info(f"Initialized models: {self.models.keys()}")
+
+ # register normalized timestamps
+ self.register_normalized_timestamps(self.num_timesteps)
+ for class_name in self.gaussian_classes.keys():
+ model = self.models[class_name]
+ if hasattr(model, 'register_normalized_timestamps'):
+ model.register_normalized_timestamps(self.normalized_timestamps)
+ if hasattr(model, 'set_bbox'):
+ model.set_bbox(self.aabb)
+
+ def safe_init_models(
+ self,
+ model: torch.nn.Module,
+ instance_pts_dict: Dict[str, Dict[str, torch.Tensor]]
+ ) -> None:
+ if len(instance_pts_dict.keys()) > 0:
+ model.create_from_pcd(
+ instance_pts_dict=instance_pts_dict
+ )
+ return False
+ else:
+ return True
+
+ def init_gaussians_from_dataset(
+ self,
+ dataset: DrivingDataset,
+ ) -> None:
+ # get instance points
+ rigidnode_pts_dict, deformnode_pts_dict, smplnode_pts_dict = {}, {}, {}
+ if "RigidNodes" in self.model_config:
+ rigidnode_pts_dict = dataset.get_init_objects(
+ cur_node_type='RigidNodes',
+ **self.model_config["RigidNodes"]["init"]
+ )
+
+ if "DeformableNodes" in self.model_config:
+ deformnode_pts_dict = dataset.get_init_objects(
+ cur_node_type='DeformableNodes',
+ exclude_smpl="SMPLNodes" in self.model_config,
+ **self.model_config["DeformableNodes"]["init"]
+ )
+
+ if "SMPLNodes" in self.model_config:
+ smplnode_pts_dict = dataset.get_init_smpl_objects(
+ **self.model_config["SMPLNodes"]["init"]
+ )
+ allnode_pts_dict = {**rigidnode_pts_dict, **deformnode_pts_dict, **smplnode_pts_dict}
+
+ # NOTE: Some gaussian classes may be empty (because no points for initialization)
+ # We will delete these classes from the model_config and models
+ empty_classes = []
+
+ # collect models
+ for class_name in self.gaussian_classes:
+ model_cfg = self.model_config[class_name]
+ model = self.models[class_name]
+
+ empty = False
+ if class_name == 'Background':
+ # ------ initialize gaussians ------
+ init_cfg = model_cfg.pop('init')
+ # sample points from the lidar point clouds
+ if init_cfg.get("from_lidar", None) is not None:
+ sampled_pts, sampled_color, sampled_time = dataset.get_lidar_samples(
+ **init_cfg.from_lidar, device=self.device
+ )
+ else:
+ sampled_pts, sampled_color, sampled_time = \
+ torch.empty(0, 3).to(self.device), torch.empty(0, 3).to(self.device), None
+
+ random_pts = []
+ num_near_pts = init_cfg.get('near_randoms', 0)
+ if num_near_pts > 0: # uniformly sample points inside the scene's sphere
+ num_near_pts *= 3 # since some invisible points will be filtered out
+ random_pts.append(uniform_sample_sphere(num_near_pts, self.device))
+ num_far_pts = init_cfg.get('far_randoms', 0)
+ if num_far_pts > 0: # inverse distances uniformly from (0, 1 / scene_radius)
+ num_far_pts *= 3
+ random_pts.append(uniform_sample_sphere(num_far_pts, self.device, inverse=True))
+
+ if num_near_pts + num_far_pts > 0:
+ random_pts = torch.cat(random_pts, dim=0)
+ random_pts = random_pts * self.scene_radius + self.scene_origin
+ visible_mask = dataset.check_pts_visibility(random_pts)
+ valid_pts = random_pts[visible_mask]
+
+ sampled_pts = torch.cat([sampled_pts, valid_pts], dim=0)
+ sampled_color = torch.cat([sampled_color, torch.rand(valid_pts.shape, ).to(self.device)], dim=0)
+
+ processed_init_pts = dataset.filter_pts_in_boxes(
+ seed_pts=sampled_pts,
+ seed_colors=sampled_color,
+ valid_instances_dict=allnode_pts_dict
+ )
+
+ model.create_from_pcd(
+ init_means=processed_init_pts["pts"], init_colors=processed_init_pts["colors"]
+ )
+
+ if class_name == 'RigidNodes':
+ empty = self.safe_init_models(
+ model=model,
+ instance_pts_dict=rigidnode_pts_dict
+ )
+
+ if class_name == 'DeformableNodes':
+ empty = self.safe_init_models(
+ model=model,
+ instance_pts_dict=deformnode_pts_dict
+ )
+
+ if class_name == 'SMPLNodes':
+ empty = self.safe_init_models(
+ model=model,
+ instance_pts_dict=smplnode_pts_dict
+ )
+
+ if empty:
+ empty_classes.append(class_name)
+ logger.warning(f"No points for {class_name} found, will remove the model")
+ else:
+ logger.info(f"Initialized {class_name} gaussians")
+
+ if len(empty_classes) > 0:
+ for class_name in empty_classes:
+ del self.models[class_name]
+ del self.model_config[class_name]
+ del self.gaussian_classes[class_name]
+ logger.warning(f"Model for {class_name} is removed")
+
+ logger.info(f"Initialized gaussians from pcd")
+
+ def forward(
+ self,
+ image_infos: Dict[str, torch.Tensor],
+ camera_infos: Dict[str, torch.Tensor],
+ novel_view: bool = False
+ ) -> Dict[str, torch.Tensor]:
+ """Forward pass of the model
+
+ Args:
+ image_infos (Dict[str, torch.Tensor]): image and pixels information
+ camera_infos (Dict[str, torch.Tensor]): camera information
+ novel_view: whether the view is novel, if True, disable the camera refinement
+
+ Returns:
+ Dict[str, torch.Tensor]: output of the model
+ """
+
+ # set current time or use temporal smoothing
+ normed_time = image_infos["normed_time"].flatten()[0]
+ self.cur_frame = torch.argmin(
+ torch.abs(self.normalized_timestamps - normed_time)
+ )
+
+ # for evaluation
+ for model in self.models.values():
+ if hasattr(model, 'in_test_set'):
+ model.in_test_set = self.in_test_set
+
+ # assigne current frame to gaussian models
+ for class_name in self.gaussian_classes.keys():
+ model = self.models[class_name]
+ if hasattr(model, 'set_cur_frame'):
+ model.set_cur_frame(self.cur_frame)
+
+ # prapare data
+ processed_cam = self.process_camera(
+ camera_infos=camera_infos,
+ image_ids=image_infos["img_idx"].flatten()[0],
+ novel_view=novel_view
+ )
+ gs = self.collect_gaussians(
+ cam=processed_cam,
+ image_ids=image_infos["img_idx"].flatten()[0]
+ )
+
+ # render gaussians
+ outputs, render_fn = self.render_gaussians(
+ gs=gs,
+ cam=processed_cam,
+ near_plane=self.render_cfg.near_plane,
+ far_plane=self.render_cfg.far_plane,
+ render_mode="RGB+ED",
+ radius_clip=self.render_cfg.get('radius_clip', 0.)
+ )
+
+ # render sky
+ sky_model = self.models['Sky']
+ outputs["rgb_sky"] = sky_model(image_infos)
+ outputs["rgb_sky_blend"] = outputs["rgb_sky"] * (1.0 - outputs["opacity"])
+
+ # affine transformation
+ outputs["rgb"] = self.affine_transformation(
+ outputs["rgb_gaussians"] + outputs["rgb_sky"] * (1.0 - outputs["opacity"]), image_infos
+ )
+
+ if not self.training and self.render_each_class:
+ with torch.no_grad():
+ for class_name in self.gaussian_classes.keys():
+ gaussian_mask = self.pts_labels == self.gaussian_classes[class_name]
+ sep_rgb, sep_depth, sep_opacity = render_fn(gaussian_mask)
+ outputs[class_name+"_rgb"] = self.affine_transformation(sep_rgb, image_infos)
+ outputs[class_name+"_opacity"] = sep_opacity
+ outputs[class_name+"_depth"] = sep_depth
+
+ if not self.training or self.render_dynamic_mask:
+ with torch.no_grad():
+ gaussian_mask = self.pts_labels != self.gaussian_classes["Background"]
+ sep_rgb, sep_depth, sep_opacity = render_fn(gaussian_mask)
+ outputs["Dynamic_rgb"] = self.affine_transformation(sep_rgb, image_infos)
+ outputs["Dynamic_opacity"] = sep_opacity
+ outputs["Dynamic_depth"] = sep_depth
+
+ return outputs
+
+ def compute_losses(
+ self,
+ outputs: Dict[str, torch.Tensor],
+ image_infos: Dict[str, torch.Tensor],
+ cam_infos: Dict[str, torch.Tensor],
+ ) -> Dict[str, torch.Tensor]:
+ loss_dict = super().compute_losses(outputs, image_infos, cam_infos)
+
+ return loss_dict
+
+ def compute_metrics(
+ self,
+ outputs: Dict[str, torch.Tensor],
+ image_infos: Dict[str, torch.Tensor],
+ ) -> Dict[str, torch.Tensor]:
+ metric_dict = super().compute_metrics(outputs, image_infos)
+
+ return metric_dict
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/models/trainers/single.py b/worldbench/third_party/drivestudio/models/trainers/single.py
new file mode 100644
index 0000000..a0b3609
--- /dev/null
+++ b/worldbench/third_party/drivestudio/models/trainers/single.py
@@ -0,0 +1,212 @@
+from typing import Dict
+
+import torch
+import logging
+
+from datasets.driving_dataset import DrivingDataset
+from models.trainers.base import BasicTrainer, GSModelType
+from models.gaussians import PeriodicVibrationGaussians
+from utils.misc import import_str
+from utils.geometry import uniform_sample_sphere
+
+"""
+Single Gaussian Model Trainer
+
+Methods Available:
+ Static Vanilla Gaussians,
+ Defromable Gaussians for dynamic scenes,
+ Periodic Vibration Gaussians for dynamic scenes,
+"""
+
+logger = logging.getLogger()
+
+class SingleTrainer(BasicTrainer):
+ def __init__(
+ self,
+ num_timesteps: int,
+ **kwargs
+ ):
+ self.num_timesteps = num_timesteps
+ super().__init__(**kwargs)
+
+ def register_normalized_timestamps(self, num_timestamps: int):
+ self.normalized_timestamps = torch.linspace(0, 1, num_timestamps, device=self.device)
+
+ def _init_models(self):
+ # gaussian model classes
+ self.gaussian_classes["Background"] = GSModelType.Background
+
+ for class_name, model_cfg in self.model_config.items():
+ # update model config for gaussian classes
+ if class_name in self.gaussian_classes:
+ model_cfg = self.model_config.pop(class_name)
+ self.model_config[class_name] = self.update_gaussian_cfg(model_cfg)
+
+ if class_name in self.gaussian_classes.keys():
+ model = import_str(model_cfg.type)(
+ **model_cfg,
+ class_name=class_name,
+ scene_scale=self.scene_radius,
+ scene_origin=self.scene_origin,
+ num_train_images=self.num_train_images,
+ device=self.device
+ )
+
+ if class_name in self.misc_classes_keys:
+ model = import_str(model_cfg.type)(
+ class_name=class_name,
+ **model_cfg.get('params', {}),
+ n=self.num_full_images,
+ device=self.device
+ ).to(self.device)
+
+ self.models[class_name] = model
+
+ logger.info(f"Initialized models: {self.models.keys()}")
+
+ # register normalized timestamps
+ self.register_normalized_timestamps(self.num_timesteps)
+ for class_name in self.gaussian_classes.keys():
+ model = self.models[class_name]
+ if hasattr(model, 'register_normalized_timestamps'):
+ model.register_normalized_timestamps(self.normalized_timestamps)
+ if hasattr(model, 'set_bbox'):
+ model.set_bbox(self.aabb)
+
+ def init_gaussians_from_dataset(
+ self,
+ dataset: DrivingDataset,
+ ) -> None:
+ for class_name in self.gaussian_classes:
+ model_cfg = self.model_config[class_name]
+ model = self.models[class_name]
+ if class_name == 'Background':
+ # ------ initialize gaussians ------
+ init_cfg = model_cfg.pop('init')
+ # sample points from the lidar point clouds
+ if init_cfg.get("from_lidar", None) is not None:
+ sampled_pts, sampled_color, sampled_time = dataset.get_lidar_samples(
+ **init_cfg.from_lidar, device=self.device
+ )
+ else:
+ sampled_pts, sampled_color, sampled_time = \
+ torch.empty(0, 3).to(self.device), torch.empty(0, 3).to(self.device), None
+
+ random_pts = []
+ num_near_pts = init_cfg.get('near_randoms', 0)
+ if num_near_pts > 0: # uniformly sample points inside the scene's sphere
+ num_near_pts *= 3 # since some invisible points will be filtered out
+ random_pts.append(uniform_sample_sphere(num_near_pts, self.device))
+ num_far_pts = init_cfg.get('far_randoms', 0)
+ if num_far_pts > 0: # inverse distances uniformly from (0, 1 / scene_radius)
+ num_far_pts *= 3
+ random_pts.append(uniform_sample_sphere(num_far_pts, self.device, inverse=True))
+
+ if num_near_pts + num_far_pts > 0:
+ random_pts = torch.cat(random_pts, dim=0)
+ random_pts = random_pts * self.scene_radius + self.scene_origin
+ visible_mask = dataset.check_pts_visibility(random_pts)
+ valid_pts = random_pts[visible_mask]
+
+ sampled_pts = torch.cat([sampled_pts, valid_pts], dim=0)
+ sampled_color = torch.cat([sampled_color, torch.rand(valid_pts.shape, ).to(self.device)], dim=0)
+ if sampled_time is not None:
+ sampled_time = torch.cat([sampled_time, torch.zeros(valid_pts.shape[0], 1).to(self.device)], dim=0)
+
+ if isinstance(model, PeriodicVibrationGaussians):
+ model.create_from_pcd(
+ init_means=sampled_pts, init_colors=sampled_color, init_times=sampled_time
+ )
+ else:
+ model.create_from_pcd(
+ init_means=sampled_pts, init_colors=sampled_color
+ )
+
+ logger.info(f"Initialized {class_name} gaussians")
+ logger.info(f"Initialized gaussians from pcd")
+
+ def forward(
+ self,
+ image_infos: Dict[str, torch.Tensor],
+ camera_infos: Dict[str, torch.Tensor],
+ novel_view: bool = False
+ ) -> Dict[str, torch.Tensor]:
+ """Forward pass of the model
+
+ Args:
+ image_infos (Dict[str, torch.Tensor]): image and pixels information
+ camera_infos (Dict[str, torch.Tensor]): camera information
+ novel_view: whether the view is novel, if True, disable the camera refinement
+
+ Returns:
+ Dict[str, torch.Tensor]: output of the model
+ """
+
+ # set current time or use temporal smoothing
+ normed_time = image_infos["normed_time"].flatten()[0]
+ self.cur_frame = torch.argmin(
+ torch.abs(self.normalized_timestamps - normed_time)
+ )
+
+ # for evaluation
+ for model in self.models.values():
+ if hasattr(model, 'in_test_set'):
+ model.in_test_set = self.in_test_set
+
+ # assigne current frame to gaussian models
+ for class_name in self.gaussian_classes.keys():
+ model = self.models[class_name]
+ if hasattr(model, 'set_cur_frame'):
+ model.set_cur_frame(self.cur_frame)
+
+ # prapare data
+ processed_cam = self.process_camera(
+ camera_infos=camera_infos,
+ image_ids=image_infos["img_idx"].flatten()[0],
+ novel_view=novel_view
+ )
+ gs = self.collect_gaussians(
+ cam=processed_cam,
+ image_ids=image_infos["img_idx"].flatten()[0]
+ )
+
+ # render gaussians
+ outputs, _ = self.render_gaussians(
+ gs=gs,
+ cam=processed_cam,
+ near_plane=self.render_cfg.near_plane,
+ far_plane=self.render_cfg.far_plane,
+ render_mode="RGB+ED",
+ radius_clip=self.render_cfg.get('radius_clip', 0.)
+ )
+
+ # render sky
+ sky_model = self.models['Sky']
+ outputs["rgb_sky"] = sky_model(image_infos)
+ outputs["rgb_sky_blend"] = outputs["rgb_sky"] * (1.0 - outputs["opacity"])
+
+ # affine transformation
+ outputs["rgb"] = self.affine_transformation(
+ outputs["rgb_gaussians"] + outputs["rgb_sky"] * (1.0 - outputs["opacity"]), image_infos
+ )
+
+ return outputs
+
+ def compute_losses(
+ self,
+ outputs: Dict[str, torch.Tensor],
+ image_infos: Dict[str, torch.Tensor],
+ cam_infos: Dict[str, torch.Tensor],
+ ) -> Dict[str, torch.Tensor]:
+ loss_dict = super().compute_losses(outputs, image_infos, cam_infos)
+
+ return loss_dict
+
+ def compute_metrics(
+ self,
+ outputs: Dict[str, torch.Tensor],
+ image_infos: Dict[str, torch.Tensor],
+ ) -> Dict[str, torch.Tensor]:
+ metric_dict = super().compute_metrics(outputs, image_infos)
+
+ return metric_dict
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/models/video_utils.py b/worldbench/third_party/drivestudio/models/video_utils.py
new file mode 100644
index 0000000..4d7c0aa
--- /dev/null
+++ b/worldbench/third_party/drivestudio/models/video_utils.py
@@ -0,0 +1,622 @@
+from typing import Literal, Dict, List, Optional, Callable
+from tqdm import tqdm, trange
+import numpy as np
+import os
+import logging
+import imageio
+from PIL import Image
+
+import torch
+from torch import Tensor
+from torch.nn import functional as F
+from skimage.metrics import structural_similarity as ssim
+
+from datasets.base import SplitWrapper
+from models.trainers.base import BasicTrainer
+from utils.visualization import (
+ to8b,
+ depth_visualizer,
+)
+
+logger = logging.getLogger()
+
+def get_numpy(x: Tensor) -> np.ndarray:
+ return x.squeeze().cpu().numpy()
+
+def non_zero_mean(x: Tensor) -> float:
+ return sum(x) / len(x) if len(x) > 0 else -1
+
+def compute_psnr(prediction: Tensor, target: Tensor) -> float:
+ """
+ Computes the Peak Signal-to-Noise Ratio (PSNR) between the prediction and target tensors.
+
+ Args:
+ prediction (torch.Tensor): The predicted tensor.
+ target (torch.Tensor): The target tensor.
+
+ Returns:
+ float: The PSNR value between the prediction and target tensors.
+ """
+ if not isinstance(prediction, Tensor):
+ prediction = Tensor(prediction)
+ if not isinstance(target, Tensor):
+ target = Tensor(target).to(prediction.device)
+ return (-10 * torch.log10(F.mse_loss(prediction, target))).item()
+
+
+def render_images(
+ trainer: BasicTrainer,
+ dataset: SplitWrapper,
+ compute_metrics: bool = False,
+ compute_error_map: bool = False,
+ vis_indices: Optional[List[int]] = None,
+ save_image_pairs: Optional[bool] = False
+):
+ """
+ Render pixel-related outputs from a model.
+
+ Args:
+ ....skip obvious args
+ compute_metrics (bool, optional): Whether to compute metrics. Defaults to False.
+ vis_indices (Optional[List[int]], optional): Indices to visualize. Defaults to None.
+ """
+ trainer.set_eval()
+ render_results = render(
+ dataset,
+ trainer=trainer,
+ compute_metrics=compute_metrics,
+ compute_error_map=compute_error_map,
+ vis_indices=vis_indices,
+ save_image_pairs=save_image_pairs
+ )
+ if compute_metrics:
+ num_samples = len(dataset) if vis_indices is None else len(vis_indices)
+ logger.info(f"Eval over {num_samples} images:")
+ logger.info(f"\t Full Image PSNR: {render_results['psnr']:.4f}")
+ logger.info(f"\t Full Image SSIM: {render_results['ssim']:.4f}")
+ logger.info(f"\t Full Image LPIPS: {render_results['lpips']:.4f}")
+ logger.info(f"\t Non-Sky PSNR: {render_results['occupied_psnr']:.4f}")
+ logger.info(f"\t Non-Sky SSIM: {render_results['occupied_ssim']:.4f}")
+ logger.info(f"\tDynamic-Only PSNR: {render_results['masked_psnr']:.4f}")
+ logger.info(f"\tDynamic-Only SSIM: {render_results['masked_ssim']:.4f}")
+ logger.info(f"\t Human-Only PSNR: {render_results['human_psnr']:.4f}")
+ logger.info(f"\t Human-Only SSIM: {render_results['human_ssim']:.4f}")
+ logger.info(f"\tVehicle-Only PSNR: {render_results['vehicle_psnr']:.4f}")
+ logger.info(f"\tVehicle-Only SSIM: {render_results['vehicle_ssim']:.4f}")
+
+ return render_results
+
+
+def render(
+ dataset: SplitWrapper,
+ trainer: BasicTrainer = None,
+ compute_metrics: bool = False,
+ compute_error_map: bool = False,
+ vis_indices: Optional[List[int]] = None,
+ save_image_pairs: Optional[bool] = False
+):
+ """
+ Renders a dataset utilizing a specified render function.
+
+ Parameters:
+ dataset: Dataset to render.
+ trainer: Gaussian trainer, includes gaussian models and rendering modules
+ compute_metrics: Optional; if True, the function will compute and return metrics. Default is False.
+ compute_error_map: Optional; if True, the function will compute and return error maps. Default is False.
+ vis_indices: Optional; if not None, the function will only render the specified indices. Default is None.
+ """
+ # rgbs
+ rgbs, gt_rgbs, rgb_sky_blend, rgb_sky = [], [], [], []
+ Background_rgbs, RigidNodes_rgbs, DeformableNodes_rgbs, SMPLNodes_rgbs, Dynamic_rgbs = [], [], [], [], []
+ error_maps = []
+
+ # depths
+ depths, lidar_on_images = [], []
+ Background_depths, RigidNodes_depths, DeformableNodes_depths, SMPLNodes_depths, Dynamic_depths = [], [], [], [], []
+
+ # sky
+ opacities, sky_masks = [], []
+ Background_opacities, RigidNodes_opacities, DeformableNodes_opacities, SMPLNodes_opacities, Dynamic_opacities = [], [], [], [], []
+
+ # misc
+ cam_names, cam_ids = [], []
+
+ if compute_metrics:
+ psnrs, ssim_scores, lpipss = [], [], []
+ masked_psnrs, masked_ssims = [], []
+ human_psnrs, human_ssims = [], []
+ vehicle_psnrs, vehicle_ssims = [], []
+ occupied_psnrs, occupied_ssims = [], []
+
+ with torch.no_grad():
+ indices = vis_indices if vis_indices is not None else range(len(dataset))
+ camera_downscale = trainer._get_downscale_factor()
+ for i in tqdm(indices, desc=f"rendering {dataset.split}", dynamic_ncols=True):
+ # get image and camera infos
+ image_infos, cam_infos = dataset.get_image(i, camera_downscale)
+ for k, v in image_infos.items():
+ if isinstance(v, Tensor):
+ image_infos[k] = v.cuda(non_blocking=True)
+ for k, v in cam_infos.items():
+ if isinstance(v, Tensor):
+ cam_infos[k] = v.cuda(non_blocking=True)
+ # render the image
+ results = trainer(image_infos, cam_infos)
+
+ # ------------- clip rgb ------------- #
+ for k, v in results.items():
+ if isinstance(v, Tensor) and "rgb" in k:
+ results[k] = v.clamp(0., 1.)
+
+ # ------------- cam names ------------- #
+ cam_names.append(cam_infos["cam_name"])
+ cam_ids.append(
+ cam_infos["cam_id"].flatten()[0].cpu().numpy()
+ )
+
+ # ------------- rgb ------------- #
+ rgb = results["rgb"]
+ rgbs.append(get_numpy(rgb))
+ if "pixels" in image_infos:
+ gt_rgbs.append(get_numpy(image_infos["pixels"]))
+
+ green_background = torch.tensor([0.0, 177, 64]) / 255.0
+ green_background = green_background.to(rgb.device)
+ if "Background_rgb" in results:
+ Background_rgb = results["Background_rgb"] * results[
+ "Background_opacity"
+ ] + green_background * (1 - results["Background_opacity"])
+ Background_rgbs.append(get_numpy(Background_rgb))
+ if "RigidNodes_rgb" in results:
+ RigidNodes_rgb = results["RigidNodes_rgb"] * results[
+ "RigidNodes_opacity"
+ ] + green_background * (1 - results["RigidNodes_opacity"])
+ RigidNodes_rgbs.append(get_numpy(RigidNodes_rgb))
+ if "DeformableNodes_rgb" in results:
+ DeformableNodes_rgb = results["DeformableNodes_rgb"] * results[
+ "DeformableNodes_opacity"
+ ] + green_background * (1 - results["DeformableNodes_opacity"])
+ DeformableNodes_rgbs.append(get_numpy(DeformableNodes_rgb))
+ if "SMPLNodes_rgb" in results:
+ SMPLNodes_rgb = results["SMPLNodes_rgb"] * results[
+ "SMPLNodes_opacity"
+ ] + green_background * (1 - results["SMPLNodes_opacity"])
+ SMPLNodes_rgbs.append(get_numpy(SMPLNodes_rgb))
+ if "Dynamic_rgb" in results:
+ Dynamic_rgb = results["Dynamic_rgb"] * results[
+ "Dynamic_opacity"
+ ] + green_background * (1 - results["Dynamic_opacity"])
+ Dynamic_rgbs.append(get_numpy(Dynamic_rgb))
+ if compute_error_map:
+ # cal mean squared error
+ error_map = (rgb - image_infos["pixels"]) ** 2
+ error_map = error_map.mean(dim=-1, keepdim=True)
+ # scale
+ error_map = (error_map - error_map.min()) / (error_map.max() - error_map.min())
+ error_map = error_map.repeat_interleave(3, dim=-1)
+ error_maps.append(get_numpy(error_map))
+ if "rgb_sky_blend" in results:
+ rgb_sky_blend.append(get_numpy(results["rgb_sky_blend"]))
+ if "rgb_sky" in results:
+ rgb_sky.append(get_numpy(results["rgb_sky"]))
+ # ------------- depth ------------- #
+ depth = results["depth"]
+ depths.append(get_numpy(depth))
+ # ------------- mask ------------- #
+ if "opacity" in results:
+ opacities.append(get_numpy(results["opacity"]))
+ if "Background_depth" in results:
+ Background_depths.append(get_numpy(results["Background_depth"]))
+ Background_opacities.append(get_numpy(results["Background_opacity"]))
+ if "RigidNodes_depth" in results:
+ RigidNodes_depths.append(get_numpy(results["RigidNodes_depth"]))
+ RigidNodes_opacities.append(get_numpy(results["RigidNodes_opacity"]))
+ if "DeformableNodes_depth" in results:
+ DeformableNodes_depths.append(get_numpy(results["DeformableNodes_depth"]))
+ DeformableNodes_opacities.append(get_numpy(results["DeformableNodes_opacity"]))
+ if "SMPLNodes_depth" in results:
+ SMPLNodes_depths.append(get_numpy(results["SMPLNodes_depth"]))
+ SMPLNodes_opacities.append(get_numpy(results["SMPLNodes_opacity"]))
+ if "Dynamic_depth" in results:
+ Dynamic_depths.append(get_numpy(results["Dynamic_depth"]))
+ Dynamic_opacities.append(get_numpy(results["Dynamic_opacity"]))
+ if "sky_masks" in image_infos:
+ sky_masks.append(get_numpy(image_infos["sky_masks"]))
+
+ # ------------- lidar ------------- #
+ if "lidar_depth_map" in image_infos:
+ depth_map = image_infos["lidar_depth_map"]
+ depth_img = depth_map.cpu().numpy()
+ depth_img = depth_visualizer(depth_img, depth_img > 0)
+ mask = (depth_map.unsqueeze(-1) > 0).cpu().numpy()
+ lidar_on_image = image_infos["pixels"].cpu().numpy() * (1 - mask) + depth_img * mask
+ lidar_on_images.append(lidar_on_image)
+
+ if compute_metrics:
+ psnr = compute_psnr(rgb, image_infos["pixels"])
+ ssim_score = ssim(
+ get_numpy(rgb),
+ get_numpy(image_infos["pixels"]),
+ data_range=1.0,
+ channel_axis=-1,
+ )
+ if save_image_pairs:
+ save_path = 'image_pairs'
+ os.makedirs(os.path.join(save_path, 'gt'), exist_ok=True)
+ os.makedirs(os.path.join(save_path, 'pred'), exist_ok=True)
+ Image.fromarray((get_numpy(rgb)*255).astype('uint8')).save(f'{save_path}/pred/{i}.png')
+ Image.fromarray((get_numpy(image_infos["pixels"])*255).astype('uint8')).save(f'{save_path}/gt/{i}.png')
+
+ lpips = trainer.lpips(
+ rgb[None, ...].permute(0, 3, 1, 2),
+ image_infos["pixels"][None, ...].permute(0, 3, 1, 2)
+ )
+ logger.info(f"Frame {i}: PSNR {psnr:.4f}, SSIM {ssim_score:.4f}")
+ psnrs.append(psnr)
+ ssim_scores.append(ssim_score)
+ lpipss.append(lpips.item())
+
+ if "sky_masks" in image_infos:
+ occupied_mask = ~get_numpy(image_infos["sky_masks"]).astype(bool)
+ if occupied_mask.sum() > 0:
+ occupied_psnrs.append(
+ compute_psnr(
+ rgb[occupied_mask], image_infos["pixels"][occupied_mask]
+ )
+ )
+ occupied_ssims.append(
+ ssim(
+ get_numpy(rgb),
+ get_numpy(image_infos["pixels"]),
+ data_range=1.0,
+ channel_axis=-1,
+ full=True,
+ )[1][occupied_mask].mean()
+ )
+
+ if "dynamic_masks" in image_infos:
+ dynamic_mask = get_numpy(image_infos["dynamic_masks"]).astype(bool)
+ if dynamic_mask.sum() > 0:
+ masked_psnrs.append(
+ compute_psnr(
+ rgb[dynamic_mask], image_infos["pixels"][dynamic_mask]
+ )
+ )
+ masked_ssims.append(
+ ssim(
+ get_numpy(rgb),
+ get_numpy(image_infos["pixels"]),
+ data_range=1.0,
+ channel_axis=-1,
+ full=True,
+ )[1][dynamic_mask].mean()
+ )
+
+ if "human_masks" in image_infos:
+ human_mask = get_numpy(image_infos["human_masks"]).astype(bool)
+ if human_mask.sum() > 0:
+ human_psnrs.append(
+ compute_psnr(
+ rgb[human_mask], image_infos["pixels"][human_mask]
+ )
+ )
+ human_ssims.append(
+ ssim(
+ get_numpy(rgb),
+ get_numpy(image_infos["pixels"]),
+ data_range=1.0,
+ channel_axis=-1,
+ full=True,
+ )[1][human_mask].mean()
+ )
+
+ if "vehicle_masks" in image_infos:
+ vehicle_mask = get_numpy(image_infos["vehicle_masks"]).astype(bool)
+ if vehicle_mask.sum() > 0:
+ vehicle_psnrs.append(
+ compute_psnr(
+ rgb[vehicle_mask], image_infos["pixels"][vehicle_mask]
+ )
+ )
+ vehicle_ssims.append(
+ ssim(
+ get_numpy(rgb),
+ get_numpy(image_infos["pixels"]),
+ data_range=1.0,
+ channel_axis=-1,
+ full=True,
+ )[1][vehicle_mask].mean()
+ )
+
+ # messy aggregation...
+ results_dict = {}
+ results_dict["psnr"] = non_zero_mean(psnrs) if compute_metrics else -1
+ results_dict["ssim"] = non_zero_mean(ssim_scores) if compute_metrics else -1
+ results_dict["lpips"] = non_zero_mean(lpipss) if compute_metrics else -1
+ results_dict["occupied_psnr"] = non_zero_mean(occupied_psnrs) if compute_metrics else -1
+ results_dict["occupied_ssim"] = non_zero_mean(occupied_ssims) if compute_metrics else -1
+ results_dict["masked_psnr"] = non_zero_mean(masked_psnrs) if compute_metrics else -1
+ results_dict["masked_ssim"] = non_zero_mean(masked_ssims) if compute_metrics else -1
+ results_dict["human_psnr"] = non_zero_mean(human_psnrs) if compute_metrics else -1
+ results_dict["human_ssim"] = non_zero_mean(human_ssims) if compute_metrics else -1
+ results_dict["vehicle_psnr"] = non_zero_mean(vehicle_psnrs) if compute_metrics else -1
+ results_dict["vehicle_ssim"] = non_zero_mean(vehicle_ssims) if compute_metrics else -1
+ results_dict["rgbs"] = rgbs
+ results_dict["depths"] = depths
+ results_dict["cam_names"] = cam_names
+ results_dict["cam_ids"] = cam_ids
+ if len(opacities) > 0:
+ results_dict["opacities"] = opacities
+ if len(gt_rgbs) > 0:
+ results_dict["gt_rgbs"] = gt_rgbs
+ if len(error_maps) > 0:
+ results_dict["rgb_error_maps"] = error_maps
+ if len(rgb_sky_blend) > 0:
+ results_dict["rgb_sky_blend"] = rgb_sky_blend
+ if len(rgb_sky) > 0:
+ results_dict["rgb_sky"] = rgb_sky
+ if len(sky_masks) > 0:
+ results_dict["gt_sky_masks"] = sky_masks
+ if len(lidar_on_images) > 0:
+ results_dict["lidar_on_images"] = lidar_on_images
+ if len(Background_rgbs) > 0:
+ results_dict["Background_rgbs"] = Background_rgbs
+ if len(RigidNodes_rgbs) > 0:
+ results_dict["RigidNodes_rgbs"] = RigidNodes_rgbs
+ if len(DeformableNodes_rgbs) > 0:
+ results_dict["DeformableNodes_rgbs"] = DeformableNodes_rgbs
+ if len(SMPLNodes_rgbs) > 0:
+ results_dict["SMPLNodes_rgbs"] = SMPLNodes_rgbs
+ if len(Dynamic_rgbs) > 0:
+ results_dict["Dynamic_rgbs"] = Dynamic_rgbs
+ if len(Background_depths) > 0:
+ results_dict["Background_depths"] = Background_depths
+ if len(RigidNodes_depths) > 0:
+ results_dict["RigidNodes_depths"] = RigidNodes_depths
+ if len(DeformableNodes_depths) > 0:
+ results_dict["DeformableNodes_depths"] = DeformableNodes_depths
+ if len(SMPLNodes_depths) > 0:
+ results_dict["SMPLNodes_depths"] = SMPLNodes_depths
+ if len(Dynamic_depths) > 0:
+ results_dict["Dynamic_depths"] = Dynamic_depths
+ if len(Background_opacities) > 0:
+ results_dict["Background_opacities"] = Background_opacities
+ if len(RigidNodes_opacities) > 0:
+ results_dict["RigidNodes_opacities"] = RigidNodes_opacities
+ if len(DeformableNodes_opacities) > 0:
+ results_dict["DeformableNodes_opacities"] = DeformableNodes_opacities
+ if len(SMPLNodes_opacities) > 0:
+ results_dict["SMPLNodes_opacities"] = SMPLNodes_opacities
+ if len(Dynamic_opacities) > 0:
+ results_dict["Dynamic_opacities"] = Dynamic_opacities
+ return results_dict
+
+
+def save_videos(
+ render_results: Dict[str, List[Tensor]],
+ save_pth: str,
+ layout: Callable,
+ num_timestamps: int,
+ keys: List[str] = ["gt_rgbs", "rgbs", "depths"],
+ num_cams: int = 3,
+ save_seperate_video: bool = False,
+ save_images: bool = False,
+ fps: int = 10,
+ verbose: bool = True,
+):
+ if save_seperate_video:
+ return_frame = save_seperate_videos(
+ render_results,
+ save_pth,
+ layout,
+ num_timestamps=num_timestamps,
+ keys=keys,
+ num_cams=num_cams,
+ save_images=save_images,
+ fps=fps,
+ verbose=verbose,
+ )
+ else:
+ return_frame = save_concatenated_videos(
+ render_results,
+ save_pth,
+ layout,
+ num_timestamps=num_timestamps,
+ keys=keys,
+ num_cams=num_cams,
+ save_images=save_images,
+ fps=fps,
+ verbose=verbose,
+ )
+ return return_frame
+
+
+def render_novel_views(trainer, render_data: list, save_path: str, fps: int = 30) -> None:
+ """
+ Perform rendering and save the result as a video.
+
+ Args:
+ trainer: Trainer object containing the rendering method
+ render_data (list): List of dicts, each containing elements required for rendering a single frame
+ save_path (str): Path to save the output video
+ fps (int): Frames per second for the output video
+ """
+ trainer.set_eval()
+
+ writer = imageio.get_writer(save_path, mode='I', fps=fps)
+
+ with torch.no_grad():
+ for frame_data in render_data:
+ # Move data to GPU
+ for key, value in frame_data["cam_infos"].items():
+ frame_data["cam_infos"][key] = value.cuda(non_blocking=True)
+ for key, value in frame_data["image_infos"].items():
+ frame_data["image_infos"][key] = value.cuda(non_blocking=True)
+
+ # Perform rendering
+ outputs = trainer(
+ image_infos=frame_data["image_infos"],
+ camera_infos=frame_data["cam_infos"],
+ novel_view=True
+ )
+
+ # Extract RGB image and mask
+ rgb = outputs["rgb"].cpu().numpy().clip(
+ min=1.e-6, max=1-1.e-6
+ )
+
+ # Convert to uint8 and write to video
+ rgb_uint8 = (rgb * 255).astype(np.uint8)
+ writer.append_data(rgb_uint8)
+
+ writer.close()
+ print(f"Video saved to {save_path}")
+
+
+def save_concatenated_videos(
+ render_results: Dict[str, List[Tensor]],
+ save_pth: str,
+ layout: Callable,
+ num_timestamps: int,
+ keys: List[str] = ["gt_rgbs", "rgbs", "depths"],
+ num_cams: int = 3,
+ save_images: bool = False,
+ fps: int = 10,
+ verbose: bool = True,
+):
+ if num_timestamps == 1: # it's an image
+ writer = imageio.get_writer(save_pth, mode="I")
+ return_frame_id = 0
+ else:
+ return_frame_id = num_timestamps // 2
+ writer = imageio.get_writer(save_pth, mode="I", fps=fps)
+ for i in trange(num_timestamps, desc="saving video", dynamic_ncols=True):
+ merged_list = []
+ cam_names = render_results["cam_names"][i * num_cams : (i + 1) * num_cams]
+ for key in keys:
+ # skip if the key is not in render_results
+ if "mask" in key:
+ new_key = key.replace("mask", "opacities")
+ if new_key not in render_results or len(render_results[new_key]) == 0:
+ continue
+ frames = render_results[new_key][i * num_cams : (i + 1) * num_cams]
+ else:
+ if key not in render_results or len(render_results[key]) == 0:
+ continue
+ frames = render_results[key][i * num_cams : (i + 1) * num_cams]
+ # convert to rgb if necessary
+ if key == "gt_sky_masks":
+ frames = [np.stack([frame, frame, frame], axis=-1) for frame in frames]
+ elif "mask" in key:
+ frames = [
+ np.stack([frame, frame, frame], axis=-1) for frame in frames
+ ]
+ elif "depth" in key:
+ try:
+ opacities = render_results[key.replace("depths", "opacities")][
+ i * num_cams : (i + 1) * num_cams
+ ]
+ except:
+ if "median" in key:
+ opacities = render_results[
+ key.replace("median_depths", "opacities")
+ ][i * num_cams : (i + 1) * num_cams]
+ else:
+ continue
+ frames = [
+ depth_visualizer(frame, opacity)
+ for frame, opacity in zip(frames, opacities)
+ ]
+ tiled_img = layout(frames, cam_names)
+ # frames = np.concatenate(frames, axis=1)
+ merged_list.append(tiled_img)
+ merged_frame = to8b(np.concatenate(merged_list, axis=0))
+ if i == return_frame_id:
+ return_frame = merged_frame
+ writer.append_data(merged_frame)
+ writer.close()
+ if verbose:
+ logger.info(f"saved video to {save_pth}")
+ del render_results
+ return {"concatenated_frame": return_frame}
+
+
+def save_seperate_videos(
+ render_results: Dict[str, List[Tensor]],
+ save_pth: str,
+ layout: Callable,
+ num_timestamps: int,
+ keys: List[str] = ["gt_rgbs", "rgbs", "depths"],
+ num_cams: int = 3,
+ fps: int = 10,
+ verbose: bool = False,
+ save_images: bool = False,
+):
+ return_frame_id = num_timestamps // 2
+ return_frame_dict = {}
+ for key in keys:
+ tmp_save_pth = save_pth.replace(".mp4", f"_{key}.mp4")
+ tmp_save_pth = tmp_save_pth.replace(".png", f"_{key}.png")
+ if num_timestamps == 1: # it's an image
+ writer = imageio.get_writer(tmp_save_pth, mode="I")
+ else:
+ writer = imageio.get_writer(tmp_save_pth, mode="I", fps=fps)
+ if "mask" not in key:
+ if key not in render_results or len(render_results[key]) == 0:
+ continue
+ for i in range(num_timestamps):
+ cam_names = render_results["cam_names"][i * num_cams : (i + 1) * num_cams]
+ # skip if the key is not in render_results
+ if "mask" in key:
+ new_key = key.replace("mask", "opacities")
+ if new_key not in render_results or len(render_results[new_key]) == 0:
+ continue
+ frames = render_results[new_key][i * num_cams : (i + 1) * num_cams]
+ else:
+ if key not in render_results or len(render_results[key]) == 0:
+ continue
+ frames = render_results[key][i * num_cams : (i + 1) * num_cams]
+ # convert to rgb if necessary
+ if key == "gt_sky_masks":
+ frames = [np.stack([frame, frame, frame], axis=-1) for frame in frames]
+ elif "mask" in key:
+ frames = [
+ np.stack([frame, frame, frame], axis=-1) for frame in frames
+ ]
+ elif "depth" in key:
+ try:
+ opacities = render_results[key.replace("depths", "opacities")][
+ i * num_cams : (i + 1) * num_cams
+ ]
+ except:
+ if "median" in key:
+ opacities = render_results[
+ key.replace("median_depths", "opacities")
+ ][i * num_cams : (i + 1) * num_cams]
+ else:
+ continue
+ frames = [
+ depth_visualizer(frame, opacity)
+ for frame, opacity in zip(frames, opacities)
+ ]
+ tiled_img = layout(frames, cam_names)
+ if save_images:
+ if i == 0:
+ os.makedirs(tmp_save_pth.replace(".mp4", ""), exist_ok=True)
+ for j, frame in enumerate(frames):
+ imageio.imwrite(
+ tmp_save_pth.replace(".mp4", f"/{i:03d}_{j:03d}.png"),
+ to8b(frame),
+ )
+ # frames = to8b(np.concatenate(frames, axis=1))
+ frames = to8b(tiled_img)
+ writer.append_data(frames)
+ if i == return_frame_id:
+ return_frame_dict[key] = frames
+ # close the writer
+ writer.close()
+ del writer
+ if verbose:
+ logger.info(f"saved video to {tmp_save_pth}")
+ del render_results
+ return return_frame_dict
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/requirements.txt b/worldbench/third_party/drivestudio/requirements.txt
new file mode 100644
index 0000000..77e9168
--- /dev/null
+++ b/worldbench/third_party/drivestudio/requirements.txt
@@ -0,0 +1,33 @@
+# PyTorch and related libraries
+--extra-index-url https://download.pytorch.org/whl/cu117
+torch==2.0.0+cu117
+torchvision==0.15.0+cu117
+xformers==0.0.18
+timm==0.9.5
+pytorch_msssim==1.0.0
+
+omegaconf==2.3.0
+torchmetrics==0.10.3
+
+tensorboard==2.11.0
+wandb==0.15.8
+matplotlib==3.5.3
+plotly==5.13.1
+viser==0.2.1
+
+imageio
+imageio-ffmpeg
+scikit-image==0.20.0
+opencv-python
+
+open3d==0.16.0
+pyquaternion==0.9.9
+chumpy
+numpy==1.23.1
+kornia==0.7.2
+
+tqdm
+gdown
+
+nerfview==0.0.3
+lpips==0.1.4
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/setup.py b/worldbench/third_party/drivestudio/setup.py
new file mode 100644
index 0000000..18535b8
--- /dev/null
+++ b/worldbench/third_party/drivestudio/setup.py
@@ -0,0 +1,7 @@
+from setuptools import setup, find_packages
+
+setup(
+ name='drivestudio', # 可以自定义
+ version='0.1',
+ packages=find_packages(), # 自动查找当前目录下的所有包（含 __init__.py 的目录）
+)
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/.gitignore b/worldbench/third_party/drivestudio/third_party/smplx/.gitignore
new file mode 100644
index 0000000..095c81a
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/.gitignore
@@ -0,0 +1,108 @@
+#### joe made this: http://goel.io/joe
+
+#####=== Python ===#####
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+# Usually these files are written by a python script from a template
+# before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/LICENSE b/worldbench/third_party/drivestudio/third_party/smplx/LICENSE
new file mode 100644
index 0000000..920431e
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/LICENSE
@@ -0,0 +1,58 @@
+License
+
+Software Copyright License for non-commercial scientific research purposes
+Please read carefully the following terms and conditions and any accompanying documentation before you download and/or use the SMPL-X/SMPLify-X model, data and software, (the "Model & Software"), including 3D meshes, blend weights, blend shapes, textures, software, scripts, and animations. By downloading and/or using the Model & Software (including downloading, cloning, installing, and any other use of this github repository), you acknowledge that you have read these terms and conditions, understand them, and agree to be bound by them. If you do not agree with these terms and conditions, you must not download and/or use the Model & Software. Any infringement of the terms of this agreement will automatically terminate your rights under this License
+
+Ownership / Licensees
+The Software and the associated materials has been developed at the
+
+Max Planck Institute for Intelligent Systems (hereinafter "MPI").
+
+Any copyright or patent right is owned by and proprietary material of the
+
+Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (hereinafter “MPG”; MPI and MPG hereinafter collectively “Max-Planck”)
+
+hereinafter the “Licensor”.
+
+License Grant
+Licensor grants you (Licensee) personally a single-user, non-exclusive, non-transferable, free of charge right:
+
+To install the Model & Software on computers owned, leased or otherwise controlled by you and/or your organization;
+To use the Model & Software for the sole purpose of performing non-commercial scientific research, non-commercial education, or non-commercial artistic projects;
+Any other use, in particular any use for commercial, pornographic, military, or surveillance, purposes is prohibited. This includes, without limitation, incorporation in a commercial product, use in a commercial service, or production of other artifacts for commercial purposes. The Data & Software may not be used to create fake, libelous, misleading, or defamatory content of any kind excluding analyses in peer-reviewed scientific research. The Data & Software may not be reproduced, modified and/or made available in any form to any third party without Max-Planck’s prior written permission.
+
+The Data & Software may not be used for pornographic purposes or to generate pornographic material whether commercial or not. This license also prohibits the use of the Software to train methods/algorithms/neural networks/etc. for commercial, pornographic, military, surveillance, or defamatory use of any kind. By downloading the Data & Software, you agree not to reverse engineer it.
+
+No Distribution
+The Model & Software and the license herein granted shall not be copied, shared, distributed, re-sold, offered for re-sale, transferred or sub-licensed in whole or in part except that you may make one copy for archive purposes only.
+
+Disclaimer of Representations and Warranties
+You expressly acknowledge and agree that the Model & Software results from basic research, is provided “AS IS”, may contain errors, and that any use of the Model & Software is at your sole risk. LICENSOR MAKES NO REPRESENTATIONS OR WARRANTIES OF ANY KIND CONCERNING THE MODEL & SOFTWARE, NEITHER EXPRESS NOR IMPLIED, AND THE ABSENCE OF ANY LEGAL OR ACTUAL DEFECTS, WHETHER DISCOVERABLE OR NOT. Specifically, and not to limit the foregoing, licensor makes no representations or warranties (i) regarding the merchantability or fitness for a particular purpose of the Model & Software, (ii) that the use of the Model & Software will not infringe any patents, copyrights or other intellectual property rights of a third party, and (iii) that the use of the Model & Software will not cause any damage of any kind to you or a third party.
+
+Limitation of Liability
+Because this Model & Software License Agreement qualifies as a donation, according to Section 521 of the German Civil Code (Bürgerliches Gesetzbuch – BGB) Licensor as a donor is liable for intent and gross negligence only. If the Licensor fraudulently conceals a legal or material defect, they are obliged to compensate the Licensee for the resulting damage.
+Licensor shall be liable for loss of data only up to the amount of typical recovery costs which would have arisen had proper and regular data backup measures been taken. For the avoidance of doubt Licensor shall be liable in accordance with the German Product Liability Act in the event of product liability. The foregoing applies also to Licensor’s legal representatives or assistants in performance. Any further liability shall be excluded.
+Patent claims generated through the usage of the Model & Software cannot be directed towards the copyright holders.
+The Model & Software is provided in the state of development the licensor defines. If modified or extended by Licensee, the Licensor makes no claims about the fitness of the Model & Software and is not responsible for any problems such modifications cause.
+
+No Maintenance Services
+You understand and agree that Licensor is under no obligation to provide either maintenance services, update services, notices of latent defects, or corrections of defects with regard to the Model & Software. Licensor nevertheless reserves the right to update, modify, or discontinue the Model & Software at any time.
+
+Defects of the Model & Software must be notified in writing to the Licensor with a comprehensible description of the error symptoms. The notification of the defect should enable the reproduction of the error. The Licensee is encouraged to communicate any use, results, modification or publication.
+
+Publications using the Model & Software
+You acknowledge that the Model & Software is a valuable scientific resource and agree to appropriately reference the following paper in any publication making use of the Model & Software.
+
+Citation:
+
+
+@inproceedings{SMPL-X:2019,
+ title = {Expressive Body Capture: 3D Hands, Face, and Body from a Single Image},
+ author = {Pavlakos, Georgios and Choutas, Vasileios and Ghorbani, Nima and Bolkart, Timo and Osman, Ahmed A. A. and Tzionas, Dimitrios and Black, Michael J.},
+ booktitle = {Proceedings IEEE Conf. on Computer Vision and Pattern Recognition (CVPR)},
+ year = {2019}
+}
+Commercial licensing opportunities
+For commercial uses of the Software, please send email to ps-license@tue.mpg.de
+
+This Agreement shall be governed by the laws of the Federal Republic of Germany except for the UN Sales Convention.
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/README.md b/worldbench/third_party/drivestudio/third_party/smplx/README.md
new file mode 100644
index 0000000..0a59aab
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/README.md
@@ -0,0 +1,203 @@
+## SMPL-X: A new joint 3D model of the human body, face and hands together
+
+[[Paper Page](https://smpl-x.is.tue.mpg.de)] [[Paper](https://ps.is.tuebingen.mpg.de/uploads_file/attachment/attachment/497/SMPL-X.pdf)]
+[[Supp. Mat.](https://ps.is.tuebingen.mpg.de/uploads_file/attachment/attachment/498/SMPL-X-supp.pdf)]
+
+## Table of Contents
+ * [License](#license)
+ * [Description](#description)
+ * [News](#news)
+ * [Installation](#installation)
+ * [Downloading the model](#downloading-the-model)
+ * [Loading SMPL-X, SMPL+H and SMPL](#loading-smpl-x-smplh-and-smpl)
+ * [SMPL and SMPL+H setup](#smpl-and-smplh-setup)
+ * [Model loading](https://github.com/vchoutas/smplx#model-loading)
+ * [MANO and FLAME correspondences](#mano-and-flame-correspondences)
+ * [Example](#example)
+ * [Modifying the global pose of the model](#modifying-the-global-pose-of-the-model)
+ * [Citation](#citation)
+ * [Acknowledgments](#acknowledgments)
+ * [Contact](#contact)
+
+## License
+
+Software Copyright License for **non-commercial scientific research purposes**.
+Please read carefully the [terms and conditions](https://github.com/vchoutas/smplx/blob/master/LICENSE) and any accompanying documentation before you download and/or use the SMPL-X/SMPLify-X model, data and software, (the "Model & Software"), including 3D meshes, blend weights, blend shapes, textures, software, scripts, and animations. By downloading and/or using the Model & Software (including downloading, cloning, installing, and any other use of this github repository), you acknowledge that you have read these terms and conditions, understand them, and agree to be bound by them. If you do not agree with these terms and conditions, you must not download and/or use the Model & Software. Any infringement of the terms of this agreement will automatically terminate your rights under this [License](./LICENSE).
+
+## Disclaimer
+
+The original images used for the figures 1 and 2 of the paper can be found in this link.
+The images in the paper are used under license from gettyimages.com.
+We have acquired the right to use them in the publication, but redistribution is not allowed.
+Please follow the instructions on the given link to acquire right of usage.
+Our results are obtained on the 483 × 724 pixels resolution of the original images.
+
+## Description
+
+*SMPL-X* (SMPL eXpressive) is a unified body model with shape parameters trained jointly for the
+face, hands and body. *SMPL-X* uses standard vertex based linear blend skinning with learned corrective blend
+shapes, has N = 10, 475 vertices and K = 54 joints,
+which include joints for the neck, jaw, eyeballs and fingers.
+SMPL-X is defined by a function M(θ, β, ψ), where θ is the pose parameters, β the shape parameters and
+ψ the facial expression parameters.
+
+## News
+
+- 3 November 2020: We release the code to transfer between the models in the
+ SMPL family. For more details on the code, go to this [readme
+ file](./transfer_model/README.md). A detailed explanation on how the mappings
+ were extracted can be found [here](./transfer_model/docs/transfer.md).
+- 23 September 2020: A UV map is now available for SMPL-X, please check the
+ Downloads section of the website.
+- 20 August 2020: The full shape and expression space of SMPL-X are now available.
+
+## Installation
+
+To install the model please follow the next steps in the specified order:
+1. To install from PyPi simply run:
+ ```Shell
+ pip install smplx[all]
+ ```
+2. Clone this repository and install it using the *setup.py* script:
+```Shell
+git clone https://github.com/vchoutas/smplx
+python setup.py install
+```
+
+## Downloading the model
+
+To download the *SMPL-X* model go to [this project website](https://smpl-x.is.tue.mpg.de) and register to get access to the downloads section.
+
+To download the *SMPL+H* model go to [this project website](http://mano.is.tue.mpg.de) and register to get access to the downloads section.
+
+To download the *SMPL* model go to [this](http://smpl.is.tue.mpg.de) (male and female models) and [this](http://smplify.is.tue.mpg.de) (gender neutral model) project website and register to get access to the downloads section.
+
+## Loading SMPL-X, SMPL+H and SMPL
+
+### SMPL and SMPL+H setup
+
+The loader gives the option to use any of the SMPL-X, SMPL+H, SMPL, and MANO models. Depending on the model you want to use, please follow the respective download instructions. To switch between MANO, SMPL, SMPL+H and SMPL-X just change the *model_path* or *model_type* parameters. For more details please check the docs of the model classes.
+Before using SMPL and SMPL+H you should follow the instructions in [tools/README.md](./tools/README.md) to remove the
+Chumpy objects from both model pkls, as well as merge the MANO parameters with SMPL+H.
+
+### Model loading
+
+You can either use the [create](https://github.com/vchoutas/smplx/blob/c63c02b478c5c6f696491ed9167e3af6b08d89b1/smplx/body_models.py#L54)
+function from [body_models](./smplx/body_models.py) or directly call the constructor for the
+[SMPL](https://github.com/vchoutas/smplx/blob/c63c02b478c5c6f696491ed9167e3af6b08d89b1/smplx/body_models.py#L106),
+[SMPL+H](https://github.com/vchoutas/smplx/blob/c63c02b478c5c6f696491ed9167e3af6b08d89b1/smplx/body_models.py#L395) and
+[SMPL-X](https://github.com/vchoutas/smplx/blob/c63c02b478c5c6f696491ed9167e3af6b08d89b1/smplx/body_models.py#L628) model. The path to the model can either be the path to the file with the parameters or a directory with the following structure:
+```bash
+models
+├── smpl
+│ ├── SMPL_FEMALE.pkl
+│ └── SMPL_MALE.pkl
+│ └── SMPL_NEUTRAL.pkl
+├── smplh
+│ ├── SMPLH_FEMALE.pkl
+│ └── SMPLH_MALE.pkl
+├── mano
+| ├── MANO_RIGHT.pkl
+| └── MANO_LEFT.pkl
+└── smplx
+ ├── SMPLX_FEMALE.npz
+ ├── SMPLX_FEMALE.pkl
+ ├── SMPLX_MALE.npz
+ ├── SMPLX_MALE.pkl
+ ├── SMPLX_NEUTRAL.npz
+ └── SMPLX_NEUTRAL.pkl
+```
+
+
+## MANO and FLAME correspondences
+
+The vertex correspondences between SMPL-X and MANO, FLAME can be downloaded
+from [the project website](https://smpl-x.is.tue.mpg.de). If you have extracted
+the correspondence data in the folder *correspondences*, then use the following
+scripts to visualize them:
+
+1. To view MANO correspondences run the following command:
+
+```
+python examples/vis_mano_vertices.py --model-folder $SMPLX_FOLDER --corr-fname correspondences/MANO_SMPLX_vertex_ids.pkl
+```
+
+2. To view FLAME correspondences run the following command:
+
+```
+python examples/vis_flame_vertices.py --model-folder $SMPLX_FOLDER --corr-fname correspondences/SMPL-X__FLAME_vertex_ids.npy
+```
+
+## Example
+
+After installing the *smplx* package and downloading the model parameters you should be able to run the *demo.py*
+script to visualize the results. For this step you have to install the [pyrender](https://pyrender.readthedocs.io/en/latest/index.html) and [trimesh](https://trimsh.org/) packages.
+
+`python examples/demo.py --model-folder $SMPLX_FOLDER --plot-joints=True --gender="neutral"`
+
+## Modifying the global pose of the model
+
+If you want to modify the global pose of the model, i.e. the root rotation and
+translation, to a new coordinate system for example, you need to take into
+account that the model rotation uses the pelvis as the center of rotation. A
+more detailed description can be found in the following
+[link](https://www.dropbox.com/scl/fi/zkatuv5shs8d4tlwr8ecc/Change-parameters-to-new-coordinate-system.paper?dl=0&rlkey=lotq1sh6wzkmyttisc05h0in0).
+If something is not clear, please let me know so that I can update the
+description.
+
+## Citation
+
+Depending on which model is loaded for your project, i.e. SMPL-X or SMPL+H or SMPL, please cite the most relevant work below, listed in the same order:
+
+```
+@inproceedings{SMPL-X:2019,
+ title = {Expressive Body Capture: 3D Hands, Face, and Body from a Single Image},
+ author = {Pavlakos, Georgios and Choutas, Vasileios and Ghorbani, Nima and Bolkart, Timo and Osman, Ahmed A. A. and Tzionas, Dimitrios and Black, Michael J.},
+ booktitle = {Proceedings IEEE Conf. on Computer Vision and Pattern Recognition (CVPR)},
+ year = {2019}
+}
+```
+
+```
+@article{MANO:SIGGRAPHASIA:2017,
+ title = {Embodied Hands: Modeling and Capturing Hands and Bodies Together},
+ author = {Romero, Javier and Tzionas, Dimitrios and Black, Michael J.},
+ journal = {ACM Transactions on Graphics, (Proc. SIGGRAPH Asia)},
+ volume = {36},
+ number = {6},
+ series = {245:1--245:17},
+ month = nov,
+ year = {2017},
+ month_numeric = {11}
+ }
+```
+
+```
+@article{SMPL:2015,
+ author = {Loper, Matthew and Mahmood, Naureen and Romero, Javier and Pons-Moll, Gerard and Black, Michael J.},
+ title = {{SMPL}: A Skinned Multi-Person Linear Model},
+ journal = {ACM Transactions on Graphics, (Proc. SIGGRAPH Asia)},
+ month = oct,
+ number = {6},
+ pages = {248:1--248:16},
+ publisher = {ACM},
+ volume = {34},
+ year = {2015}
+}
+```
+
+This repository was originally developed for SMPL-X / SMPLify-X (CVPR 2019), you might be interested in having a look: [https://smpl-x.is.tue.mpg.de](https://smpl-x.is.tue.mpg.de).
+
+## Acknowledgments
+
+### Facial Contour
+
+Special thanks to [Soubhik Sanyal](https://github.com/soubhiksanyal) for sharing the Tensorflow code used for the facial
+landmarks.
+
+## Contact
+The code of this repository was implemented by [Vassilis Choutas](vassilis.choutas@tuebingen.mpg.de).
+
+For questions, please contact [smplx@tue.mpg.de](smplx@tue.mpg.de).
+
+For commercial licensing (and all related questions for business applications), please contact [ps-licensing@tue.mpg.de](ps-licensing@tue.mpg.de).
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/config_files/smpl2smplh.yaml b/worldbench/third_party/drivestudio/third_party/smplx/config_files/smpl2smplh.yaml
new file mode 100644
index 0000000..98cc874
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/config_files/smpl2smplh.yaml
@@ -0,0 +1,25 @@
+datasets:
+ mesh_folder:
+ data_folder: 'transfer_data/meshes/smpl'
+deformation_transfer_path: 'transfer_data/smpl2smplh_def_transfer.pkl'
+mask_ids_fname: ''
+summary_steps: 100
+
+edge_fitting:
+ per_part: False
+
+optim:
+ type: 'trust-ncg'
+ maxiters: 100
+ gtol: 1e-06
+
+body_model:
+ model_type: "smplh"
+ # SMPL+H has no neutral model, so we have to manually select the gender
+ gender: "female"
+ # gender: "male"
+ folder: "transfer_data/body_models"
+ use_compressed: False
+ smplh:
+ betas:
+ num: 10
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/config_files/smpl2smplx.yaml b/worldbench/third_party/drivestudio/third_party/smplx/config_files/smpl2smplx.yaml
new file mode 100644
index 0000000..aad7ac5
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/config_files/smpl2smplx.yaml
@@ -0,0 +1,26 @@
+datasets:
+ mesh_folder:
+ data_folder: 'transfer_data/meshes/smpl'
+deformation_transfer_path: 'transfer_data/smpl2smplx_deftrafo_setup.pkl'
+mask_ids_fname: 'smplx_mask_ids.npy'
+summary_steps: 100
+
+edge_fitting:
+ per_part: False
+
+optim:
+ type: 'trust-ncg'
+ maxiters: 100
+ gtol: 1e-06
+
+body_model:
+ model_type: "smplx"
+ gender: "neutral"
+ folder: "transfer_data/body_models"
+ use_compressed: False
+ use_face_contour: True
+ smplx:
+ betas:
+ num: 10
+ expression:
+ num: 10
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/config_files/smplh2smpl.yaml b/worldbench/third_party/drivestudio/third_party/smplx/config_files/smplh2smpl.yaml
new file mode 100644
index 0000000..223d337
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/config_files/smplh2smpl.yaml
@@ -0,0 +1,24 @@
+datasets:
+ mesh_folder:
+ data_folder: 'transfer_data/meshes/smplh'
+deformation_transfer_path: 'transfer_data/smplh2smpl_def_transfer.pkl'
+mask_ids_fname: ''
+summary_steps: 100
+
+edge_fitting:
+ per_part: False
+
+optim:
+ type: 'trust-ncg'
+ maxiters: 100
+ gtol: 1e-06
+
+body_model:
+ model_type: "smpl"
+ gender: "neutral"
+ folder: "transfer_data/body_models"
+ use_compressed: False
+ use_face_contour: True
+ smpl:
+ betas:
+ num: 10
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/config_files/smplh2smplx.yaml b/worldbench/third_party/drivestudio/third_party/smplx/config_files/smplh2smplx.yaml
new file mode 100644
index 0000000..682d0e6
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/config_files/smplh2smplx.yaml
@@ -0,0 +1,26 @@
+datasets:
+ mesh_folder:
+ data_folder: 'transfer_data/meshes/smplh'
+deformation_transfer_path: 'transfer_data/smplh2smplx_deftrafo_setup.pkl'
+mask_ids_fname: 'smplx_mask_ids.npy'
+summary_steps: 100
+
+edge_fitting:
+ per_part: False
+
+optim:
+ type: 'trust-ncg'
+ maxiters: 100
+ gtol: 1e-06
+
+body_model:
+ model_type: "smplx"
+ gender: "neutral"
+ folder: "transfer_data/body_models"
+ use_compressed: False
+ use_face_contour: True
+ smplx:
+ betas:
+ num: 10
+ expression:
+ num: 10
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/config_files/smplh2smplx_as.yaml b/worldbench/third_party/drivestudio/third_party/smplx/config_files/smplh2smplx_as.yaml
new file mode 100644
index 0000000..7031248
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/config_files/smplh2smplx_as.yaml
@@ -0,0 +1,26 @@
+datasets:
+ mesh_folder:
+ data_folder: 'transfer_data/meshes/amass_sample'
+deformation_transfer_path: 'transfer_data/smplh2smplx_deftrafo_setup.pkl'
+mask_ids_fname: 'smplx_mask_ids.npy'
+summary_steps: 100
+
+edge_fitting:
+ per_part: False
+
+optim:
+ type: 'trust-ncg'
+ maxiters: 100
+ gtol: 1e-06
+
+body_model:
+ model_type: "smplx"
+ gender: "neutral"
+ folder: "models"
+ use_compressed: False
+ use_face_contour: True
+ smplx:
+ betas:
+ num: 10
+ expression:
+ num: 10
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/config_files/smplh2smplx_onepose.yaml b/worldbench/third_party/drivestudio/third_party/smplx/config_files/smplh2smplx_onepose.yaml
new file mode 100644
index 0000000..1a592b8
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/config_files/smplh2smplx_onepose.yaml
@@ -0,0 +1,27 @@
+datasets:
+ mesh_folder:
+ data_folder: 'transfer_data/meshes/amass_onepose'
+deformation_transfer_path: 'transfer_data/smplh2smplx_deftrafo_setup.pkl'
+mask_ids_fname: 'smplx_mask_ids.npy'
+summary_steps: 100
+
+edge_fitting:
+ per_part: False
+
+optim:
+ type: 'adam'
+ lr: 0.1
+ maxiters: 10000
+ gtol: 1e-06
+
+body_model:
+ model_type: "smplx"
+ gender: "neutral"
+ folder: "models"
+ use_compressed: False
+ use_face_contour: True
+ smplx:
+ betas:
+ num: 10
+ expression:
+ num: 10
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/config_files/smplx2smpl.yaml b/worldbench/third_party/drivestudio/third_party/smplx/config_files/smplx2smpl.yaml
new file mode 100644
index 0000000..7012fef
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/config_files/smplx2smpl.yaml
@@ -0,0 +1,25 @@
+datasets:
+ mesh_folder:
+ data_folder: 'meshes/smplx'
+deformation_transfer_path: 'transfer_data/smplx2smpl_deftrafo_setup.pkl'
+mask_ids_fname: ''
+summary_steps: 100
+
+edge_fitting:
+ per_part: False
+
+optim:
+ type: 'lbfgs'
+ maxiters: 200
+ gtol: 1e-06
+
+body_model:
+ model_type: "smpl"
+ gender: "neutral"
+ ext: 'pkl'
+ folder: "transfer_data/body_models"
+ use_compressed: False
+ use_face_contour: True
+ smpl:
+ betas:
+ num: 10
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/config_files/smplx2smplh.yaml b/worldbench/third_party/drivestudio/third_party/smplx/config_files/smplx2smplh.yaml
new file mode 100644
index 0000000..76275e0
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/config_files/smplx2smplh.yaml
@@ -0,0 +1,27 @@
+datasets:
+ mesh_folder:
+ data_folder: 'meshes/smplx'
+deformation_transfer_path: 'transfer_data/smplx2smplh_deftrafo_setup.pkl'
+mask_ids_fname: ''
+summary_steps: 100
+
+edge_fitting:
+ per_part: False
+
+optim:
+ type: 'lbfgs'
+ maxiters: 200
+ gtol: 1e-06
+
+body_model:
+ model_type: "smplh"
+ # SMPL+H has no neutral model, so we have to manually select the gender
+ gender: "female"
+ # gender: "male"
+ ext: 'pkl'
+ folder: "transfer_data/body_models"
+ use_compressed: False
+ use_face_contour: True
+ smplh:
+ betas:
+ num: 10
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/examples/demo.py b/worldbench/third_party/drivestudio/third_party/smplx/examples/demo.py
new file mode 100644
index 0000000..7a6fd50
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/examples/demo.py
@@ -0,0 +1,180 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+import os.path as osp
+import argparse
+
+import numpy as np
+import torch
+
+import smplx
+
+
+def main(model_folder,
+ model_type='smplx',
+ ext='npz',
+ gender='neutral',
+ plot_joints=False,
+ num_betas=10,
+ sample_shape=True,
+ sample_expression=True,
+ num_expression_coeffs=10,
+ plotting_module='pyrender',
+ use_face_contour=False):
+
+ model = smplx.create(model_folder, model_type=model_type,
+ gender=gender, use_face_contour=use_face_contour,
+ num_betas=num_betas,
+ num_expression_coeffs=num_expression_coeffs,
+ ext=ext)
+ print(model)
+
+ betas, expression = None, None
+ if sample_shape:
+ betas = torch.randn([1, model.num_betas], dtype=torch.float32)
+ if sample_expression:
+ expression = torch.randn(
+ [1, model.num_expression_coeffs], dtype=torch.float32)
+
+ output = model(betas=betas, expression=expression,
+ return_verts=True)
+ vertices = output.vertices.detach().cpu().numpy().squeeze()
+ joints = output.joints.detach().cpu().numpy().squeeze()
+
+ print('Vertices shape =', vertices.shape)
+ print('Joints shape =', joints.shape)
+
+ if plotting_module == 'pyrender':
+ import pyrender
+ import trimesh
+ vertex_colors = np.ones([vertices.shape[0], 4]) * [0.3, 0.3, 0.3, 0.8]
+ tri_mesh = trimesh.Trimesh(vertices, model.faces,
+ vertex_colors=vertex_colors)
+
+ mesh = pyrender.Mesh.from_trimesh(tri_mesh)
+
+ scene = pyrender.Scene()
+ scene.add(mesh)
+
+ if plot_joints:
+ sm = trimesh.creation.uv_sphere(radius=0.005)
+ sm.visual.vertex_colors = [0.9, 0.1, 0.1, 1.0]
+ tfs = np.tile(np.eye(4), (len(joints), 1, 1))
+ tfs[:, :3, 3] = joints
+ joints_pcl = pyrender.Mesh.from_trimesh(sm, poses=tfs)
+ scene.add(joints_pcl)
+
+ pyrender.Viewer(scene, use_raymond_lighting=True)
+ elif plotting_module == 'matplotlib':
+ from matplotlib import pyplot as plt
+ from mpl_toolkits.mplot3d import Axes3D
+ from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+
+ fig = plt.figure()
+ ax = fig.add_subplot(111, projection='3d')
+
+ mesh = Poly3DCollection(vertices[model.faces], alpha=0.1)
+ face_color = (1.0, 1.0, 0.9)
+ edge_color = (0, 0, 0)
+ mesh.set_edgecolor(edge_color)
+ mesh.set_facecolor(face_color)
+ ax.add_collection3d(mesh)
+ ax.scatter(joints[:, 0], joints[:, 1], joints[:, 2], color='r')
+
+ if plot_joints:
+ ax.scatter(joints[:, 0], joints[:, 1], joints[:, 2], alpha=0.1)
+ plt.show()
+ elif plotting_module == 'open3d':
+ import open3d as o3d
+
+ mesh = o3d.geometry.TriangleMesh()
+ mesh.vertices = o3d.utility.Vector3dVector(
+ vertices)
+ mesh.triangles = o3d.utility.Vector3iVector(model.faces)
+ mesh.compute_vertex_normals()
+ mesh.paint_uniform_color([0.3, 0.3, 0.3])
+
+ geometry = [mesh]
+ if plot_joints:
+ joints_pcl = o3d.geometry.PointCloud()
+ joints_pcl.points = o3d.utility.Vector3dVector(joints)
+ joints_pcl.paint_uniform_color([0.7, 0.3, 0.3])
+ geometry.append(joints_pcl)
+
+ o3d.visualization.draw_geometries(geometry)
+ else:
+ raise ValueError('Unknown plotting_module: {}'.format(plotting_module))
+
+
+if __name__ == '__main__':
+ parser = argparse.ArgumentParser(description='SMPL-X Demo')
+
+ parser.add_argument('--model-folder', required=True, type=str,
+ help='The path to the model folder')
+ parser.add_argument('--model-type', default='smplx', type=str,
+ choices=['smpl', 'smplh', 'smplx', 'mano', 'flame'],
+ help='The type of model to load')
+ parser.add_argument('--gender', type=str, default='neutral',
+ help='The gender of the model')
+ parser.add_argument('--num-betas', default=10, type=int,
+ dest='num_betas',
+ help='Number of shape coefficients.')
+ parser.add_argument('--num-expression-coeffs', default=10, type=int,
+ dest='num_expression_coeffs',
+ help='Number of expression coefficients.')
+ parser.add_argument('--plotting-module', type=str, default='pyrender',
+ dest='plotting_module',
+ choices=['pyrender', 'matplotlib', 'open3d'],
+ help='The module to use for plotting the result')
+ parser.add_argument('--ext', type=str, default='npz',
+ help='Which extension to use for loading')
+ parser.add_argument('--plot-joints', default=False,
+ type=lambda arg: arg.lower() in ['true', '1'],
+ help='The path to the model folder')
+ parser.add_argument('--sample-shape', default=True,
+ dest='sample_shape',
+ type=lambda arg: arg.lower() in ['true', '1'],
+ help='Sample a random shape')
+ parser.add_argument('--sample-expression', default=True,
+ dest='sample_expression',
+ type=lambda arg: arg.lower() in ['true', '1'],
+ help='Sample a random expression')
+ parser.add_argument('--use-face-contour', default=False,
+ type=lambda arg: arg.lower() in ['true', '1'],
+ help='Compute the contour of the face')
+
+ args = parser.parse_args()
+
+ model_folder = osp.expanduser(osp.expandvars(args.model_folder))
+ model_type = args.model_type
+ plot_joints = args.plot_joints
+ use_face_contour = args.use_face_contour
+ gender = args.gender
+ ext = args.ext
+ plotting_module = args.plotting_module
+ num_betas = args.num_betas
+ num_expression_coeffs = args.num_expression_coeffs
+ sample_shape = args.sample_shape
+ sample_expression = args.sample_expression
+
+ main(model_folder, model_type, ext=ext,
+ gender=gender, plot_joints=plot_joints,
+ num_betas=num_betas,
+ num_expression_coeffs=num_expression_coeffs,
+ sample_shape=sample_shape,
+ sample_expression=sample_expression,
+ plotting_module=plotting_module,
+ use_face_contour=use_face_contour)
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/examples/demo_layers.py b/worldbench/third_party/drivestudio/third_party/smplx/examples/demo_layers.py
new file mode 100644
index 0000000..8d4e182
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/examples/demo_layers.py
@@ -0,0 +1,181 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+import os.path as osp
+import argparse
+
+import numpy as np
+import torch
+
+import smplx
+
+
+def main(model_folder,
+ model_type='smplx',
+ ext='npz',
+ gender='neutral',
+ plot_joints=False,
+ num_betas=10,
+ sample_shape=True,
+ sample_expression=True,
+ num_expression_coeffs=10,
+ plotting_module='pyrender',
+ use_face_contour=False):
+
+ model = smplx.build_layer(
+ model_folder, model_type=model_type,
+ gender=gender, use_face_contour=use_face_contour,
+ num_betas=num_betas,
+ num_expression_coeffs=num_expression_coeffs,
+ ext=ext)
+ print(model)
+
+ betas, expression = None, None
+ if sample_shape:
+ betas = torch.randn([1, model.num_betas], dtype=torch.float32)
+ if sample_expression:
+ expression = torch.randn(
+ [1, model.num_expression_coeffs], dtype=torch.float32)
+
+ output = model(betas=betas, expression=expression,
+ return_verts=True)
+ vertices = output.vertices.detach().cpu().numpy().squeeze()
+ joints = output.joints.detach().cpu().numpy().squeeze()
+
+ print('Vertices shape =', vertices.shape)
+ print('Joints shape =', joints.shape)
+
+ if plotting_module == 'pyrender':
+ import pyrender
+ import trimesh
+ vertex_colors = np.ones([vertices.shape[0], 4]) * [0.3, 0.3, 0.3, 0.8]
+ tri_mesh = trimesh.Trimesh(vertices, model.faces,
+ vertex_colors=vertex_colors)
+
+ mesh = pyrender.Mesh.from_trimesh(tri_mesh)
+
+ scene = pyrender.Scene()
+ scene.add(mesh)
+
+ if plot_joints:
+ sm = trimesh.creation.uv_sphere(radius=0.005)
+ sm.visual.vertex_colors = [0.9, 0.1, 0.1, 1.0]
+ tfs = np.tile(np.eye(4), (len(joints), 1, 1))
+ tfs[:, :3, 3] = joints
+ joints_pcl = pyrender.Mesh.from_trimesh(sm, poses=tfs)
+ scene.add(joints_pcl)
+
+ pyrender.Viewer(scene, use_raymond_lighting=True)
+ elif plotting_module == 'matplotlib':
+ from matplotlib import pyplot as plt
+ from mpl_toolkits.mplot3d import Axes3D
+ from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+
+ fig = plt.figure()
+ ax = fig.add_subplot(111, projection='3d')
+
+ mesh = Poly3DCollection(vertices[model.faces], alpha=0.1)
+ face_color = (1.0, 1.0, 0.9)
+ edge_color = (0, 0, 0)
+ mesh.set_edgecolor(edge_color)
+ mesh.set_facecolor(face_color)
+ ax.add_collection3d(mesh)
+ ax.scatter(joints[:, 0], joints[:, 1], joints[:, 2], color='r')
+
+ if plot_joints:
+ ax.scatter(joints[:, 0], joints[:, 1], joints[:, 2], alpha=0.1)
+ plt.show()
+ elif plotting_module == 'open3d':
+ import open3d as o3d
+
+ mesh = o3d.geometry.TriangleMesh()
+ mesh.vertices = o3d.utility.Vector3dVector(
+ vertices)
+ mesh.triangles = o3d.utility.Vector3iVector(model.faces)
+ mesh.compute_vertex_normals()
+ mesh.paint_uniform_color([0.3, 0.3, 0.3])
+
+ geometry = [mesh]
+ if plot_joints:
+ joints_pcl = o3d.geometry.PointCloud()
+ joints_pcl.points = o3d.utility.Vector3dVector(joints)
+ joints_pcl.paint_uniform_color([0.7, 0.3, 0.3])
+ geometry.append(joints_pcl)
+
+ o3d.visualization.draw_geometries(geometry)
+ else:
+ raise ValueError('Unknown plotting_module: {}'.format(plotting_module))
+
+
+if __name__ == '__main__':
+ parser = argparse.ArgumentParser(description='SMPL-X Demo')
+
+ parser.add_argument('--model-folder', required=True, type=str,
+ help='The path to the model folder')
+ parser.add_argument('--model-type', default='smplx', type=str,
+ choices=['smpl', 'smplh', 'smplx', 'mano', 'flame'],
+ help='The type of model to load')
+ parser.add_argument('--gender', type=str, default='neutral',
+ help='The gender of the model')
+ parser.add_argument('--num-betas', default=10, type=int,
+ dest='num_betas',
+ help='Number of shape coefficients.')
+ parser.add_argument('--num-expression-coeffs', default=10, type=int,
+ dest='num_expression_coeffs',
+ help='Number of expression coefficients.')
+ parser.add_argument('--plotting-module', type=str, default='pyrender',
+ dest='plotting_module',
+ choices=['pyrender', 'matplotlib', 'open3d'],
+ help='The module to use for plotting the result')
+ parser.add_argument('--ext', type=str, default='npz',
+ help='Which extension to use for loading')
+ parser.add_argument('--plot-joints', default=False,
+ type=lambda arg: arg.lower() in ['true', '1'],
+ help='The path to the model folder')
+ parser.add_argument('--sample-shape', default=True,
+ dest='sample_shape',
+ type=lambda arg: arg.lower() in ['true', '1'],
+ help='Sample a random shape')
+ parser.add_argument('--sample-expression', default=True,
+ dest='sample_expression',
+ type=lambda arg: arg.lower() in ['true', '1'],
+ help='Sample a random expression')
+ parser.add_argument('--use-face-contour', default=False,
+ type=lambda arg: arg.lower() in ['true', '1'],
+ help='Compute the contour of the face')
+
+ args = parser.parse_args()
+
+ model_folder = osp.expanduser(osp.expandvars(args.model_folder))
+ model_type = args.model_type
+ plot_joints = args.plot_joints
+ use_face_contour = args.use_face_contour
+ gender = args.gender
+ ext = args.ext
+ plotting_module = args.plotting_module
+ num_betas = args.num_betas
+ num_expression_coeffs = args.num_expression_coeffs
+ sample_shape = args.sample_shape
+ sample_expression = args.sample_expression
+
+ main(model_folder, model_type, ext=ext,
+ gender=gender, plot_joints=plot_joints,
+ num_betas=num_betas,
+ num_expression_coeffs=num_expression_coeffs,
+ sample_shape=sample_shape,
+ sample_expression=sample_expression,
+ plotting_module=plotting_module,
+ use_face_contour=use_face_contour)
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/examples/vis_flame_vertices.py b/worldbench/third_party/drivestudio/third_party/smplx/examples/vis_flame_vertices.py
new file mode 100644
index 0000000..b8d6b9b
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/examples/vis_flame_vertices.py
@@ -0,0 +1,92 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+import os.path as osp
+import argparse
+import pickle
+
+import numpy as np
+import torch
+import open3d as o3d
+
+import smplx
+
+
+def main(model_folder, corr_fname, ext='npz',
+ head_color=(0.3, 0.3, 0.6),
+ gender='neutral'):
+
+ head_idxs = np.load(corr_fname)
+
+ model = smplx.create(model_folder, model_type='smplx',
+ gender=gender,
+ ext=ext)
+ betas = torch.zeros([1, 10], dtype=torch.float32)
+ expression = torch.zeros([1, 10], dtype=torch.float32)
+
+ output = model(betas=betas, expression=expression,
+ return_verts=True)
+ vertices = output.vertices.detach().cpu().numpy().squeeze()
+ joints = output.joints.detach().cpu().numpy().squeeze()
+
+ print('Vertices shape =', vertices.shape)
+ print('Joints shape =', joints.shape)
+
+ mesh = o3d.geometry.TriangleMesh()
+ mesh.vertices = o3d.utility.Vector3dVector(vertices)
+ mesh.triangles = o3d.utility.Vector3iVector(model.faces)
+ mesh.compute_vertex_normals()
+
+ colors = np.ones_like(vertices) * [0.3, 0.3, 0.3]
+ colors[head_idxs] = head_color
+
+ mesh.vertex_colors = o3d.utility.Vector3dVector(colors)
+
+ o3d.visualization.draw_geometries([mesh])
+
+
+if __name__ == '__main__':
+ parser = argparse.ArgumentParser(description='SMPL-X Demo')
+
+ parser.add_argument('--model-folder', required=True, type=str,
+ help='The path to the model folder')
+ parser.add_argument('--corr-fname', required=True, type=str,
+ dest='corr_fname',
+ help='Filename with the head correspondences')
+ parser.add_argument('--gender', type=str, default='neutral',
+ help='The gender of the model')
+ parser.add_argument('--ext', type=str, default='npz',
+ help='Which extension to use for loading')
+ parser.add_argument('--head', default='right',
+ choices=['right', 'left'],
+ type=str, help='Which head to plot')
+ parser.add_argument('--head-color', type=float, nargs=3, dest='head_color',
+ default=(0.3, 0.3, 0.6),
+ help='Color for the head vertices')
+
+ args = parser.parse_args()
+
+ model_folder = osp.expanduser(osp.expandvars(args.model_folder))
+ corr_fname = args.corr_fname
+ gender = args.gender
+ ext = args.ext
+ head = args.head
+ head_color = args.head_color
+
+ main(model_folder, corr_fname, ext=ext,
+ head_color=head_color,
+ gender=gender
+ )
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/examples/vis_mano_vertices.py b/worldbench/third_party/drivestudio/third_party/smplx/examples/vis_mano_vertices.py
new file mode 100644
index 0000000..1741542
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/examples/vis_mano_vertices.py
@@ -0,0 +1,99 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+import os.path as osp
+import argparse
+import pickle
+
+import numpy as np
+import torch
+import open3d as o3d
+
+import smplx
+
+
+def main(model_folder, corr_fname, ext='npz',
+ hand_color=(0.3, 0.3, 0.6),
+ gender='neutral', hand='right'):
+
+ with open(corr_fname, 'rb') as f:
+ idxs_data = pickle.load(f)
+ if hand == 'both':
+ hand_idxs = np.concatenate(
+ [idxs_data['left_hand'], idxs_data['right_hand']]
+ )
+ else:
+ hand_idxs = idxs_data[f'{hand}_hand']
+
+ model = smplx.create(model_folder, model_type='smplx',
+ gender=gender,
+ ext=ext)
+ betas = torch.zeros([1, 10], dtype=torch.float32)
+ expression = torch.zeros([1, 10], dtype=torch.float32)
+
+ output = model(betas=betas, expression=expression,
+ return_verts=True)
+ vertices = output.vertices.detach().cpu().numpy().squeeze()
+ joints = output.joints.detach().cpu().numpy().squeeze()
+
+ print('Vertices shape =', vertices.shape)
+ print('Joints shape =', joints.shape)
+
+ mesh = o3d.geometry.TriangleMesh()
+ mesh.vertices = o3d.utility.Vector3dVector(vertices)
+ mesh.triangles = o3d.utility.Vector3iVector(model.faces)
+ mesh.compute_vertex_normals()
+
+ colors = np.ones_like(vertices) * [0.3, 0.3, 0.3]
+ colors[hand_idxs] = hand_color
+
+ mesh.vertex_colors = o3d.utility.Vector3dVector(colors)
+
+ o3d.visualization.draw_geometries([mesh])
+
+
+if __name__ == '__main__':
+ parser = argparse.ArgumentParser(description='SMPL-X Demo')
+
+ parser.add_argument('--model-folder', required=True, type=str,
+ help='The path to the model folder')
+ parser.add_argument('--corr-fname', required=True, type=str,
+ dest='corr_fname',
+ help='Filename with the hand correspondences')
+ parser.add_argument('--gender', type=str, default='neutral',
+ help='The gender of the model')
+ parser.add_argument('--ext', type=str, default='npz',
+ help='Which extension to use for loading')
+ parser.add_argument('--hand', default='right',
+ choices=['right', 'left', 'both'],
+ type=str, help='Which hand to plot')
+ parser.add_argument('--hand-color', type=float, nargs=3, dest='hand_color',
+ default=(0.3, 0.3, 0.6),
+ help='Color for the hand vertices')
+
+ args = parser.parse_args()
+
+ model_folder = osp.expanduser(osp.expandvars(args.model_folder))
+ corr_fname = args.corr_fname
+ gender = args.gender
+ ext = args.ext
+ hand = args.hand
+ hand_color = args.hand_color
+
+ main(model_folder, corr_fname, ext=ext,
+ hand_color=hand_color,
+ gender=gender, hand=hand
+ )
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/optional-requirements.txt b/worldbench/third_party/drivestudio/third_party/smplx/optional-requirements.txt
new file mode 100644
index 0000000..5042729
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/optional-requirements.txt
@@ -0,0 +1,3 @@
+pyrender>=0.1.23
+shapely
+trimesh>=2.37.6
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/requirements.txt b/worldbench/third_party/drivestudio/third_party/smplx/requirements.txt
new file mode 100644
index 0000000..5ae9325
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/requirements.txt
@@ -0,0 +1,3 @@
+numpy>=1.16.2
+torch>=1.0.1.post2
+dataclasses>=0.6
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/setup.py b/worldbench/third_party/drivestudio/third_party/smplx/setup.py
new file mode 100644
index 0000000..d010148
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/setup.py
@@ -0,0 +1,78 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems and the Max Planck Institute for Biological
+# Cybernetics. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+import io
+import os
+
+from setuptools import setup
+
+# Package meta-data.
+NAME = 'smplx'
+DESCRIPTION = 'PyTorch module for loading the SMPLX body model'
+URL = 'http://smpl-x.is.tuebingen.mpg.de'
+EMAIL = 'vassilis.choutas@tuebingen.mpg.de'
+AUTHOR = 'Vassilis Choutas'
+REQUIRES_PYTHON = '>=3.6.0'
+VERSION = '0.1.28'
+
+here = os.path.abspath(os.path.dirname(__file__))
+
+try:
+ FileNotFoundError
+except NameError:
+ FileNotFoundError = IOError
+
+# Import the README and use it as the long-description.
+# Note: this will only work if 'README.md' is present in your MANIFEST.in file!
+try:
+ with io.open(os.path.join(here, 'README.md'), encoding='utf-8') as f:
+ long_description = '\n' + f.read()
+except FileNotFoundError:
+ long_description = DESCRIPTION
+
+# Load the package's __version__.py module as a dictionary.
+about = {}
+if not VERSION:
+ with open(os.path.join(here, NAME, '__version__.py')) as f:
+ exec(f.read(), about)
+else:
+ about['__version__'] = VERSION
+
+pyrender_reqs = ['pyrender>=0.1.23', 'trimesh>=2.37.6', 'shapely']
+matplotlib_reqs = ['matplotlib']
+open3d_reqs = ['open3d-python']
+
+setup(name=NAME,
+ version=about['__version__'],
+ description=DESCRIPTION,
+ long_description=long_description,
+ long_description_content_type='text/markdown',
+ author=AUTHOR,
+ author_email=EMAIL,
+ python_requires=REQUIRES_PYTHON,
+ url=URL,
+ install_requires=[
+ 'numpy>=1.16.2',
+ 'torch>=1.0.1.post2',
+ ],
+ extras_require={
+ 'pyrender': pyrender_reqs,
+ 'open3d': open3d_reqs,
+ 'matplotlib': matplotlib_reqs,
+ 'all': pyrender_reqs + matplotlib_reqs + open3d_reqs
+ },
+ packages=['smplx'])
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/smplx/__init__.py b/worldbench/third_party/drivestudio/third_party/smplx/smplx/__init__.py
new file mode 100644
index 0000000..886949d
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/smplx/__init__.py
@@ -0,0 +1,30 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from .body_models import (
+ create,
+ SMPL,
+ SMPLH,
+ SMPLX,
+ MANO,
+ FLAME,
+ build_layer,
+ SMPLLayer,
+ SMPLHLayer,
+ SMPLXLayer,
+ MANOLayer,
+ FLAMELayer,
+)
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/smplx/body_models.py b/worldbench/third_party/drivestudio/third_party/smplx/smplx/body_models.py
new file mode 100644
index 0000000..54c00cf
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/smplx/body_models.py
@@ -0,0 +1,2529 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from typing import Optional, Dict, Union
+import os
+import os.path as osp
+
+import pickle
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+
+from .lbs import lbs, vertices2landmarks, find_dynamic_lmk_idx_and_bcoords, blend_shapes
+
+from .vertex_ids import vertex_ids as VERTEX_IDS
+from .utils import (
+ Struct,
+ to_np,
+ to_tensor,
+ Tensor,
+ Array,
+ SMPLOutput,
+ SMPLHOutput,
+ SMPLXOutput,
+ MANOOutput,
+ FLAMEOutput,
+ find_joint_kin_chain,
+)
+from .vertex_joint_selector import VertexJointSelector
+from collections import namedtuple
+
+TensorOutput = namedtuple(
+ "TensorOutput",
+ [
+ "vertices",
+ "joints",
+ "betas",
+ "expression",
+ "global_orient",
+ "body_pose",
+ "left_hand_pose",
+ "right_hand_pose",
+ "jaw_pose",
+ "transl",
+ "full_pose",
+ ],
+)
+
+
+class SMPL(nn.Module):
+ NUM_JOINTS = 23
+ NUM_BODY_JOINTS = 23
+ SHAPE_SPACE_DIM = 300
+
+ def __init__(
+ self,
+ model_path: str,
+ kid_template_path: str = "",
+ data_struct: Optional[Struct] = None,
+ create_betas: bool = True,
+ betas: Optional[Tensor] = None,
+ num_betas: int = 10,
+ create_global_orient: bool = True,
+ global_orient: Optional[Tensor] = None,
+ create_body_pose: bool = True,
+ body_pose: Optional[Tensor] = None,
+ create_transl: bool = True,
+ transl: Optional[Tensor] = None,
+ dtype=torch.float32,
+ batch_size: int = 1,
+ joint_mapper=None,
+ gender: str = "neutral",
+ age: str = "adult",
+ vertex_ids: Dict[str, int] = None,
+ v_template: Optional[Union[Tensor, Array]] = None,
+ **kwargs,
+ ) -> None:
+ """SMPL model constructor
+
+ Parameters
+ ----------
+ model_path: str
+ The path to the folder or to the file where the model
+ parameters are stored
+ data_struct: Strct
+ A struct object. If given, then the parameters of the model are
+ read from the object. Otherwise, the model tries to read the
+ parameters from the given `model_path`. (default = None)
+ create_global_orient: bool, optional
+ Flag for creating a member variable for the global orientation
+ of the body. (default = True)
+ global_orient: torch.tensor, optional, Bx3
+ The default value for the global orientation variable.
+ (default = None)
+ create_body_pose: bool, optional
+ Flag for creating a member variable for the pose of the body.
+ (default = True)
+ body_pose: torch.tensor, optional, Bx(Body Joints * 3)
+ The default value for the body pose variable.
+ (default = None)
+ num_betas: int, optional
+ Number of shape components to use
+ (default = 10).
+ create_betas: bool, optional
+ Flag for creating a member variable for the shape space
+ (default = True).
+ betas: torch.tensor, optional, Bx10
+ The default value for the shape member variable.
+ (default = None)
+ create_transl: bool, optional
+ Flag for creating a member variable for the translation
+ of the body. (default = True)
+ transl: torch.tensor, optional, Bx3
+ The default value for the transl variable.
+ (default = None)
+ dtype: torch.dtype, optional
+ The data type for the created variables
+ batch_size: int, optional
+ The batch size used for creating the member variables
+ joint_mapper: object, optional
+ An object that re-maps the joints. Useful if one wants to
+ re-order the SMPL joints to some other convention (e.g. MSCOCO)
+ (default = None)
+ gender: str, optional
+ Which gender to load
+ vertex_ids: dict, optional
+ A dictionary containing the indices of the extra vertices that
+ will be selected
+ """
+
+ self.gender = gender
+ self.age = age
+
+ if data_struct is None:
+ if osp.isdir(model_path):
+ model_fn = "SMPL_{}.{ext}".format(gender.upper(), ext="pkl")
+ smpl_path = os.path.join(model_path, model_fn)
+ else:
+ smpl_path = model_path
+ assert osp.exists(smpl_path), "Path {} does not exist!".format(smpl_path)
+
+ with open(smpl_path, "rb") as smpl_file:
+ data_struct = Struct(**pickle.load(smpl_file, encoding="latin1"))
+
+ super(SMPL, self).__init__()
+ self.batch_size = batch_size
+ shapedirs = data_struct.shapedirs
+ if shapedirs.shape[-1] < self.SHAPE_SPACE_DIM:
+ print(
+ f"WARNING: You are using a {self.name()} model, with only"
+ f" {shapedirs.shape[-1]} shape coefficients.\n"
+ f"num_betas={num_betas}, shapedirs.shape={shapedirs.shape}, "
+ f"self.SHAPE_SPACE_DIM={self.SHAPE_SPACE_DIM}"
+ )
+ num_betas = min(num_betas, shapedirs.shape[-1])
+ else:
+ num_betas = min(num_betas, self.SHAPE_SPACE_DIM)
+
+ if self.age == "kid":
+ v_template_smil = np.load(kid_template_path)
+ v_template_smil -= np.mean(v_template_smil, axis=0)
+ v_template_diff = np.expand_dims(v_template_smil - data_struct.v_template, axis=2)
+ shapedirs = np.concatenate((shapedirs[:, :, :num_betas], v_template_diff), axis=2)
+ num_betas = num_betas + 1
+
+ self._num_betas = num_betas
+ shapedirs = shapedirs[:, :, :num_betas]
+ # The shape components
+ self.register_buffer("shapedirs", to_tensor(to_np(shapedirs), dtype=dtype))
+
+ if vertex_ids is None:
+ # SMPL and SMPL-H share the same topology, so any extra joints can
+ # be drawn from the same place
+ vertex_ids = VERTEX_IDS["smplh"]
+
+ self.dtype = dtype
+
+ self.joint_mapper = joint_mapper
+
+ self.vertex_joint_selector = VertexJointSelector(vertex_ids=vertex_ids, **kwargs)
+
+ self.faces = data_struct.f
+ self.register_buffer(
+ "faces_tensor", to_tensor(to_np(self.faces, dtype=np.int64), dtype=torch.long)
+ )
+
+ if create_betas:
+ if betas is None:
+ default_betas = torch.zeros([batch_size, self.num_betas], dtype=dtype)
+ else:
+ if torch.is_tensor(betas):
+ default_betas = betas.clone().detach()
+ else:
+ default_betas = torch.tensor(betas, dtype=dtype)
+
+ self.register_parameter("betas", nn.Parameter(default_betas, requires_grad=True))
+
+ # The tensor that contains the global rotation of the model
+ # It is separated from the pose of the joints in case we wish to
+ # optimize only over one of them
+ if create_global_orient:
+ if global_orient is None:
+ default_global_orient = torch.zeros([batch_size, 3], dtype=dtype)
+ else:
+ if torch.is_tensor(global_orient):
+ default_global_orient = global_orient.clone().detach()
+ else:
+ default_global_orient = torch.tensor(global_orient, dtype=dtype)
+
+ global_orient = nn.Parameter(default_global_orient, requires_grad=True)
+ self.register_parameter("global_orient", global_orient)
+
+ if create_body_pose:
+ if body_pose is None:
+ default_body_pose = torch.zeros([batch_size, self.NUM_BODY_JOINTS * 3], dtype=dtype)
+ else:
+ if torch.is_tensor(body_pose):
+ default_body_pose = body_pose.clone().detach()
+ else:
+ default_body_pose = torch.tensor(body_pose, dtype=dtype)
+ self.register_parameter(
+ "body_pose", nn.Parameter(default_body_pose, requires_grad=True)
+ )
+
+ if create_transl:
+ if transl is None:
+ default_transl = torch.zeros([batch_size, 3], dtype=dtype, requires_grad=True)
+ else:
+ default_transl = torch.tensor(transl, dtype=dtype)
+ self.register_parameter("transl", nn.Parameter(default_transl, requires_grad=True))
+
+ if v_template is None:
+ v_template = data_struct.v_template
+ if not torch.is_tensor(v_template):
+ v_template = to_tensor(to_np(v_template), dtype=dtype)
+ # The vertices of the template model
+ self.register_buffer("v_template", v_template)
+
+ j_regressor = to_tensor(to_np(data_struct.J_regressor), dtype=dtype)
+ self.register_buffer("J_regressor", j_regressor)
+
+ # Pose blend shape basis: 6890 x 3 x 207, reshaped to 6890*3 x 207
+ num_pose_basis = data_struct.posedirs.shape[-1]
+ # 207 x 20670
+ posedirs = np.reshape(data_struct.posedirs, [-1, num_pose_basis]).T
+ self.register_buffer("posedirs", to_tensor(to_np(posedirs), dtype=dtype))
+
+ # indices of parents for each joints
+ parents = to_tensor(to_np(data_struct.kintree_table[0])).long()
+ parents[0] = -1
+ self.register_buffer("parents", parents)
+
+ lbs_weights = to_tensor(to_np(data_struct.weights), dtype=dtype)
+ self.register_buffer("lbs_weights", lbs_weights)
+
+ @property
+ def num_betas(self):
+ return self._num_betas
+
+ @property
+ def num_expression_coeffs(self):
+ return 0
+
+ def create_mean_pose(self, data_struct) -> Tensor:
+ pass
+
+ def name(self) -> str:
+ return "SMPL"
+
+ @torch.no_grad()
+ def reset_params(self, **params_dict) -> None:
+ for param_name, param in self.named_parameters():
+ if param_name in params_dict:
+ param[:] = torch.tensor(params_dict[param_name])
+ else:
+ param.fill_(0)
+
+ def get_num_verts(self) -> int:
+ return self.v_template.shape[0]
+
+ def get_num_faces(self) -> int:
+ return self.faces.shape[0]
+
+ def extra_repr(self) -> str:
+ msg = [
+ f"Gender: {self.gender.upper()}",
+ f"Number of joints: {self.J_regressor.shape[0]}",
+ f"Betas: {self.num_betas}",
+ ]
+ return "\n".join(msg)
+
+ def forward_shape(
+ self,
+ betas: Optional[Tensor] = None,
+ ) -> SMPLOutput:
+ betas = betas if betas is not None else self.betas
+ v_shaped = self.v_template + blend_shapes(betas, self.shapedirs)
+ return SMPLOutput(vertices=v_shaped, betas=betas, v_shaped=v_shaped)
+
+ def forward(
+ self,
+ betas: Optional[Tensor] = None,
+ body_pose: Optional[Tensor] = None,
+ global_orient: Optional[Tensor] = None,
+ transl: Optional[Tensor] = None,
+ return_verts=True,
+ return_full_pose: bool = False,
+ pose2rot: bool = True,
+ **kwargs,
+ ) -> SMPLOutput:
+ """Forward pass for the SMPL model
+
+ Parameters
+ ----------
+ global_orient: torch.tensor, optional, shape Bx3
+ If given, ignore the member variable and use it as the global
+ rotation of the body. Useful if someone wishes to predicts this
+ with an external model. (default=None)
+ betas: torch.tensor, optional, shape BxN_b
+ If given, ignore the member variable `betas` and use it
+ instead. For example, it can used if shape parameters
+ `betas` are predicted from some external model.
+ (default=None)
+ body_pose: torch.tensor, optional, shape Bx(J*3)
+ If given, ignore the member variable `body_pose` and use it
+ instead. For example, it can used if someone predicts the
+ pose of the body joints are predicted from some external model.
+ It should be a tensor that contains joint rotations in
+ axis-angle format. (default=None)
+ transl: torch.tensor, optional, shape Bx3
+ If given, ignore the member variable `transl` and use it
+ instead. For example, it can used if the translation
+ `transl` is predicted from some external model.
+ (default=None)
+ return_verts: bool, optional
+ Return the vertices. (default=True)
+ return_full_pose: bool, optional
+ Returns the full axis-angle pose vector (default=False)
+
+ Returns
+ -------
+ """
+ # If no shape and pose parameters are passed along, then use the
+ # ones from the module
+ global_orient = global_orient if global_orient is not None else self.global_orient
+ body_pose = body_pose if body_pose is not None else self.body_pose
+ betas = betas if betas is not None else self.betas
+
+ apply_trans = transl is not None or hasattr(self, "transl")
+ if transl is None and hasattr(self, "transl"):
+ transl = self.transl
+
+ full_pose = torch.cat([global_orient, body_pose], dim=1)
+
+ batch_size = max(betas.shape[0], global_orient.shape[0], body_pose.shape[0])
+
+ if betas.shape[0] != batch_size:
+ num_repeats = int(batch_size / betas.shape[0])
+ betas = betas.expand(num_repeats, -1)
+
+ vertices, joints = lbs(
+ betas,
+ full_pose,
+ self.v_template,
+ self.shapedirs,
+ self.posedirs,
+ self.J_regressor,
+ self.parents,
+ self.lbs_weights,
+ pose2rot=pose2rot,
+ )
+
+ joints = self.vertex_joint_selector(vertices, joints)
+ # Map the joints to the current dataset
+ if self.joint_mapper is not None:
+ joints = self.joint_mapper(joints)
+
+ if apply_trans:
+ joints += transl.unsqueeze(dim=1)
+ vertices += transl.unsqueeze(dim=1)
+
+ output = SMPLOutput(
+ vertices=vertices if return_verts else None,
+ global_orient=global_orient,
+ body_pose=body_pose,
+ joints=joints,
+ betas=betas,
+ full_pose=full_pose if return_full_pose else None,
+ )
+
+ return output
+
+
+class SMPLLayer(SMPL):
+ def __init__(self, *args, **kwargs) -> None:
+ # Just create a SMPL module without any member variables
+ super(SMPLLayer, self).__init__(
+ create_body_pose=False,
+ create_betas=False,
+ create_global_orient=False,
+ create_transl=False,
+ *args,
+ **kwargs,
+ )
+
+ def forward(
+ self,
+ betas: Optional[Tensor] = None,
+ body_pose: Optional[Tensor] = None,
+ global_orient: Optional[Tensor] = None,
+ transl: Optional[Tensor] = None,
+ return_verts=True,
+ return_full_pose: bool = False,
+ pose2rot: bool = True,
+ **kwargs,
+ ) -> SMPLOutput:
+ """Forward pass for the SMPL model
+
+ Parameters
+ ----------
+ global_orient: torch.tensor, optional, shape Bx3x3
+ Global rotation of the body. Useful if someone wishes to
+ predicts this with an external model. It is expected to be in
+ rotation matrix format. (default=None)
+ betas: torch.tensor, optional, shape BxN_b
+ Shape parameters. For example, it can used if shape parameters
+ `betas` are predicted from some external model.
+ (default=None)
+ body_pose: torch.tensor, optional, shape BxJx3x3
+ Body pose. For example, it can used if someone predicts the
+ pose of the body joints are predicted from some external model.
+ It should be a tensor that contains joint rotations in
+ rotation matrix format. (default=None)
+ transl: torch.tensor, optional, shape Bx3
+ Translation vector of the body.
+ For example, it can used if the translation
+ `transl` is predicted from some external model.
+ (default=None)
+ return_verts: bool, optional
+ Return the vertices. (default=True)
+ return_full_pose: bool, optional
+ Returns the full axis-angle pose vector (default=False)
+
+ Returns
+ -------
+ """
+ model_vars = [betas, global_orient, body_pose, transl]
+ batch_size = 1
+ for var in model_vars:
+ if var is None:
+ continue
+ batch_size = max(batch_size, len(var))
+ device, dtype = self.shapedirs.device, self.shapedirs.dtype
+ if global_orient is None:
+ global_orient = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, -1, -1, -1)
+ .contiguous()
+ )
+ if body_pose is None:
+ body_pose = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, self.NUM_BODY_JOINTS, -1, -1)
+ .contiguous()
+ )
+ if betas is None:
+ betas = torch.zeros([batch_size, self.num_betas], dtype=dtype, device=device)
+ if transl is None:
+ transl = torch.zeros([batch_size, 3], dtype=dtype, device=device)
+ full_pose = torch.cat(
+ [global_orient.reshape(-1, 1, 3, 3), body_pose.reshape(-1, self.NUM_BODY_JOINTS, 3, 3)],
+ dim=1,
+ )
+
+ vertices, joints, T, v_posed, A, J= lbs(
+ betas,
+ full_pose,
+ self.v_template,
+ self.shapedirs,
+ self.posedirs,
+ self.J_regressor,
+ self.parents,
+ self.lbs_weights,
+ pose2rot=False,
+ return_T=True,
+ return_posed_v=True,
+ return_A=True,
+ return_J=True,
+ )
+
+ joints = self.vertex_joint_selector(vertices, joints)
+ # Map the joints to the current dataset
+ if self.joint_mapper is not None:
+ joints = self.joint_mapper(joints)
+
+ if transl is not None:
+ joints += transl.unsqueeze(dim=1)
+ vertices += transl.unsqueeze(dim=1)
+
+ output = SMPLOutput(
+ vertices=vertices if return_verts else None,
+ global_orient=global_orient,
+ body_pose=body_pose,
+ joints=joints,
+ betas=betas,
+ full_pose=full_pose if return_full_pose else None,
+ T=T,
+ v_posed=v_posed, # the canonical vtx
+ A=A,
+ J=J,
+ )
+
+ return output
+
+
+class SMPLH(SMPL):
+ # The hand joints are replaced by MANO
+ NUM_BODY_JOINTS = SMPL.NUM_JOINTS - 2
+ NUM_HAND_JOINTS = 15
+ NUM_JOINTS = NUM_BODY_JOINTS + 2 * NUM_HAND_JOINTS
+
+ def __init__(
+ self,
+ model_path,
+ kid_template_path: str = "",
+ data_struct: Optional[Struct] = None,
+ create_left_hand_pose: bool = True,
+ left_hand_pose: Optional[Tensor] = None,
+ create_right_hand_pose: bool = True,
+ right_hand_pose: Optional[Tensor] = None,
+ use_pca: bool = True,
+ num_pca_comps: int = 6,
+ num_betas=16,
+ flat_hand_mean: bool = False,
+ batch_size: int = 1,
+ gender: str = "neutral",
+ age: str = "adult",
+ dtype=torch.float32,
+ vertex_ids=None,
+ use_compressed: bool = True,
+ ext: str = "pkl",
+ **kwargs,
+ ) -> None:
+ """SMPLH model constructor
+
+ Parameters
+ ----------
+ model_path: str
+ The path to the folder or to the file where the model
+ parameters are stored
+ data_struct: Strct
+ A struct object. If given, then the parameters of the model are
+ read from the object. Otherwise, the model tries to read the
+ parameters from the given `model_path`. (default = None)
+ create_left_hand_pose: bool, optional
+ Flag for creating a member variable for the pose of the left
+ hand. (default = True)
+ left_hand_pose: torch.tensor, optional, BxP
+ The default value for the left hand pose member variable.
+ (default = None)
+ create_right_hand_pose: bool, optional
+ Flag for creating a member variable for the pose of the right
+ hand. (default = True)
+ right_hand_pose: torch.tensor, optional, BxP
+ The default value for the right hand pose member variable.
+ (default = None)
+ num_pca_comps: int, optional
+ The number of PCA components to use for each hand.
+ (default = 6)
+ flat_hand_mean: bool, optional
+ If False, then the pose of the hand is initialized to False.
+ batch_size: int, optional
+ The batch size used for creating the member variables
+ gender: str, optional
+ Which gender to load
+ dtype: torch.dtype, optional
+ The data type for the created variables
+ vertex_ids: dict, optional
+ A dictionary containing the indices of the extra vertices that
+ will be selected
+ """
+
+ self.num_pca_comps = num_pca_comps
+ # If no data structure is passed, then load the data from the given
+ # model folder
+ if data_struct is None:
+ # Load the model
+ if osp.isdir(model_path):
+ model_fn = "SMPLH_{}.{ext}".format(gender.upper(), ext=ext)
+ smplh_path = os.path.join(model_path, model_fn)
+ else:
+ smplh_path = model_path
+ assert osp.exists(smplh_path), "Path {} does not exist!".format(smplh_path)
+
+ if ext == "pkl":
+ with open(smplh_path, "rb") as smplh_file:
+ model_data = pickle.load(smplh_file, encoding="latin1")
+ elif ext == "npz":
+ model_data = np.load(smplh_path, allow_pickle=True)
+ else:
+ raise ValueError("Unknown extension: {}".format(ext))
+ data_struct = Struct(**model_data)
+
+ if vertex_ids is None:
+ vertex_ids = VERTEX_IDS["smplh"]
+
+ super(SMPLH, self).__init__(
+ model_path=model_path,
+ kid_template_path=kid_template_path,
+ data_struct=data_struct,
+ num_betas=num_betas,
+ batch_size=batch_size,
+ vertex_ids=vertex_ids,
+ gender=gender,
+ age=age,
+ use_compressed=use_compressed,
+ dtype=dtype,
+ ext=ext,
+ **kwargs,
+ )
+
+ self.use_pca = use_pca
+ self.num_pca_comps = num_pca_comps
+ self.flat_hand_mean = flat_hand_mean
+
+ left_hand_components = data_struct.hands_componentsl[:num_pca_comps]
+ right_hand_components = data_struct.hands_componentsr[:num_pca_comps]
+
+ self.np_left_hand_components = left_hand_components
+ self.np_right_hand_components = right_hand_components
+ if self.use_pca:
+ self.register_buffer(
+ "left_hand_components", torch.tensor(left_hand_components, dtype=dtype)
+ )
+ self.register_buffer(
+ "right_hand_components", torch.tensor(right_hand_components, dtype=dtype)
+ )
+
+ if self.flat_hand_mean:
+ left_hand_mean = np.zeros_like(data_struct.hands_meanl)
+ else:
+ left_hand_mean = data_struct.hands_meanl
+
+ if self.flat_hand_mean:
+ right_hand_mean = np.zeros_like(data_struct.hands_meanr)
+ else:
+ right_hand_mean = data_struct.hands_meanr
+
+ self.register_buffer("left_hand_mean", to_tensor(left_hand_mean, dtype=self.dtype))
+ self.register_buffer("right_hand_mean", to_tensor(right_hand_mean, dtype=self.dtype))
+
+ # Create the buffers for the pose of the left hand
+ hand_pose_dim = num_pca_comps if use_pca else 3 * self.NUM_HAND_JOINTS
+ if create_left_hand_pose:
+ if left_hand_pose is None:
+ default_lhand_pose = torch.zeros([batch_size, hand_pose_dim], dtype=dtype)
+ else:
+ default_lhand_pose = torch.tensor(left_hand_pose, dtype=dtype)
+
+ left_hand_pose_param = nn.Parameter(default_lhand_pose, requires_grad=True)
+ self.register_parameter("left_hand_pose", left_hand_pose_param)
+
+ if create_right_hand_pose:
+ if right_hand_pose is None:
+ default_rhand_pose = torch.zeros([batch_size, hand_pose_dim], dtype=dtype)
+ else:
+ default_rhand_pose = torch.tensor(right_hand_pose, dtype=dtype)
+
+ right_hand_pose_param = nn.Parameter(default_rhand_pose, requires_grad=True)
+ self.register_parameter("right_hand_pose", right_hand_pose_param)
+
+ # Create the buffer for the mean pose.
+ pose_mean_tensor = self.create_mean_pose(data_struct, flat_hand_mean=flat_hand_mean)
+ if not torch.is_tensor(pose_mean_tensor):
+ pose_mean_tensor = torch.tensor(pose_mean_tensor, dtype=dtype)
+ self.register_buffer("pose_mean", pose_mean_tensor)
+
+ def create_mean_pose(self, data_struct, flat_hand_mean=False):
+ # Create the array for the mean pose. If flat_hand is false, then use
+ # the mean that is given by the data, rather than the flat open hand
+ global_orient_mean = torch.zeros([3], dtype=self.dtype)
+ body_pose_mean = torch.zeros([self.NUM_BODY_JOINTS * 3], dtype=self.dtype)
+
+ pose_mean = torch.cat(
+ [global_orient_mean, body_pose_mean, self.left_hand_mean, self.right_hand_mean], dim=0
+ )
+ return pose_mean
+
+ def name(self) -> str:
+ return "SMPL+H"
+
+ def extra_repr(self):
+ msg = super(SMPLH, self).extra_repr()
+ msg = [msg]
+ if self.use_pca:
+ msg.append(f"Number of PCA components: {self.num_pca_comps}")
+ msg.append(f"Flat hand mean: {self.flat_hand_mean}")
+ return "\n".join(msg)
+
+ def forward(
+ self,
+ betas: Optional[Tensor] = None,
+ global_orient: Optional[Tensor] = None,
+ body_pose: Optional[Tensor] = None,
+ left_hand_pose: Optional[Tensor] = None,
+ right_hand_pose: Optional[Tensor] = None,
+ transl: Optional[Tensor] = None,
+ return_verts: bool = True,
+ return_full_pose: bool = False,
+ pose2rot: bool = True,
+ **kwargs,
+ ) -> SMPLHOutput:
+ """ """
+
+ # If no shape and pose parameters are passed along, then use the
+ # ones from the module
+ global_orient = global_orient if global_orient is not None else self.global_orient
+ body_pose = body_pose if body_pose is not None else self.body_pose
+ betas = betas if betas is not None else self.betas
+ left_hand_pose = left_hand_pose if left_hand_pose is not None else self.left_hand_pose
+ right_hand_pose = right_hand_pose if right_hand_pose is not None else self.right_hand_pose
+
+ apply_trans = transl is not None or hasattr(self, "transl")
+ if transl is None:
+ if hasattr(self, "transl"):
+ transl = self.transl
+
+ if self.use_pca:
+ left_hand_pose = torch.einsum("bi,ij->bj", [left_hand_pose, self.left_hand_components])
+ right_hand_pose = torch.einsum(
+ "bi,ij->bj", [right_hand_pose, self.right_hand_components]
+ )
+
+ full_pose = torch.cat([global_orient, body_pose, left_hand_pose, right_hand_pose], dim=1)
+ full_pose += self.pose_mean
+
+ vertices, joints = lbs(
+ betas,
+ full_pose,
+ self.v_template,
+ self.shapedirs,
+ self.posedirs,
+ self.J_regressor,
+ self.parents,
+ self.lbs_weights,
+ pose2rot=pose2rot,
+ )
+
+ # Add any extra joints that might be needed
+ joints = self.vertex_joint_selector(vertices, joints)
+ if self.joint_mapper is not None:
+ joints = self.joint_mapper(joints)
+
+ if apply_trans:
+ joints += transl.unsqueeze(dim=1)
+ vertices += transl.unsqueeze(dim=1)
+
+ output = SMPLHOutput(
+ vertices=vertices if return_verts else None,
+ joints=joints,
+ betas=betas,
+ global_orient=global_orient,
+ body_pose=body_pose,
+ left_hand_pose=left_hand_pose,
+ right_hand_pose=right_hand_pose,
+ full_pose=full_pose if return_full_pose else None,
+ )
+
+ return output
+
+
+class SMPLHLayer(SMPLH):
+ def __init__(self, *args, **kwargs) -> None:
+ """SMPL+H as a layer model constructor"""
+ super(SMPLHLayer, self).__init__(
+ create_global_orient=False,
+ create_body_pose=False,
+ create_left_hand_pose=False,
+ create_right_hand_pose=False,
+ create_betas=False,
+ create_transl=False,
+ *args,
+ **kwargs,
+ )
+
+ def forward(
+ self,
+ betas: Optional[Tensor] = None,
+ global_orient: Optional[Tensor] = None,
+ body_pose: Optional[Tensor] = None,
+ left_hand_pose: Optional[Tensor] = None,
+ right_hand_pose: Optional[Tensor] = None,
+ transl: Optional[Tensor] = None,
+ return_verts: bool = True,
+ return_full_pose: bool = False,
+ pose2rot: bool = True,
+ **kwargs,
+ ) -> SMPLHOutput:
+ """Forward pass for the SMPL+H model
+
+ Parameters
+ ----------
+ global_orient: torch.tensor, optional, shape Bx3x3
+ Global rotation of the body. Useful if someone wishes to
+ predicts this with an external model. It is expected to be in
+ rotation matrix format. (default=None)
+ betas: torch.tensor, optional, shape BxN_b
+ Shape parameters. For example, it can used if shape parameters
+ `betas` are predicted from some external model.
+ (default=None)
+ body_pose: torch.tensor, optional, shape BxJx3x3
+ If given, ignore the member variable `body_pose` and use it
+ instead. For example, it can used if someone predicts the
+ pose of the body joints are predicted from some external model.
+ It should be a tensor that contains joint rotations in
+ rotation matrix format. (default=None)
+ left_hand_pose: torch.tensor, optional, shape Bx15x3x3
+ If given, contains the pose of the left hand.
+ It should be a tensor that contains joint rotations in
+ rotation matrix format. (default=None)
+ right_hand_pose: torch.tensor, optional, shape Bx15x3x3
+ If given, contains the pose of the right hand.
+ It should be a tensor that contains joint rotations in
+ rotation matrix format. (default=None)
+ transl: torch.tensor, optional, shape Bx3
+ Translation vector of the body.
+ For example, it can used if the translation
+ `transl` is predicted from some external model.
+ (default=None)
+ return_verts: bool, optional
+ Return the vertices. (default=True)
+ return_full_pose: bool, optional
+ Returns the full axis-angle pose vector (default=False)
+
+ Returns
+ -------
+ """
+ model_vars = [betas, global_orient, body_pose, transl, left_hand_pose, right_hand_pose]
+ batch_size = 1
+ for var in model_vars:
+ if var is None:
+ continue
+ batch_size = max(batch_size, len(var))
+ device, dtype = self.shapedirs.device, self.shapedirs.dtype
+ if global_orient is None:
+ global_orient = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, -1, -1, -1)
+ .contiguous()
+ )
+ if body_pose is None:
+ body_pose = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, 21, -1, -1)
+ .contiguous()
+ )
+ if left_hand_pose is None:
+ left_hand_pose = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, 15, -1, -1)
+ .contiguous()
+ )
+ if right_hand_pose is None:
+ right_hand_pose = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, 15, -1, -1)
+ .contiguous()
+ )
+ if betas is None:
+ betas = torch.zeros([batch_size, self.num_betas], dtype=dtype, device=device)
+ if transl is None:
+ transl = torch.zeros([batch_size, 3], dtype=dtype, device=device)
+
+ # Concatenate all pose vectors
+ full_pose = torch.cat(
+ [
+ global_orient.reshape(-1, 1, 3, 3),
+ body_pose.reshape(-1, self.NUM_BODY_JOINTS, 3, 3),
+ left_hand_pose.reshape(-1, self.NUM_HAND_JOINTS, 3, 3),
+ right_hand_pose.reshape(-1, self.NUM_HAND_JOINTS, 3, 3),
+ ],
+ dim=1,
+ )
+
+ vertices, joints = lbs(
+ betas,
+ full_pose,
+ self.v_template,
+ self.shapedirs,
+ self.posedirs,
+ self.J_regressor,
+ self.parents,
+ self.lbs_weights,
+ pose2rot=False,
+ )
+
+ # Add any extra joints that might be needed
+ joints = self.vertex_joint_selector(vertices, joints)
+ if self.joint_mapper is not None:
+ joints = self.joint_mapper(joints)
+
+ if transl is not None:
+ joints += transl.unsqueeze(dim=1)
+ vertices += transl.unsqueeze(dim=1)
+
+ output = SMPLHOutput(
+ vertices=vertices if return_verts else None,
+ joints=joints,
+ betas=betas,
+ global_orient=global_orient,
+ body_pose=body_pose,
+ left_hand_pose=left_hand_pose,
+ right_hand_pose=right_hand_pose,
+ full_pose=full_pose if return_full_pose else None,
+ )
+
+ return output
+
+
+class SMPLX(SMPLH):
+ """
+ SMPL-X (SMPL eXpressive) is a unified body model, with shape parameters
+ trained jointly for the face, hands and body.
+ SMPL-X uses standard vertex based linear blend skinning with learned
+ corrective blend shapes, has N=10475 vertices and K=54 joints,
+ which includes joints for the neck, jaw, eyeballs and fingers.
+ """
+
+ NUM_BODY_JOINTS = SMPLH.NUM_BODY_JOINTS
+ NUM_HAND_JOINTS = 15
+ NUM_FACE_JOINTS = 3
+ NUM_JOINTS = NUM_BODY_JOINTS + 2 * NUM_HAND_JOINTS + NUM_FACE_JOINTS
+ EXPRESSION_SPACE_DIM = 100
+ NECK_IDX = 12
+
+ def __init__(
+ self,
+ model_path: str,
+ kid_template_path: str = "",
+ num_expression_coeffs: int = 10,
+ create_expression: bool = True,
+ expression: Optional[Tensor] = None,
+ create_jaw_pose: bool = True,
+ jaw_pose: Optional[Tensor] = None,
+ create_leye_pose: bool = True,
+ leye_pose: Optional[Tensor] = None,
+ create_reye_pose=True,
+ reye_pose: Optional[Tensor] = None,
+ use_face_contour: bool = False,
+ batch_size: int = 1,
+ gender: str = "neutral",
+ age: str = "adult",
+ dtype=torch.float32,
+ ext: str = "npz",
+ **kwargs,
+ ) -> None:
+ """SMPLX model constructor
+
+ Parameters
+ ----------
+ model_path: str
+ The path to the folder or to the file where the model
+ parameters are stored
+ num_expression_coeffs: int, optional
+ Number of expression components to use
+ (default = 10).
+ create_expression: bool, optional
+ Flag for creating a member variable for the expression space
+ (default = True).
+ expression: torch.tensor, optional, Bx10
+ The default value for the expression member variable.
+ (default = None)
+ create_jaw_pose: bool, optional
+ Flag for creating a member variable for the jaw pose.
+ (default = False)
+ jaw_pose: torch.tensor, optional, Bx3
+ The default value for the jaw pose variable.
+ (default = None)
+ create_leye_pose: bool, optional
+ Flag for creating a member variable for the left eye pose.
+ (default = False)
+ leye_pose: torch.tensor, optional, Bx10
+ The default value for the left eye pose variable.
+ (default = None)
+ create_reye_pose: bool, optional
+ Flag for creating a member variable for the right eye pose.
+ (default = False)
+ reye_pose: torch.tensor, optional, Bx10
+ The default value for the right eye pose variable.
+ (default = None)
+ use_face_contour: bool, optional
+ Whether to compute the keypoints that form the facial contour
+ batch_size: int, optional
+ The batch size used for creating the member variables
+ gender: str, optional
+ Which gender to load
+ dtype: torch.dtype
+ The data type for the created variables
+ """
+
+ # Load the model
+ if osp.isdir(model_path):
+ model_fn = "SMPLX_{}.{ext}".format(gender.upper(), ext=ext)
+ smplx_path = os.path.join(model_path, model_fn)
+ else:
+ smplx_path = model_path
+ assert osp.exists(smplx_path), "Path {} does not exist!".format(smplx_path)
+
+ if ext == "pkl":
+ with open(smplx_path, "rb") as smplx_file:
+ model_data = pickle.load(smplx_file, encoding="latin1")
+ elif ext == "npz":
+ model_data = np.load(smplx_path, allow_pickle=True)
+ else:
+ raise ValueError("Unknown extension: {}".format(ext))
+
+ data_struct = Struct(**model_data)
+
+ super(SMPLX, self).__init__(
+ model_path=model_path,
+ kid_template_path=kid_template_path,
+ data_struct=data_struct,
+ dtype=dtype,
+ batch_size=batch_size,
+ vertex_ids=VERTEX_IDS["smplx"],
+ gender=gender,
+ age=age,
+ ext=ext,
+ **kwargs,
+ )
+
+ lmk_faces_idx = data_struct.lmk_faces_idx
+ self.register_buffer("lmk_faces_idx", torch.tensor(lmk_faces_idx, dtype=torch.long))
+ lmk_bary_coords = data_struct.lmk_bary_coords
+ self.register_buffer("lmk_bary_coords", torch.tensor(lmk_bary_coords, dtype=dtype))
+
+ self.use_face_contour = use_face_contour
+ if self.use_face_contour:
+ dynamic_lmk_faces_idx = data_struct.dynamic_lmk_faces_idx
+ dynamic_lmk_faces_idx = torch.tensor(dynamic_lmk_faces_idx, dtype=torch.long)
+ self.register_buffer("dynamic_lmk_faces_idx", dynamic_lmk_faces_idx)
+
+ dynamic_lmk_bary_coords = data_struct.dynamic_lmk_bary_coords
+ dynamic_lmk_bary_coords = torch.tensor(dynamic_lmk_bary_coords, dtype=dtype)
+ self.register_buffer("dynamic_lmk_bary_coords", dynamic_lmk_bary_coords)
+
+ neck_kin_chain = find_joint_kin_chain(self.NECK_IDX, self.parents)
+ self.register_buffer("neck_kin_chain", torch.tensor(neck_kin_chain, dtype=torch.long))
+
+ if create_jaw_pose:
+ if jaw_pose is None:
+ default_jaw_pose = torch.zeros([batch_size, 3], dtype=dtype)
+ else:
+ default_jaw_pose = torch.tensor(jaw_pose, dtype=dtype)
+ jaw_pose_param = nn.Parameter(default_jaw_pose, requires_grad=True)
+ self.register_parameter("jaw_pose", jaw_pose_param)
+
+ if create_leye_pose:
+ if leye_pose is None:
+ default_leye_pose = torch.zeros([batch_size, 3], dtype=dtype)
+ else:
+ default_leye_pose = torch.tensor(leye_pose, dtype=dtype)
+ leye_pose_param = nn.Parameter(default_leye_pose, requires_grad=True)
+ self.register_parameter("leye_pose", leye_pose_param)
+
+ if create_reye_pose:
+ if reye_pose is None:
+ default_reye_pose = torch.zeros([batch_size, 3], dtype=dtype)
+ else:
+ default_reye_pose = torch.tensor(reye_pose, dtype=dtype)
+ reye_pose_param = nn.Parameter(default_reye_pose, requires_grad=True)
+ self.register_parameter("reye_pose", reye_pose_param)
+
+ shapedirs = data_struct.shapedirs
+ if len(shapedirs.shape) < 3:
+ shapedirs = shapedirs[:, :, None]
+ if shapedirs.shape[-1] < self.SHAPE_SPACE_DIM + self.EXPRESSION_SPACE_DIM:
+ print(
+ f"WARNING: You are using a {self.name()} model, with only"
+ " 10 shape and 10 expression coefficients."
+ )
+ expr_start_idx = 10
+ expr_end_idx = 20
+ num_expression_coeffs = min(num_expression_coeffs, 10)
+ else:
+ expr_start_idx = self.SHAPE_SPACE_DIM
+ expr_end_idx = self.SHAPE_SPACE_DIM + num_expression_coeffs
+ num_expression_coeffs = min(num_expression_coeffs, self.EXPRESSION_SPACE_DIM)
+
+ self._num_expression_coeffs = num_expression_coeffs
+
+ expr_dirs = shapedirs[:, :, expr_start_idx:expr_end_idx]
+ self.register_buffer("expr_dirs", to_tensor(to_np(expr_dirs), dtype=dtype))
+
+ if create_expression:
+ if expression is None:
+ default_expression = torch.zeros(
+ [batch_size, self.num_expression_coeffs], dtype=dtype
+ )
+ else:
+ default_expression = torch.tensor(expression, dtype=dtype)
+ expression_param = nn.Parameter(default_expression, requires_grad=True)
+ self.register_parameter("expression", expression_param)
+
+ def name(self) -> str:
+ return "SMPL-X"
+
+ @property
+ def num_expression_coeffs(self):
+ return self._num_expression_coeffs
+
+ def create_mean_pose(self, data_struct, flat_hand_mean=False):
+ # Create the array for the mean pose. If flat_hand is false, then use
+ # the mean that is given by the data, rather than the flat open hand
+ global_orient_mean = torch.zeros([3], dtype=self.dtype)
+ body_pose_mean = torch.zeros([self.NUM_BODY_JOINTS * 3], dtype=self.dtype)
+ jaw_pose_mean = torch.zeros([3], dtype=self.dtype)
+ leye_pose_mean = torch.zeros([3], dtype=self.dtype)
+ reye_pose_mean = torch.zeros([3], dtype=self.dtype)
+
+ pose_mean = np.concatenate(
+ [
+ global_orient_mean,
+ body_pose_mean,
+ jaw_pose_mean,
+ leye_pose_mean,
+ reye_pose_mean,
+ self.left_hand_mean,
+ self.right_hand_mean,
+ ],
+ axis=0,
+ )
+
+ return pose_mean
+
+ def extra_repr(self):
+ msg = super(SMPLX, self).extra_repr()
+ msg = [msg, f"Number of Expression Coefficients: {self.num_expression_coeffs}"]
+ return "\n".join(msg)
+
+ def forward(
+ self,
+ betas: Optional[Tensor] = None,
+ global_orient: Optional[Tensor] = None,
+ body_pose: Optional[Tensor] = None,
+ left_hand_pose: Optional[Tensor] = None,
+ right_hand_pose: Optional[Tensor] = None,
+ transl: Optional[Tensor] = None,
+ expression: Optional[Tensor] = None,
+ jaw_pose: Optional[Tensor] = None,
+ leye_pose: Optional[Tensor] = None,
+ reye_pose: Optional[Tensor] = None,
+ return_verts: bool = True,
+ return_full_pose: bool = False,
+ pose2rot: bool = True,
+ return_shaped: bool = True,
+ **kwargs,
+ ) -> SMPLXOutput:
+ """
+ Forward pass for the SMPLX model
+
+ Parameters
+ ----------
+ global_orient: torch.tensor, optional, shape Bx3
+ If given, ignore the member variable and use it as the global
+ rotation of the body. Useful if someone wishes to predicts this
+ with an external model. (default=None)
+ betas: torch.tensor, optional, shape BxN_b
+ If given, ignore the member variable `betas` and use it
+ instead. For example, it can used if shape parameters
+ `betas` are predicted from some external model.
+ (default=None)
+ expression: torch.tensor, optional, shape BxN_e
+ If given, ignore the member variable `expression` and use it
+ instead. For example, it can used if expression parameters
+ `expression` are predicted from some external model.
+ body_pose: torch.tensor, optional, shape Bx(J*3)
+ If given, ignore the member variable `body_pose` and use it
+ instead. For example, it can used if someone predicts the
+ pose of the body joints are predicted from some external model.
+ It should be a tensor that contains joint rotations in
+ axis-angle format. (default=None)
+ left_hand_pose: torch.tensor, optional, shape BxP
+ If given, ignore the member variable `left_hand_pose` and
+ use this instead. It should either contain PCA coefficients or
+ joint rotations in axis-angle format.
+ right_hand_pose: torch.tensor, optional, shape BxP
+ If given, ignore the member variable `right_hand_pose` and
+ use this instead. It should either contain PCA coefficients or
+ joint rotations in axis-angle format.
+ jaw_pose: torch.tensor, optional, shape Bx3
+ If given, ignore the member variable `jaw_pose` and
+ use this instead. It should either joint rotations in
+ axis-angle format.
+ transl: torch.tensor, optional, shape Bx3
+ If given, ignore the member variable `transl` and use it
+ instead. For example, it can used if the translation
+ `transl` is predicted from some external model.
+ (default=None)
+ return_verts: bool, optional
+ Return the vertices. (default=True)
+ return_full_pose: bool, optional
+ Returns the full axis-angle pose vector (default=False)
+
+ Returns
+ -------
+ output: ModelOutput
+ A named tuple of type `ModelOutput`
+ """
+
+ # If no shape and pose parameters are passed along, then use the
+ # ones from the module
+ global_orient = global_orient if global_orient is not None else self.global_orient
+ body_pose = body_pose if body_pose is not None else self.body_pose
+ betas = betas if betas is not None else self.betas
+
+ left_hand_pose = left_hand_pose if left_hand_pose is not None else self.left_hand_pose
+ right_hand_pose = right_hand_pose if right_hand_pose is not None else self.right_hand_pose
+ jaw_pose = jaw_pose if jaw_pose is not None else self.jaw_pose
+ leye_pose = leye_pose if leye_pose is not None else self.leye_pose
+ reye_pose = reye_pose if reye_pose is not None else self.reye_pose
+ expression = expression if expression is not None else self.expression
+
+ apply_trans = transl is not None or hasattr(self, "transl")
+ if transl is None:
+ if hasattr(self, "transl"):
+ transl = self.transl
+
+ if self.use_pca:
+ left_hand_pose = torch.einsum("bi,ij->bj", [left_hand_pose, self.left_hand_components])
+ right_hand_pose = torch.einsum(
+ "bi,ij->bj", [right_hand_pose, self.right_hand_components]
+ )
+
+ full_pose = torch.cat(
+ [
+ global_orient.reshape(-1, 1, 3),
+ body_pose.reshape(-1, self.NUM_BODY_JOINTS, 3),
+ jaw_pose.reshape(-1, 1, 3),
+ leye_pose.reshape(-1, 1, 3),
+ reye_pose.reshape(-1, 1, 3),
+ left_hand_pose.reshape(-1, 15, 3),
+ right_hand_pose.reshape(-1, 15, 3),
+ ],
+ dim=1,
+ ).reshape(-1, 165)
+
+ # Add the mean pose of the model. Does not affect the body, only the
+ # hands when flat_hand_mean == False
+ full_pose += self.pose_mean
+
+ batch_size = max(betas.shape[0], global_orient.shape[0], body_pose.shape[0])
+ # Concatenate the shape and expression coefficients
+ scale = int(batch_size / betas.shape[0])
+ if scale > 1:
+ betas = betas.expand(scale, -1)
+ shape_components = torch.cat([betas, expression], dim=-1)
+
+ shapedirs = torch.cat([self.shapedirs, self.expr_dirs], dim=-1)
+
+ vertices, joints = lbs(
+ shape_components,
+ full_pose,
+ self.v_template,
+ shapedirs,
+ self.posedirs,
+ self.J_regressor,
+ self.parents,
+ self.lbs_weights,
+ pose2rot=pose2rot,
+ )
+
+ lmk_faces_idx = self.lmk_faces_idx.unsqueeze(dim=0).expand(batch_size, -1).contiguous()
+ lmk_bary_coords = self.lmk_bary_coords.unsqueeze(dim=0).repeat(self.batch_size, 1, 1)
+ if self.use_face_contour:
+ lmk_idx_and_bcoords = find_dynamic_lmk_idx_and_bcoords(
+ vertices,
+ full_pose,
+ self.dynamic_lmk_faces_idx,
+ self.dynamic_lmk_bary_coords,
+ self.neck_kin_chain,
+ pose2rot=True,
+ )
+ dyn_lmk_faces_idx, dyn_lmk_bary_coords = lmk_idx_and_bcoords
+
+ lmk_faces_idx = torch.cat([lmk_faces_idx, dyn_lmk_faces_idx], 1)
+ lmk_bary_coords = torch.cat(
+ [lmk_bary_coords.expand(batch_size, -1, -1), dyn_lmk_bary_coords], 1
+ )
+
+ landmarks = vertices2landmarks(vertices, self.faces_tensor, lmk_faces_idx, lmk_bary_coords)
+
+ # Add any extra joints that might be needed
+ joints = self.vertex_joint_selector(vertices, joints)
+ # Add the landmarks to the joints
+ joints = torch.cat([joints, landmarks], dim=1)
+ # Map the joints to the current dataset
+
+ if self.joint_mapper is not None:
+ joints = self.joint_mapper(joints=joints, vertices=vertices)
+
+ if apply_trans:
+ joints += transl.unsqueeze(dim=1)
+ vertices += transl.unsqueeze(dim=1)
+
+ v_shaped = None
+ if return_shaped:
+ v_shaped = self.v_template + blend_shapes(betas, self.shapedirs)
+ else:
+ v_shaped = Tensor(0)
+ output = SMPLXOutput(
+ vertices=vertices if return_verts else None,
+ joints=joints,
+ betas=betas,
+ expression=expression,
+ global_orient=global_orient,
+ body_pose=body_pose,
+ left_hand_pose=left_hand_pose,
+ right_hand_pose=right_hand_pose,
+ jaw_pose=jaw_pose,
+ v_shaped=v_shaped,
+ full_pose=full_pose if return_full_pose else None,
+ )
+ return output
+
+
+class SMPLXLayer(SMPLX):
+ def __init__(self, *args, **kwargs) -> None:
+ # Just create a SMPLX module without any member variables
+ super(SMPLXLayer, self).__init__(
+ create_global_orient=False,
+ create_body_pose=False,
+ create_left_hand_pose=False,
+ create_right_hand_pose=False,
+ create_jaw_pose=False,
+ create_leye_pose=False,
+ create_reye_pose=False,
+ create_betas=False,
+ create_expression=False,
+ create_transl=False,
+ *args,
+ **kwargs,
+ )
+
+ def forward(
+ self,
+ betas: Optional[Tensor] = None,
+ global_orient: Optional[Tensor] = None,
+ body_pose: Optional[Tensor] = None,
+ left_hand_pose: Optional[Tensor] = None,
+ right_hand_pose: Optional[Tensor] = None,
+ transl: Optional[Tensor] = None,
+ expression: Optional[Tensor] = None,
+ jaw_pose: Optional[Tensor] = None,
+ leye_pose: Optional[Tensor] = None,
+ reye_pose: Optional[Tensor] = None,
+ return_verts: bool = True,
+ return_full_pose: bool = True,
+ **kwargs,
+ ) -> TensorOutput:
+ """
+ Forward pass for the SMPLX model
+
+ Parameters
+ ----------
+ global_orient: torch.tensor, optional, shape Bx3x3
+ If given, ignore the member variable and use it as the global
+ rotation of the body. Useful if someone wishes to predicts this
+ with an external model. It is expected to be in rotation matrix
+ format. (default=None)
+ betas: torch.tensor, optional, shape BxN_b
+ If given, ignore the member variable `betas` and use it
+ instead. For example, it can used if shape parameters
+ `betas` are predicted from some external model.
+ (default=None)
+ expression: torch.tensor, optional, shape BxN_e
+ Expression coefficients.
+ For example, it can used if expression parameters
+ `expression` are predicted from some external model.
+ body_pose: torch.tensor, optional, shape BxJx3x3
+ If given, ignore the member variable `body_pose` and use it
+ instead. For example, it can used if someone predicts the
+ pose of the body joints are predicted from some external model.
+ It should be a tensor that contains joint rotations in
+ rotation matrix format. (default=None)
+ left_hand_pose: torch.tensor, optional, shape Bx15x3x3
+ If given, contains the pose of the left hand.
+ It should be a tensor that contains joint rotations in
+ rotation matrix format. (default=None)
+ right_hand_pose: torch.tensor, optional, shape Bx15x3x3
+ If given, contains the pose of the right hand.
+ It should be a tensor that contains joint rotations in
+ rotation matrix format. (default=None)
+ jaw_pose: torch.tensor, optional, shape Bx3x3
+ Jaw pose. It should either joint rotations in
+ rotation matrix format.
+ transl: torch.tensor, optional, shape Bx3
+ Translation vector of the body.
+ For example, it can used if the translation
+ `transl` is predicted from some external model.
+ (default=None)
+ return_verts: bool, optional
+ Return the vertices. (default=True)
+ return_full_pose: bool, optional
+ Returns the full pose vector (default=False)
+ Returns
+ -------
+ output: ModelOutput
+ A data class that contains the posed vertices and joints
+ """
+ device, dtype = self.shapedirs.device, self.shapedirs.dtype
+
+ model_vars = [
+ betas,
+ global_orient,
+ body_pose,
+ transl,
+ expression,
+ left_hand_pose,
+ right_hand_pose,
+ jaw_pose,
+ ]
+ batch_size = 1
+ for var in model_vars:
+ if var is None:
+ continue
+ batch_size = max(batch_size, len(var))
+
+ if global_orient is None:
+ global_orient = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, -1, -1, -1)
+ .contiguous()
+ )
+ if body_pose is None:
+ body_pose = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, self.NUM_BODY_JOINTS, -1, -1)
+ .contiguous()
+ )
+ if left_hand_pose is None:
+ left_hand_pose = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, 15, -1, -1)
+ .contiguous()
+ )
+ if right_hand_pose is None:
+ right_hand_pose = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, 15, -1, -1)
+ .contiguous()
+ )
+ if jaw_pose is None:
+ jaw_pose = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, -1, -1, -1)
+ .contiguous()
+ )
+ if leye_pose is None:
+ leye_pose = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, -1, -1, -1)
+ .contiguous()
+ )
+ if reye_pose is None:
+ reye_pose = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, -1, -1, -1)
+ .contiguous()
+ )
+ if expression is None:
+ expression = torch.zeros(
+ [batch_size, self.num_expression_coeffs], dtype=dtype, device=device
+ )
+ if betas is None:
+ betas = torch.zeros([batch_size, self.num_betas], dtype=dtype, device=device)
+ if transl is None:
+ transl = torch.zeros([batch_size, 3], dtype=dtype, device=device)
+
+ # Concatenate all pose vectors
+ full_pose = torch.cat(
+ [
+ global_orient.reshape(-1, 1, 3, 3),
+ body_pose.reshape(-1, self.NUM_BODY_JOINTS, 3, 3),
+ jaw_pose.reshape(-1, 1, 3, 3),
+ leye_pose.reshape(-1, 1, 3, 3),
+ reye_pose.reshape(-1, 1, 3, 3),
+ left_hand_pose.reshape(-1, self.NUM_HAND_JOINTS, 3, 3),
+ right_hand_pose.reshape(-1, self.NUM_HAND_JOINTS, 3, 3),
+ ],
+ dim=1,
+ )
+ shape_components = torch.cat([betas, expression], dim=-1)
+
+ shapedirs = torch.cat([self.shapedirs, self.expr_dirs], dim=-1)
+
+ vertices, joints = lbs(
+ shape_components,
+ full_pose,
+ self.v_template,
+ shapedirs,
+ self.posedirs,
+ self.J_regressor,
+ self.parents,
+ self.lbs_weights,
+ pose2rot=False,
+ )
+
+ lmk_faces_idx = self.lmk_faces_idx.unsqueeze(dim=0).expand(batch_size, -1).contiguous()
+ lmk_bary_coords = self.lmk_bary_coords.unsqueeze(dim=0).repeat(batch_size, 1, 1)
+ if self.use_face_contour:
+ lmk_idx_and_bcoords = find_dynamic_lmk_idx_and_bcoords(
+ vertices,
+ full_pose,
+ self.dynamic_lmk_faces_idx,
+ self.dynamic_lmk_bary_coords,
+ self.neck_kin_chain,
+ pose2rot=False,
+ )
+ dyn_lmk_faces_idx, dyn_lmk_bary_coords = lmk_idx_and_bcoords
+
+ lmk_faces_idx = torch.cat([lmk_faces_idx, dyn_lmk_faces_idx], 1)
+ lmk_bary_coords = torch.cat(
+ [lmk_bary_coords.expand(batch_size, -1, -1), dyn_lmk_bary_coords], 1
+ )
+
+ landmarks = vertices2landmarks(vertices, self.faces_tensor, lmk_faces_idx, lmk_bary_coords)
+
+ # Add any extra joints that might be needed
+ joints = self.vertex_joint_selector(vertices, joints)
+ # Add the landmarks to the joints
+ joints = torch.cat([joints, landmarks], dim=1)
+ # Map the joints to the current dataset
+
+ if self.joint_mapper is not None:
+ joints = self.joint_mapper(joints=joints, vertices=vertices)
+
+ if transl is not None:
+ joints += transl.unsqueeze(dim=1)
+ vertices += transl.unsqueeze(dim=1)
+
+ output = TensorOutput(
+ vertices=vertices if return_verts else Tensor(0),
+ joints=joints,
+ betas=betas,
+ expression=expression,
+ global_orient=global_orient,
+ body_pose=body_pose,
+ left_hand_pose=left_hand_pose,
+ right_hand_pose=right_hand_pose,
+ jaw_pose=jaw_pose,
+ transl=transl if transl != None else Tensor(0),
+ full_pose=full_pose if return_full_pose else Tensor(0),
+ )
+
+ return output
+
+
+class MANO(SMPL):
+ # The hand joints are replaced by MANO
+ NUM_BODY_JOINTS = 1
+ NUM_HAND_JOINTS = 15
+ NUM_JOINTS = NUM_BODY_JOINTS + NUM_HAND_JOINTS
+
+ def __init__(
+ self,
+ model_path: str,
+ is_rhand: bool = True,
+ data_struct: Optional[Struct] = None,
+ create_hand_pose: bool = True,
+ hand_pose: Optional[Tensor] = None,
+ use_pca: bool = True,
+ num_pca_comps: int = 6,
+ flat_hand_mean: bool = False,
+ batch_size: int = 1,
+ dtype=torch.float32,
+ vertex_ids=None,
+ use_compressed: bool = True,
+ ext: str = "pkl",
+ **kwargs,
+ ) -> None:
+ """MANO model constructor
+
+ Parameters
+ ----------
+ model_path: str
+ The path to the folder or to the file where the model
+ parameters are stored
+ data_struct: Strct
+ A struct object. If given, then the parameters of the model are
+ read from the object. Otherwise, the model tries to read the
+ parameters from the given `model_path`. (default = None)
+ create_hand_pose: bool, optional
+ Flag for creating a member variable for the pose of the right
+ hand. (default = True)
+ hand_pose: torch.tensor, optional, BxP
+ The default value for the right hand pose member variable.
+ (default = None)
+ num_pca_comps: int, optional
+ The number of PCA components to use for each hand.
+ (default = 6)
+ flat_hand_mean: bool, optional
+ If False, then the pose of the hand is initialized to False.
+ batch_size: int, optional
+ The batch size used for creating the member variables
+ dtype: torch.dtype, optional
+ The data type for the created variables
+ vertex_ids: dict, optional
+ A dictionary containing the indices of the extra vertices that
+ will be selected
+ """
+
+ self.num_pca_comps = num_pca_comps
+ self.is_rhand = is_rhand
+ # If no data structure is passed, then load the data from the given
+ # model folder
+ if data_struct is None:
+ # Load the model
+ if osp.isdir(model_path):
+ model_fn = "MANO_{}.{ext}".format("RIGHT" if is_rhand else "LEFT", ext=ext)
+ mano_path = os.path.join(model_path, model_fn)
+ else:
+ mano_path = model_path
+ self.is_rhand = True if "RIGHT" in os.path.basename(model_path) else False
+ assert osp.exists(mano_path), "Path {} does not exist!".format(mano_path)
+
+ if ext == "pkl":
+ with open(mano_path, "rb") as mano_file:
+ model_data = pickle.load(mano_file, encoding="latin1")
+ elif ext == "npz":
+ model_data = np.load(mano_path, allow_pickle=True)
+ else:
+ raise ValueError("Unknown extension: {}".format(ext))
+ data_struct = Struct(**model_data)
+
+ if vertex_ids is None:
+ vertex_ids = VERTEX_IDS["smplh"]
+
+ super(MANO, self).__init__(
+ model_path=model_path,
+ data_struct=data_struct,
+ batch_size=batch_size,
+ vertex_ids=vertex_ids,
+ use_compressed=use_compressed,
+ dtype=dtype,
+ ext=ext,
+ **kwargs,
+ )
+
+ # add only MANO tips to the extra joints
+ self.vertex_joint_selector.extra_joints_idxs = to_tensor(
+ list(VERTEX_IDS["mano"].values()), dtype=torch.long
+ )
+
+ self.use_pca = use_pca
+ self.num_pca_comps = num_pca_comps
+ if self.num_pca_comps == 45:
+ self.use_pca = False
+ self.flat_hand_mean = flat_hand_mean
+
+ hand_components = data_struct.hands_components[:num_pca_comps]
+
+ self.np_hand_components = hand_components
+
+ if self.use_pca:
+ self.register_buffer("hand_components", torch.tensor(hand_components, dtype=dtype))
+
+ if self.flat_hand_mean:
+ hand_mean = np.zeros_like(data_struct.hands_mean)
+ else:
+ hand_mean = data_struct.hands_mean
+
+ self.register_buffer("hand_mean", to_tensor(hand_mean, dtype=self.dtype))
+
+ # Create the buffers for the pose of the left hand
+ hand_pose_dim = num_pca_comps if use_pca else 3 * self.NUM_HAND_JOINTS
+ if create_hand_pose:
+ if hand_pose is None:
+ default_hand_pose = torch.zeros([batch_size, hand_pose_dim], dtype=dtype)
+ else:
+ default_hand_pose = torch.tensor(hand_pose, dtype=dtype)
+
+ hand_pose_param = nn.Parameter(default_hand_pose, requires_grad=True)
+ self.register_parameter("hand_pose", hand_pose_param)
+
+ # Create the buffer for the mean pose.
+ pose_mean = self.create_mean_pose(data_struct, flat_hand_mean=flat_hand_mean)
+ pose_mean_tensor = pose_mean.clone().to(dtype)
+ # pose_mean_tensor = torch.tensor(pose_mean, dtype=dtype)
+ self.register_buffer("pose_mean", pose_mean_tensor)
+
+ def name(self) -> str:
+ return "MANO"
+
+ def create_mean_pose(self, data_struct, flat_hand_mean=False):
+ # Create the array for the mean pose. If flat_hand is false, then use
+ # the mean that is given by the data, rather than the flat open hand
+ global_orient_mean = torch.zeros([3], dtype=self.dtype)
+ pose_mean = torch.cat([global_orient_mean, self.hand_mean], dim=0)
+ return pose_mean
+
+ def extra_repr(self):
+ msg = [super(MANO, self).extra_repr()]
+ if self.use_pca:
+ msg.append(f"Number of PCA components: {self.num_pca_comps}")
+ msg.append(f"Flat hand mean: {self.flat_hand_mean}")
+ return "\n".join(msg)
+
+ def forward(
+ self,
+ betas: Optional[Tensor] = None,
+ global_orient: Optional[Tensor] = None,
+ hand_pose: Optional[Tensor] = None,
+ transl: Optional[Tensor] = None,
+ return_verts: bool = True,
+ return_full_pose: bool = False,
+ **kwargs,
+ ) -> MANOOutput:
+ """Forward pass for the MANO model"""
+ # If no shape and pose parameters are passed along, then use the
+ # ones from the module
+ global_orient = global_orient if global_orient is not None else self.global_orient
+ betas = betas if betas is not None else self.betas
+ hand_pose = hand_pose if hand_pose is not None else self.hand_pose
+
+ apply_trans = transl is not None or hasattr(self, "transl")
+ if transl is None:
+ if hasattr(self, "transl"):
+ transl = self.transl
+
+ if self.use_pca:
+ hand_pose = torch.einsum("bi,ij->bj", [hand_pose, self.hand_components])
+
+ full_pose = torch.cat([global_orient, hand_pose], dim=1)
+ full_pose += self.pose_mean
+
+ vertices, joints = lbs(
+ betas,
+ full_pose,
+ self.v_template,
+ self.shapedirs,
+ self.posedirs,
+ self.J_regressor,
+ self.parents,
+ self.lbs_weights,
+ pose2rot=True,
+ )
+
+ # # Add pre-selected extra joints that might be needed
+ # joints = self.vertex_joint_selector(vertices, joints)
+
+ if self.joint_mapper is not None:
+ joints = self.joint_mapper(joints)
+
+ if apply_trans:
+ joints = joints + transl.unsqueeze(dim=1)
+ vertices = vertices + transl.unsqueeze(dim=1)
+
+ output = MANOOutput(
+ vertices=vertices if return_verts else None,
+ joints=joints if return_verts else None,
+ betas=betas,
+ global_orient=global_orient,
+ hand_pose=hand_pose,
+ full_pose=full_pose if return_full_pose else None,
+ )
+
+ return output
+
+
+class MANOLayer(MANO):
+ def __init__(self, *args, **kwargs) -> None:
+ """MANO as a layer model constructor"""
+ super(MANOLayer, self).__init__(
+ create_global_orient=False,
+ create_hand_pose=False,
+ create_betas=False,
+ create_transl=False,
+ *args,
+ **kwargs,
+ )
+
+ def name(self) -> str:
+ return "MANO"
+
+ def forward(
+ self,
+ betas: Optional[Tensor] = None,
+ global_orient: Optional[Tensor] = None,
+ hand_pose: Optional[Tensor] = None,
+ transl: Optional[Tensor] = None,
+ return_verts: bool = True,
+ return_full_pose: bool = False,
+ **kwargs,
+ ) -> MANOOutput:
+ """Forward pass for the MANO model"""
+ device, dtype = self.shapedirs.device, self.shapedirs.dtype
+ if global_orient is None:
+ batch_size = 1
+ global_orient = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, -1, -1, -1)
+ .contiguous()
+ )
+ else:
+ batch_size = global_orient.shape[0]
+ if hand_pose is None:
+ hand_pose = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, 15, -1, -1)
+ .contiguous()
+ )
+ if betas is None:
+ betas = torch.zeros([batch_size, self.num_betas], dtype=dtype, device=device)
+ if transl is None:
+ transl = torch.zeros([batch_size, 3], dtype=dtype, device=device)
+
+ full_pose = torch.cat([global_orient, hand_pose], dim=1)
+ vertices, joints = lbs(
+ betas,
+ full_pose,
+ self.v_template,
+ self.shapedirs,
+ self.posedirs,
+ self.J_regressor,
+ self.parents,
+ self.lbs_weights,
+ pose2rot=False,
+ )
+
+ if self.joint_mapper is not None:
+ joints = self.joint_mapper(joints)
+
+ if transl is not None:
+ joints = joints + transl.unsqueeze(dim=1)
+ vertices = vertices + transl.unsqueeze(dim=1)
+
+ output = MANOOutput(
+ vertices=vertices if return_verts else None,
+ joints=joints if return_verts else None,
+ betas=betas,
+ global_orient=global_orient,
+ hand_pose=hand_pose,
+ full_pose=full_pose if return_full_pose else None,
+ )
+
+ return output
+
+
+class FLAME(SMPL):
+ NUM_JOINTS = 5
+ SHAPE_SPACE_DIM = 300
+ EXPRESSION_SPACE_DIM = 100
+ NECK_IDX = 0
+
+ def __init__(
+ self,
+ model_path: str,
+ data_struct=None,
+ num_expression_coeffs=10,
+ create_expression: bool = True,
+ expression: Optional[Tensor] = None,
+ create_neck_pose: bool = True,
+ neck_pose: Optional[Tensor] = None,
+ create_jaw_pose: bool = True,
+ jaw_pose: Optional[Tensor] = None,
+ create_leye_pose: bool = True,
+ leye_pose: Optional[Tensor] = None,
+ create_reye_pose=True,
+ reye_pose: Optional[Tensor] = None,
+ use_face_contour=False,
+ batch_size: int = 1,
+ gender: str = "neutral",
+ dtype: torch.dtype = torch.float32,
+ ext="pkl",
+ **kwargs,
+ ) -> None:
+ """FLAME model constructor
+
+ Parameters
+ ----------
+ model_path: str
+ The path to the folder or to the file where the model
+ parameters are stored
+ num_expression_coeffs: int, optional
+ Number of expression components to use
+ (default = 10).
+ create_expression: bool, optional
+ Flag for creating a member variable for the expression space
+ (default = True).
+ expression: torch.tensor, optional, Bx10
+ The default value for the expression member variable.
+ (default = None)
+ create_neck_pose: bool, optional
+ Flag for creating a member variable for the neck pose.
+ (default = False)
+ neck_pose: torch.tensor, optional, Bx3
+ The default value for the neck pose variable.
+ (default = None)
+ create_jaw_pose: bool, optional
+ Flag for creating a member variable for the jaw pose.
+ (default = False)
+ jaw_pose: torch.tensor, optional, Bx3
+ The default value for the jaw pose variable.
+ (default = None)
+ create_leye_pose: bool, optional
+ Flag for creating a member variable for the left eye pose.
+ (default = False)
+ leye_pose: torch.tensor, optional, Bx10
+ The default value for the left eye pose variable.
+ (default = None)
+ create_reye_pose: bool, optional
+ Flag for creating a member variable for the right eye pose.
+ (default = False)
+ reye_pose: torch.tensor, optional, Bx10
+ The default value for the right eye pose variable.
+ (default = None)
+ use_face_contour: bool, optional
+ Whether to compute the keypoints that form the facial contour
+ batch_size: int, optional
+ The batch size used for creating the member variables
+ gender: str, optional
+ Which gender to load
+ dtype: torch.dtype
+ The data type for the created variables
+ """
+ model_fn = f"FLAME_{gender.upper()}.{ext}"
+ flame_path = os.path.join(model_path, model_fn)
+ assert osp.exists(flame_path), "Path {} does not exist!".format(flame_path)
+ if ext == "npz":
+ file_data = np.load(flame_path, allow_pickle=True)
+ elif ext == "pkl":
+ with open(flame_path, "rb") as smpl_file:
+ file_data = pickle.load(smpl_file, encoding="latin1")
+ else:
+ raise ValueError("Unknown extension: {}".format(ext))
+ data_struct = Struct(**file_data)
+
+ super(FLAME, self).__init__(
+ model_path=model_path,
+ data_struct=data_struct,
+ dtype=dtype,
+ batch_size=batch_size,
+ gender=gender,
+ ext=ext,
+ **kwargs,
+ )
+
+ self.use_face_contour = use_face_contour
+
+ self.vertex_joint_selector.extra_joints_idxs = to_tensor([], dtype=torch.long)
+
+ if create_neck_pose:
+ if neck_pose is None:
+ default_neck_pose = torch.zeros([batch_size, 3], dtype=dtype)
+ else:
+ default_neck_pose = torch.tensor(neck_pose, dtype=dtype)
+ neck_pose_param = nn.Parameter(default_neck_pose, requires_grad=True)
+ self.register_parameter("neck_pose", neck_pose_param)
+
+ if create_jaw_pose:
+ if jaw_pose is None:
+ default_jaw_pose = torch.zeros([batch_size, 3], dtype=dtype)
+ else:
+ default_jaw_pose = torch.tensor(jaw_pose, dtype=dtype)
+ jaw_pose_param = nn.Parameter(default_jaw_pose, requires_grad=True)
+ self.register_parameter("jaw_pose", jaw_pose_param)
+
+ if create_leye_pose:
+ if leye_pose is None:
+ default_leye_pose = torch.zeros([batch_size, 3], dtype=dtype)
+ else:
+ default_leye_pose = torch.tensor(leye_pose, dtype=dtype)
+ leye_pose_param = nn.Parameter(default_leye_pose, requires_grad=True)
+ self.register_parameter("leye_pose", leye_pose_param)
+
+ if create_reye_pose:
+ if reye_pose is None:
+ default_reye_pose = torch.zeros([batch_size, 3], dtype=dtype)
+ else:
+ default_reye_pose = torch.tensor(reye_pose, dtype=dtype)
+ reye_pose_param = nn.Parameter(default_reye_pose, requires_grad=True)
+ self.register_parameter("reye_pose", reye_pose_param)
+
+ shapedirs = data_struct.shapedirs
+ if len(shapedirs.shape) < 3:
+ shapedirs = shapedirs[:, :, None]
+ if shapedirs.shape[-1] < self.SHAPE_SPACE_DIM + self.EXPRESSION_SPACE_DIM:
+ print(
+ f"WARNING: You are using a {self.name()} model, with only"
+ " 10 shape and 10 expression coefficients."
+ )
+ expr_start_idx = 10
+ expr_end_idx = 20
+ num_expression_coeffs = min(num_expression_coeffs, 10)
+ else:
+ expr_start_idx = self.SHAPE_SPACE_DIM
+ expr_end_idx = self.SHAPE_SPACE_DIM + num_expression_coeffs
+ num_expression_coeffs = min(num_expression_coeffs, self.EXPRESSION_SPACE_DIM)
+
+ self._num_expression_coeffs = num_expression_coeffs
+
+ expr_dirs = shapedirs[:, :, expr_start_idx:expr_end_idx]
+ self.register_buffer("expr_dirs", to_tensor(to_np(expr_dirs), dtype=dtype))
+
+ if create_expression:
+ if expression is None:
+ default_expression = torch.zeros(
+ [batch_size, self.num_expression_coeffs], dtype=dtype
+ )
+ else:
+ default_expression = torch.tensor(expression, dtype=dtype)
+ expression_param = nn.Parameter(default_expression, requires_grad=True)
+ self.register_parameter("expression", expression_param)
+
+ # The pickle file that contains the barycentric coordinates for
+ # regressing the landmarks
+ landmark_bcoord_filename = osp.join(model_path, "flame_static_embedding.pkl")
+
+ with open(landmark_bcoord_filename, "rb") as fp:
+ landmarks_data = pickle.load(fp, encoding="latin1")
+
+ lmk_faces_idx = landmarks_data["lmk_face_idx"].astype(np.int64)
+ self.register_buffer("lmk_faces_idx", torch.tensor(lmk_faces_idx, dtype=torch.long))
+ lmk_bary_coords = landmarks_data["lmk_b_coords"]
+ self.register_buffer("lmk_bary_coords", torch.tensor(lmk_bary_coords, dtype=dtype))
+ if self.use_face_contour:
+ face_contour_path = os.path.join(model_path, "flame_dynamic_embedding.npy")
+ contour_embeddings = np.load(face_contour_path, allow_pickle=True, encoding="latin1")[
+ ()
+ ]
+
+ dynamic_lmk_faces_idx = np.array(contour_embeddings["lmk_face_idx"], dtype=np.int64)
+ dynamic_lmk_faces_idx = torch.tensor(dynamic_lmk_faces_idx, dtype=torch.long)
+ self.register_buffer("dynamic_lmk_faces_idx", dynamic_lmk_faces_idx)
+
+ dynamic_lmk_b_coords = torch.tensor(contour_embeddings["lmk_b_coords"], dtype=dtype)
+ self.register_buffer("dynamic_lmk_bary_coords", dynamic_lmk_b_coords)
+
+ neck_kin_chain = find_joint_kin_chain(self.NECK_IDX, self.parents)
+ self.register_buffer("neck_kin_chain", torch.tensor(neck_kin_chain, dtype=torch.long))
+
+ @property
+ def num_expression_coeffs(self):
+ return self._num_expression_coeffs
+
+ def name(self) -> str:
+ return "FLAME"
+
+ def extra_repr(self):
+ msg = [
+ super(FLAME, self).extra_repr(),
+ f"Number of Expression Coefficients: {self.num_expression_coeffs}",
+ f"Use face contour: {self.use_face_contour}",
+ ]
+ return "\n".join(msg)
+
+ def forward(
+ self,
+ betas: Optional[Tensor] = None,
+ global_orient: Optional[Tensor] = None,
+ neck_pose: Optional[Tensor] = None,
+ transl: Optional[Tensor] = None,
+ expression: Optional[Tensor] = None,
+ jaw_pose: Optional[Tensor] = None,
+ leye_pose: Optional[Tensor] = None,
+ reye_pose: Optional[Tensor] = None,
+ return_verts: bool = True,
+ return_full_pose: bool = False,
+ pose2rot: bool = True,
+ **kwargs,
+ ) -> FLAMEOutput:
+ """
+ Forward pass for the SMPLX model
+
+ Parameters
+ ----------
+ global_orient: torch.tensor, optional, shape Bx3
+ If given, ignore the member variable and use it as the global
+ rotation of the body. Useful if someone wishes to predicts this
+ with an external model. (default=None)
+ betas: torch.tensor, optional, shape Bx10
+ If given, ignore the member variable `betas` and use it
+ instead. For example, it can used if shape parameters
+ `betas` are predicted from some external model.
+ (default=None)
+ expression: torch.tensor, optional, shape Bx10
+ If given, ignore the member variable `expression` and use it
+ instead. For example, it can used if expression parameters
+ `expression` are predicted from some external model.
+ jaw_pose: torch.tensor, optional, shape Bx3
+ If given, ignore the member variable `jaw_pose` and
+ use this instead. It should either joint rotations in
+ axis-angle format.
+ jaw_pose: torch.tensor, optional, shape Bx3
+ If given, ignore the member variable `jaw_pose` and
+ use this instead. It should either joint rotations in
+ axis-angle format.
+ transl: torch.tensor, optional, shape Bx3
+ If given, ignore the member variable `transl` and use it
+ instead. For example, it can used if the translation
+ `transl` is predicted from some external model.
+ (default=None)
+ return_verts: bool, optional
+ Return the vertices. (default=True)
+ return_full_pose: bool, optional
+ Returns the full axis-angle pose vector (default=False)
+
+ Returns
+ -------
+ output: ModelOutput
+ A named tuple of type `ModelOutput`
+ """
+
+ # If no shape and pose parameters are passed along, then use the
+ # ones from the module
+ global_orient = global_orient if global_orient is not None else self.global_orient
+ jaw_pose = jaw_pose if jaw_pose is not None else self.jaw_pose
+ neck_pose = neck_pose if neck_pose is not None else self.neck_pose
+
+ leye_pose = leye_pose if leye_pose is not None else self.leye_pose
+ reye_pose = reye_pose if reye_pose is not None else self.reye_pose
+
+ betas = betas if betas is not None else self.betas
+ expression = expression if expression is not None else self.expression
+
+ apply_trans = transl is not None or hasattr(self, "transl")
+ if transl is None:
+ if hasattr(self, "transl"):
+ transl = self.transl
+
+ full_pose = torch.cat([global_orient, neck_pose, jaw_pose, leye_pose, reye_pose], dim=1)
+
+ batch_size = max(betas.shape[0], global_orient.shape[0], jaw_pose.shape[0])
+ # Concatenate the shape and expression coefficients
+ scale = int(batch_size / betas.shape[0])
+ if scale > 1:
+ betas = betas.expand(scale, -1)
+ shape_components = torch.cat([betas, expression], dim=-1)
+ shapedirs = torch.cat([self.shapedirs, self.expr_dirs], dim=-1)
+
+ vertices, joints = lbs(
+ shape_components,
+ full_pose,
+ self.v_template,
+ shapedirs,
+ self.posedirs,
+ self.J_regressor,
+ self.parents,
+ self.lbs_weights,
+ pose2rot=pose2rot,
+ )
+
+ lmk_faces_idx = self.lmk_faces_idx.unsqueeze(dim=0).expand(batch_size, -1).contiguous()
+ lmk_bary_coords = self.lmk_bary_coords.unsqueeze(dim=0).repeat(batch_size, 1, 1)
+ if self.use_face_contour:
+ lmk_idx_and_bcoords = find_dynamic_lmk_idx_and_bcoords(
+ vertices,
+ full_pose,
+ self.dynamic_lmk_faces_idx,
+ self.dynamic_lmk_bary_coords,
+ self.neck_kin_chain,
+ pose2rot=True,
+ )
+ dyn_lmk_faces_idx, dyn_lmk_bary_coords = lmk_idx_and_bcoords
+ lmk_faces_idx = torch.cat([lmk_faces_idx, dyn_lmk_faces_idx], 1)
+ lmk_bary_coords = torch.cat(
+ [lmk_bary_coords.expand(batch_size, -1, -1), dyn_lmk_bary_coords], 1
+ )
+
+ landmarks = vertices2landmarks(vertices, self.faces_tensor, lmk_faces_idx, lmk_bary_coords)
+
+ # Add any extra joints that might be needed
+ joints = self.vertex_joint_selector(vertices, joints)
+ # Add the landmarks to the joints
+ joints = torch.cat([joints, landmarks], dim=1)
+
+ # Map the joints to the current dataset
+ if self.joint_mapper is not None:
+ joints = self.joint_mapper(joints=joints, vertices=vertices)
+
+ if apply_trans:
+ joints += transl.unsqueeze(dim=1)
+ vertices += transl.unsqueeze(dim=1)
+
+ output = FLAMEOutput(
+ vertices=vertices if return_verts else None,
+ joints=joints,
+ betas=betas,
+ expression=expression,
+ global_orient=global_orient,
+ neck_pose=neck_pose,
+ jaw_pose=jaw_pose,
+ full_pose=full_pose if return_full_pose else None,
+ )
+ return output
+
+
+class FLAMELayer(FLAME):
+ def __init__(self, *args, **kwargs) -> None:
+ """FLAME as a layer model constructor"""
+ super(FLAMELayer, self).__init__(
+ create_betas=False,
+ create_expression=False,
+ create_global_orient=False,
+ create_neck_pose=False,
+ create_jaw_pose=False,
+ create_leye_pose=False,
+ create_reye_pose=False,
+ *args,
+ **kwargs,
+ )
+
+ def forward(
+ self,
+ betas: Optional[Tensor] = None,
+ global_orient: Optional[Tensor] = None,
+ neck_pose: Optional[Tensor] = None,
+ transl: Optional[Tensor] = None,
+ expression: Optional[Tensor] = None,
+ jaw_pose: Optional[Tensor] = None,
+ leye_pose: Optional[Tensor] = None,
+ reye_pose: Optional[Tensor] = None,
+ return_verts: bool = True,
+ return_full_pose: bool = False,
+ pose2rot: bool = True,
+ **kwargs,
+ ) -> FLAMEOutput:
+ """
+ Forward pass for the SMPLX model
+
+ Parameters
+ ----------
+ global_orient: torch.tensor, optional, shape Bx3x3
+ Global rotation of the body. Useful if someone wishes to
+ predicts this with an external model. It is expected to be in
+ rotation matrix format. (default=None)
+ betas: torch.tensor, optional, shape BxN_b
+ Shape parameters. For example, it can used if shape parameters
+ `betas` are predicted from some external model.
+ (default=None)
+ expression: torch.tensor, optional, shape BxN_e
+ If given, ignore the member variable `expression` and use it
+ instead. For example, it can used if expression parameters
+ `expression` are predicted from some external model.
+ jaw_pose: torch.tensor, optional, shape Bx3x3
+ Jaw pose. It should either joint rotations in
+ rotation matrix format.
+ transl: torch.tensor, optional, shape Bx3
+ Translation vector of the body.
+ For example, it can used if the translation
+ `transl` is predicted from some external model.
+ (default=None)
+ return_verts: bool, optional
+ Return the vertices. (default=True)
+ return_full_pose: bool, optional
+ Returns the full axis-angle pose vector (default=False)
+
+ Returns
+ -------
+ output: ModelOutput
+ A named tuple of type `ModelOutput`
+ """
+ device, dtype = self.shapedirs.device, self.shapedirs.dtype
+ if global_orient is None:
+ batch_size = 1
+ global_orient = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, -1, -1, -1)
+ .contiguous()
+ )
+ else:
+ batch_size = global_orient.shape[0]
+ if neck_pose is None:
+ neck_pose = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, 1, -1, -1)
+ .contiguous()
+ )
+ if jaw_pose is None:
+ jaw_pose = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, -1, -1, -1)
+ .contiguous()
+ )
+ if leye_pose is None:
+ leye_pose = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, -1, -1, -1)
+ .contiguous()
+ )
+ if reye_pose is None:
+ reye_pose = (
+ torch.eye(3, device=device, dtype=dtype)
+ .view(1, 1, 3, 3)
+ .expand(batch_size, -1, -1, -1)
+ .contiguous()
+ )
+ if betas is None:
+ betas = torch.zeros([batch_size, self.num_betas], dtype=dtype, device=device)
+ if expression is None:
+ expression = torch.zeros(
+ [batch_size, self.num_expression_coeffs], dtype=dtype, device=device
+ )
+ if transl is None:
+ transl = torch.zeros([batch_size, 3], dtype=dtype, device=device)
+
+ full_pose = torch.cat([global_orient, neck_pose, jaw_pose, leye_pose, reye_pose], dim=1)
+
+ shape_components = torch.cat([betas, expression], dim=-1)
+ shapedirs = torch.cat([self.shapedirs, self.expr_dirs], dim=-1)
+
+ vertices, joints = lbs(
+ shape_components,
+ full_pose,
+ self.v_template,
+ shapedirs,
+ self.posedirs,
+ self.J_regressor,
+ self.parents,
+ self.lbs_weights,
+ pose2rot=False,
+ )
+
+ lmk_faces_idx = self.lmk_faces_idx.unsqueeze(dim=0).expand(batch_size, -1).contiguous()
+ lmk_bary_coords = self.lmk_bary_coords.unsqueeze(dim=0).repeat(batch_size, 1, 1)
+ if self.use_face_contour:
+ lmk_idx_and_bcoords = find_dynamic_lmk_idx_and_bcoords(
+ vertices,
+ full_pose,
+ self.dynamic_lmk_faces_idx,
+ self.dynamic_lmk_bary_coords,
+ self.neck_kin_chain,
+ pose2rot=False,
+ )
+ dyn_lmk_faces_idx, dyn_lmk_bary_coords = lmk_idx_and_bcoords
+ lmk_faces_idx = torch.cat([lmk_faces_idx, dyn_lmk_faces_idx], 1)
+ lmk_bary_coords = torch.cat(
+ [lmk_bary_coords.expand(batch_size, -1, -1), dyn_lmk_bary_coords], 1
+ )
+
+ landmarks = vertices2landmarks(vertices, self.faces_tensor, lmk_faces_idx, lmk_bary_coords)
+
+ # Add any extra joints that might be needed
+ joints = self.vertex_joint_selector(vertices, joints)
+ # Add the landmarks to the joints
+ joints = torch.cat([joints, landmarks], dim=1)
+
+ # Map the joints to the current dataset
+ if self.joint_mapper is not None:
+ joints = self.joint_mapper(joints=joints, vertices=vertices)
+
+ joints += transl.unsqueeze(dim=1)
+ vertices += transl.unsqueeze(dim=1)
+
+ output = FLAMEOutput(
+ vertices=vertices if return_verts else None,
+ joints=joints,
+ betas=betas,
+ expression=expression,
+ global_orient=global_orient,
+ neck_pose=neck_pose,
+ jaw_pose=jaw_pose,
+ full_pose=full_pose if return_full_pose else None,
+ )
+ return output
+
+
+def build_layer(
+ model_path: str, model_type: str = "smpl", **kwargs
+) -> Union[SMPLLayer, SMPLHLayer, SMPLXLayer, MANOLayer, FLAMELayer]:
+ """Method for creating a model from a path and a model type
+
+ Parameters
+ ----------
+ model_path: str
+ Either the path to the model you wish to load or a folder,
+ where each subfolder contains the differents types, i.e.:
+ model_path:
+ |
+ |-- smpl
+ |-- SMPL_FEMALE
+ |-- SMPL_NEUTRAL
+ |-- SMPL_MALE
+ |-- smplh
+ |-- SMPLH_FEMALE
+ |-- SMPLH_MALE
+ |-- smplx
+ |-- SMPLX_FEMALE
+ |-- SMPLX_NEUTRAL
+ |-- SMPLX_MALE
+ |-- mano
+ |-- MANO RIGHT
+ |-- MANO LEFT
+ |-- flame
+ |-- FLAME_FEMALE
+ |-- FLAME_MALE
+ |-- FLAME_NEUTRAL
+
+ model_type: str, optional
+ When model_path is a folder, then this parameter specifies the
+ type of model to be loaded
+ **kwargs: dict
+ Keyword arguments
+
+ Returns
+ -------
+ body_model: nn.Module
+ The PyTorch module that implements the corresponding body model
+ Raises
+ ------
+ ValueError: In case the model type is not one of SMPL, SMPLH,
+ SMPLX, MANO or FLAME
+ """
+
+ if osp.isdir(model_path):
+ model_path = os.path.join(model_path, model_type)
+ else:
+ model_type = osp.basename(model_path).split("_")[0].lower()
+
+ if model_type.lower() == "smpl":
+ return SMPLLayer(model_path, **kwargs)
+ elif model_type.lower() == "smplh":
+ return SMPLHLayer(model_path, **kwargs)
+ elif model_type.lower() == "smplx":
+ return SMPLXLayer(model_path, **kwargs)
+ elif "mano" in model_type.lower():
+ return MANOLayer(model_path, **kwargs)
+ elif "flame" in model_type.lower():
+ return FLAMELayer(model_path, **kwargs)
+ else:
+ raise ValueError(f"Unknown model type {model_type}, exiting!")
+
+
+def create(
+ model_path: str, model_type: str = "smpl", **kwargs
+) -> Union[SMPL, SMPLH, SMPLX, MANO, FLAME]:
+ """Method for creating a model from a path and a model type
+
+ Parameters
+ ----------
+ model_path: str
+ Either the path to the model you wish to load or a folder,
+ where each subfolder contains the differents types, i.e.:
+ model_path:
+ |
+ |-- smpl
+ |-- SMPL_FEMALE
+ |-- SMPL_NEUTRAL
+ |-- SMPL_MALE
+ |-- smplh
+ |-- SMPLH_FEMALE
+ |-- SMPLH_MALE
+ |-- smplx
+ |-- SMPLX_FEMALE
+ |-- SMPLX_NEUTRAL
+ |-- SMPLX_MALE
+ |-- mano
+ |-- MANO RIGHT
+ |-- MANO LEFT
+
+ model_type: str, optional
+ When model_path is a folder, then this parameter specifies the
+ type of model to be loaded
+ **kwargs: dict
+ Keyword arguments
+
+ Returns
+ -------
+ body_model: nn.Module
+ The PyTorch module that implements the corresponding body model
+ Raises
+ ------
+ ValueError: In case the model type is not one of SMPL, SMPLH,
+ SMPLX, MANO or FLAME
+ """
+
+ # If it's a folder, assume
+ if osp.isdir(model_path):
+ model_path = os.path.join(model_path, model_type)
+ else:
+ model_type = osp.basename(model_path).split("_")[0].lower()
+
+ if model_type.lower() == "smpl":
+ return SMPL(model_path, **kwargs)
+ elif model_type.lower() == "smplh":
+ return SMPLH(model_path, **kwargs)
+ elif model_type.lower() == "smplx":
+ return SMPLX(model_path, **kwargs)
+ elif "mano" in model_type.lower():
+ return MANO(model_path, **kwargs)
+ elif "flame" in model_type.lower():
+ return FLAME(model_path, **kwargs)
+ else:
+ raise ValueError(f"Unknown model type {model_type}, exiting!")
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/smplx/joint_names.py b/worldbench/third_party/drivestudio/third_party/smplx/smplx/joint_names.py
new file mode 100644
index 0000000..fcdb2e1
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/smplx/joint_names.py
@@ -0,0 +1,320 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+import numpy as np
+
+JOINT_NAMES = [
+ "pelvis",
+ "left_hip",
+ "right_hip",
+ "spine1",
+ "left_knee",
+ "right_knee",
+ "spine2",
+ "left_ankle",
+ "right_ankle",
+ "spine3",
+ "left_foot",
+ "right_foot",
+ "neck",
+ "left_collar",
+ "right_collar",
+ "head",
+ "left_shoulder",
+ "right_shoulder",
+ "left_elbow",
+ "right_elbow",
+ "left_wrist",
+ "right_wrist",
+ "jaw",
+ "left_eye_smplhf",
+ "right_eye_smplhf",
+ "left_index1",
+ "left_index2",
+ "left_index3",
+ "left_middle1",
+ "left_middle2",
+ "left_middle3",
+ "left_pinky1",
+ "left_pinky2",
+ "left_pinky3",
+ "left_ring1",
+ "left_ring2",
+ "left_ring3",
+ "left_thumb1",
+ "left_thumb2",
+ "left_thumb3",
+ "right_index1",
+ "right_index2",
+ "right_index3",
+ "right_middle1",
+ "right_middle2",
+ "right_middle3",
+ "right_pinky1",
+ "right_pinky2",
+ "right_pinky3",
+ "right_ring1",
+ "right_ring2",
+ "right_ring3",
+ "right_thumb1",
+ "right_thumb2",
+ "right_thumb3",
+ "nose",
+ "right_eye",
+ "left_eye",
+ "right_ear",
+ "left_ear",
+ "left_big_toe",
+ "left_small_toe",
+ "left_heel",
+ "right_big_toe",
+ "right_small_toe",
+ "right_heel",
+ "left_thumb",
+ "left_index",
+ "left_middle",
+ "left_ring",
+ "left_pinky",
+ "right_thumb",
+ "right_index",
+ "right_middle",
+ "right_ring",
+ "right_pinky",
+ "right_eye_brow1",
+ "right_eye_brow2",
+ "right_eye_brow3",
+ "right_eye_brow4",
+ "right_eye_brow5",
+ "left_eye_brow5",
+ "left_eye_brow4",
+ "left_eye_brow3",
+ "left_eye_brow2",
+ "left_eye_brow1",
+ "nose1",
+ "nose2",
+ "nose3",
+ "nose4",
+ "right_nose_2",
+ "right_nose_1",
+ "nose_middle",
+ "left_nose_1",
+ "left_nose_2",
+ "right_eye1",
+ "right_eye2",
+ "right_eye3",
+ "right_eye4",
+ "right_eye5",
+ "right_eye6",
+ "left_eye4",
+ "left_eye3",
+ "left_eye2",
+ "left_eye1",
+ "left_eye6",
+ "left_eye5",
+ "right_mouth_1",
+ "right_mouth_2",
+ "right_mouth_3",
+ "mouth_top",
+ "left_mouth_3",
+ "left_mouth_2",
+ "left_mouth_1",
+ "left_mouth_5", # 59 in OpenPose output
+ "left_mouth_4", # 58 in OpenPose output
+ "mouth_bottom",
+ "right_mouth_4",
+ "right_mouth_5",
+ "right_lip_1",
+ "right_lip_2",
+ "lip_top",
+ "left_lip_2",
+ "left_lip_1",
+ "left_lip_3",
+ "lip_bottom",
+ "right_lip_3",
+ # Face contour
+ "right_contour_1",
+ "right_contour_2",
+ "right_contour_3",
+ "right_contour_4",
+ "right_contour_5",
+ "right_contour_6",
+ "right_contour_7",
+ "right_contour_8",
+ "contour_middle",
+ "left_contour_8",
+ "left_contour_7",
+ "left_contour_6",
+ "left_contour_5",
+ "left_contour_4",
+ "left_contour_3",
+ "left_contour_2",
+ "left_contour_1",
+]
+
+
+SMPLH_JOINT_NAMES = [
+ "pelvis",
+ "left_hip",
+ "right_hip",
+ "spine1",
+ "left_knee",
+ "right_knee",
+ "spine2",
+ "left_ankle",
+ "right_ankle",
+ "spine3",
+ "left_foot",
+ "right_foot",
+ "neck",
+ "left_collar",
+ "right_collar",
+ "head",
+ "left_shoulder",
+ "right_shoulder",
+ "left_elbow",
+ "right_elbow",
+ "left_wrist",
+ "right_wrist",
+ "left_index1",
+ "left_index2",
+ "left_index3",
+ "left_middle1",
+ "left_middle2",
+ "left_middle3",
+ "left_pinky1",
+ "left_pinky2",
+ "left_pinky3",
+ "left_ring1",
+ "left_ring2",
+ "left_ring3",
+ "left_thumb1",
+ "left_thumb2",
+ "left_thumb3",
+ "right_index1",
+ "right_index2",
+ "right_index3",
+ "right_middle1",
+ "right_middle2",
+ "right_middle3",
+ "right_pinky1",
+ "right_pinky2",
+ "right_pinky3",
+ "right_ring1",
+ "right_ring2",
+ "right_ring3",
+ "right_thumb1",
+ "right_thumb2",
+ "right_thumb3",
+ "nose",
+ "right_eye",
+ "left_eye",
+ "right_ear",
+ "left_ear",
+ "left_big_toe",
+ "left_small_toe",
+ "left_heel",
+ "right_big_toe",
+ "right_small_toe",
+ "right_heel",
+ "left_thumb",
+ "left_index",
+ "left_middle",
+ "left_ring",
+ "left_pinky",
+ "right_thumb",
+ "right_index",
+ "right_middle",
+ "right_ring",
+ "right_pinky",
+]
+
+SMPL_JOINT_NAMES = [
+ "pelvis",
+ "left_hip",
+ "right_hip",
+ "spine1",
+ "left_knee",
+ "right_knee",
+ "spine2",
+ "left_ankle",
+ "right_ankle",
+ "spine3",
+ "left_foot",
+ "right_foot",
+ "neck",
+ "left_collar",
+ "right_collar",
+ "head",
+ "left_shoulder",
+ "right_shoulder",
+ "left_elbow",
+ "right_elbow",
+ "left_wrist",
+ "right_wrist",
+ "left_hand",
+ "right_hand",
+]
+
+
+class Body:
+ """
+ Class for storing a single body pose.
+ """
+
+ def __init__(self, joints, joint_names):
+ assert joints.ndim > 1
+ assert joints.shape[0] == len(joint_names)
+ self.joints = {}
+ for i, j in enumerate(joint_names):
+ self.joints[j] = joints[i]
+
+ @staticmethod
+ def from_smpl(joints):
+ """
+ Create a Body object from SMPL joints.
+ """
+ return Body(joints, SMPL_JOINT_NAMES)
+
+ @staticmethod
+ def from_smplh(joints):
+ """
+ Create a Body object from SMPLH joints.
+ """
+ return Body(joints, SMPLH_JOINT_NAMES)
+
+ def _as(self, joint_names):
+ """
+ Return a Body object with the specified joint names.
+ """
+ joint_list = []
+ for j in joint_names:
+ if j not in self.joints:
+ joint_list.append(np.zeros_like(self.joints["spine1"]))
+ else:
+ joint_list.append(self.joints[j])
+ return np.stack(joint_list, axis=0)
+
+ def as_smpl(self):
+ """
+ Convert the body to SMPL joints.
+ """
+ return self._as(SMPL_JOINT_NAMES)
+
+ def as_smplh(self):
+ """
+ Convert the body to SMPLH joints.
+ """
+ return self._as(SMPLH_JOINT_NAMES)
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/smplx/lbs.py b/worldbench/third_party/drivestudio/third_party/smplx/smplx/lbs.py
new file mode 100644
index 0000000..025fa87
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/smplx/lbs.py
@@ -0,0 +1,418 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+from typing import Tuple, List
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+
+from .utils import rot_mat_to_euler, Tensor
+
+
+def find_dynamic_lmk_idx_and_bcoords(
+ vertices: Tensor,
+ pose: Tensor,
+ dynamic_lmk_faces_idx: Tensor,
+ dynamic_lmk_b_coords: Tensor,
+ neck_kin_chain: List[int],
+ pose2rot: bool = True,
+) -> Tuple[Tensor, Tensor]:
+ ''' Compute the faces, barycentric coordinates for the dynamic landmarks
+
+
+ To do so, we first compute the rotation of the neck around the y-axis
+ and then use a pre-computed look-up table to find the faces and the
+ barycentric coordinates that will be used.
+
+ Special thanks to Soubhik Sanyal (soubhik.sanyal@tuebingen.mpg.de)
+ for providing the original TensorFlow implementation and for the LUT.
+
+ Parameters
+ ----------
+ vertices: torch.tensor BxVx3, dtype = torch.float32
+ The tensor of input vertices
+ pose: torch.tensor Bx(Jx3), dtype = torch.float32
+ The current pose of the body model
+ dynamic_lmk_faces_idx: torch.tensor L, dtype = torch.long
+ The look-up table from neck rotation to faces
+ dynamic_lmk_b_coords: torch.tensor Lx3, dtype = torch.float32
+ The look-up table from neck rotation to barycentric coordinates
+ neck_kin_chain: list
+ A python list that contains the indices of the joints that form the
+ kinematic chain of the neck.
+ dtype: torch.dtype, optional
+
+ Returns
+ -------
+ dyn_lmk_faces_idx: torch.tensor, dtype = torch.long
+ A tensor of size BxL that contains the indices of the faces that
+ will be used to compute the current dynamic landmarks.
+ dyn_lmk_b_coords: torch.tensor, dtype = torch.float32
+ A tensor of size BxL that contains the indices of the faces that
+ will be used to compute the current dynamic landmarks.
+ '''
+
+ dtype = vertices.dtype
+ batch_size = vertices.shape[0]
+
+ if pose2rot:
+ aa_pose = torch.index_select(pose.view(batch_size, -1, 3), 1,
+ neck_kin_chain)
+ rot_mats = batch_rodrigues(
+ aa_pose.view(-1, 3)).view(batch_size, -1, 3, 3)
+ else:
+ rot_mats = torch.index_select(
+ pose.view(batch_size, -1, 3, 3), 1, neck_kin_chain)
+
+ rel_rot_mat = torch.eye(
+ 3, device=vertices.device, dtype=dtype).unsqueeze_(dim=0).repeat(
+ batch_size, 1, 1)
+ for idx in range(len(neck_kin_chain)):
+ rel_rot_mat = torch.bmm(rot_mats[:, idx], rel_rot_mat)
+
+ y_rot_angle = torch.round(
+ torch.clamp(-rot_mat_to_euler(rel_rot_mat) * 180.0 / np.pi,
+ max=39)).to(dtype=torch.long)
+ neg_mask = y_rot_angle.lt(0).to(dtype=torch.long)
+ mask = y_rot_angle.lt(-39).to(dtype=torch.long)
+ neg_vals = mask * 78 + (1 - mask) * (39 - y_rot_angle)
+ y_rot_angle = (neg_mask * neg_vals +
+ (1 - neg_mask) * y_rot_angle)
+
+ dyn_lmk_faces_idx = torch.index_select(dynamic_lmk_faces_idx,
+ 0, y_rot_angle)
+ dyn_lmk_b_coords = torch.index_select(dynamic_lmk_b_coords,
+ 0, y_rot_angle)
+
+ return dyn_lmk_faces_idx, dyn_lmk_b_coords
+
+
+def vertices2landmarks(
+ vertices: Tensor,
+ faces: Tensor,
+ lmk_faces_idx: Tensor,
+ lmk_bary_coords: Tensor
+) -> Tensor:
+ ''' Calculates landmarks by barycentric interpolation
+
+ Parameters
+ ----------
+ vertices: torch.tensor BxVx3, dtype = torch.float32
+ The tensor of input vertices
+ faces: torch.tensor Fx3, dtype = torch.long
+ The faces of the mesh
+ lmk_faces_idx: torch.tensor L, dtype = torch.long
+ The tensor with the indices of the faces used to calculate the
+ landmarks.
+ lmk_bary_coords: torch.tensor Lx3, dtype = torch.float32
+ The tensor of barycentric coordinates that are used to interpolate
+ the landmarks
+
+ Returns
+ -------
+ landmarks: torch.tensor BxLx3, dtype = torch.float32
+ The coordinates of the landmarks for each mesh in the batch
+ '''
+ # Extract the indices of the vertices for each face
+ # BxLx3
+ batch_size, num_verts = vertices.shape[:2]
+ device = vertices.device
+
+ lmk_faces = torch.index_select(faces, 0, lmk_faces_idx.view(-1).to(torch.long)).view(
+ batch_size, -1, 3)
+ #The '.to(torch.long)'.
+ # added to make the trace work in c++,
+ # otherwise you get a runtime error in c++:
+ # 'index_select(): Expected dtype int32 or int64 for index'
+
+ lmk_faces += torch.arange(
+ batch_size, dtype=torch.long, device=device).view(-1, 1, 1) * num_verts
+
+ lmk_vertices = vertices.view(-1, 3)[lmk_faces].view(
+ batch_size, -1, 3, 3)
+
+ landmarks = torch.einsum('blfi,blf->bli', [lmk_vertices, lmk_bary_coords])
+ return landmarks
+
+
+def lbs(
+ betas: Tensor,
+ pose: Tensor,
+ v_template: Tensor,
+ shapedirs: Tensor,
+ posedirs: Tensor,
+ J_regressor: Tensor,
+ parents: Tensor,
+ lbs_weights: Tensor,
+ pose2rot: bool = True,
+ return_T: bool = False, # ! JH 2023.9 modified here
+ return_posed_v: bool = False, # ! JH 2023.9 modified here
+ return_A = False, # ! JH 2023.10 modified here
+ return_J = False, # ! JH 2023.10 modified here
+) -> Tuple[Tensor, Tensor]:
+ ''' Performs Linear Blend Skinning with the given shape and pose parameters
+
+ Parameters
+ ----------
+ betas : torch.tensor BxNB
+ The tensor of shape parameters
+ pose : torch.tensor Bx(J + 1) * 3
+ The pose parameters in axis-angle format
+ v_template torch.tensor BxVx3
+ The template mesh that will be deformed
+ shapedirs : torch.tensor 1xNB
+ The tensor of PCA shape displacements
+ posedirs : torch.tensor Px(V * 3)
+ The pose PCA coefficients
+ J_regressor : torch.tensor JxV
+ The regressor array that is used to calculate the joints from
+ the position of the vertices
+ parents: torch.tensor J
+ The array that describes the kinematic tree for the model
+ lbs_weights: torch.tensor N x V x (J + 1)
+ The linear blend skinning weights that represent how much the
+ rotation matrix of each part affects each vertex
+ pose2rot: bool, optional
+ Flag on whether to convert the input pose tensor to rotation
+ matrices. The default value is True. If False, then the pose tensor
+ should already contain rotation matrices and have a size of
+ Bx(J + 1)x9
+ dtype: torch.dtype, optional
+
+ Returns
+ -------
+ verts: torch.tensor BxVx3
+ The vertices of the mesh after applying the shape and pose
+ displacements.
+ joints: torch.tensor BxJx3
+ The joints of the model
+ '''
+
+ batch_size = max(betas.shape[0], pose.shape[0])
+ device, dtype = betas.device, betas.dtype
+
+ # Add shape contribution
+ v_shaped = v_template + blend_shapes(betas, shapedirs)
+
+ # Get the joints
+ # NxJx3 array
+ J = vertices2joints(J_regressor, v_shaped)
+
+ # 3. Add pose blend shapes
+ # N x J x 3 x 3
+ ident = torch.eye(3, dtype=dtype, device=device)
+ if pose2rot:
+ rot_mats = batch_rodrigues(pose.view(-1, 3)).view(
+ [batch_size, -1, 3, 3])
+
+ pose_feature = (rot_mats[:, 1:, :, :] - ident).view([batch_size, -1])
+ # (N x P) x (P, V * 3) -> N x V x 3
+ pose_offsets = torch.matmul(
+ pose_feature, posedirs).view(batch_size, -1, 3)
+ else:
+ pose_feature = pose[:, 1:].view(batch_size, -1, 3, 3) - ident
+ rot_mats = pose.view(batch_size, -1, 3, 3)
+
+ pose_offsets = torch.matmul(pose_feature.view(batch_size, -1),
+ posedirs).view(batch_size, -1, 3)
+
+ v_posed = pose_offsets + v_shaped
+ # 4. Get the global joint location
+ J_transformed, A = batch_rigid_transform(rot_mats, J, parents, dtype=dtype)
+
+ # 5. Do skinning:
+ # W is N x V x (J + 1)
+ W = lbs_weights.unsqueeze(dim=0).expand([batch_size, -1, -1])
+ # (N x V x (J + 1)) x (N x (J + 1) x 16)
+ num_joints = J_regressor.shape[0]
+ T = torch.matmul(W, A.view(batch_size, num_joints, 16)) \
+ .view(batch_size, -1, 4, 4)
+
+ homogen_coord = torch.ones([batch_size, v_posed.shape[1], 1],
+ dtype=dtype, device=device)
+ v_posed_homo = torch.cat([v_posed, homogen_coord], dim=2)
+ v_homo = torch.matmul(T, torch.unsqueeze(v_posed_homo, dim=-1))
+
+ verts = v_homo[:, :, :3, 0]
+
+ ret_list = [verts, J_transformed]
+ if return_T:
+ ret_list.append(T)
+ if return_posed_v:
+ ret_list.append(v_posed)
+ if return_A:
+ ret_list.append(A)
+ if return_J:
+ ret_list.append(J)
+ return ret_list
+
+
+def vertices2joints(J_regressor: Tensor, vertices: Tensor) -> Tensor:
+ ''' Calculates the 3D joint locations from the vertices
+
+ Parameters
+ ----------
+ J_regressor : torch.tensor JxV
+ The regressor array that is used to calculate the joints from the
+ position of the vertices
+ vertices : torch.tensor BxVx3
+ The tensor of mesh vertices
+
+ Returns
+ -------
+ torch.tensor BxJx3
+ The location of the joints
+ '''
+
+ return torch.einsum('bik,ji->bjk', [vertices, J_regressor])
+
+
+def blend_shapes(betas: Tensor, shape_disps: Tensor) -> Tensor:
+ ''' Calculates the per vertex displacement due to the blend shapes
+
+
+ Parameters
+ ----------
+ betas : torch.tensor Bx(num_betas)
+ Blend shape coefficients
+ shape_disps: torch.tensor Vx3x(num_betas)
+ Blend shapes
+
+ Returns
+ -------
+ torch.tensor BxVx3
+ The per-vertex displacement due to shape deformation
+ '''
+
+ # Displacement[b, m, k] = sum_{l} betas[b, l] * shape_disps[m, k, l]
+ # i.e. Multiply each shape displacement by its corresponding beta and
+ # then sum them.
+ blend_shape = torch.einsum('bl,mkl->bmk', [betas, shape_disps])
+ return blend_shape
+
+
+def batch_rodrigues(
+ rot_vecs: Tensor,
+ epsilon: float = 1e-8,
+) -> Tensor:
+ ''' Calculates the rotation matrices for a batch of rotation vectors
+ Parameters
+ ----------
+ rot_vecs: torch.tensor Nx3
+ array of N axis-angle vectors
+ Returns
+ -------
+ R: torch.tensor Nx3x3
+ The rotation matrices for the given axis-angle parameters
+ '''
+
+ batch_size = rot_vecs.shape[0]
+ device, dtype = rot_vecs.device, rot_vecs.dtype
+
+ angle = torch.norm(rot_vecs + 1e-8, dim=1, keepdim=True)
+ rot_dir = rot_vecs / angle
+
+ cos = torch.unsqueeze(torch.cos(angle), dim=1)
+ sin = torch.unsqueeze(torch.sin(angle), dim=1)
+
+ # Bx1 arrays
+ rx, ry, rz = torch.split(rot_dir, 1, dim=1)
+ K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device)
+
+ zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device)
+ K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
+ .view((batch_size, 3, 3))
+
+ ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
+ rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
+ return rot_mat
+
+
+def transform_mat(R: Tensor, t: Tensor) -> Tensor:
+ ''' Creates a batch of transformation matrices
+ Args:
+ - R: Bx3x3 array of a batch of rotation matrices
+ - t: Bx3x1 array of a batch of translation vectors
+ Returns:
+ - T: Bx4x4 Transformation matrix
+ '''
+ # No padding left or right, only add an extra row
+ return torch.cat([F.pad(R, [0, 0, 0, 1]),
+ F.pad(t, [0, 0, 0, 1], value=1)], dim=2)
+
+
+def batch_rigid_transform(
+ rot_mats: Tensor,
+ joints: Tensor,
+ parents: Tensor,
+ dtype=torch.float32
+) -> Tensor:
+ """
+ Applies a batch of rigid transformations to the joints
+
+ Parameters
+ ----------
+ rot_mats : torch.tensor BxNx3x3
+ Tensor of rotation matrices
+ joints : torch.tensor BxNx3
+ Locations of joints
+ parents : torch.tensor BxN
+ The kinematic tree of each object
+ dtype : torch.dtype, optional:
+ The data type of the created tensors, the default is torch.float32
+
+ Returns
+ -------
+ posed_joints : torch.tensor BxNx3
+ The locations of the joints after applying the pose rotations
+ rel_transforms : torch.tensor BxNx4x4
+ The relative (with respect to the root joint) rigid transformations
+ for all the joints
+ """
+
+ joints = torch.unsqueeze(joints, dim=-1)
+
+ rel_joints = joints.clone()
+ rel_joints[:, 1:] -= joints[:, parents[1:]]
+
+ transforms_mat = transform_mat(
+ rot_mats.reshape(-1, 3, 3),
+ rel_joints.reshape(-1, 3, 1)).reshape(-1, joints.shape[1], 4, 4)
+
+ transform_chain = [transforms_mat[:, 0]]
+ for i in range(1, parents.shape[0]):
+ # Subtract the joint location at the rest pose
+ # No need for rotation, since it's identity when at rest
+ curr_res = torch.matmul(transform_chain[parents[i]],
+ transforms_mat[:, i])
+ transform_chain.append(curr_res)
+
+ transforms = torch.stack(transform_chain, dim=1)
+
+ # The last column of the transformations contains the posed joints
+ posed_joints = transforms[:, :, :3, 3]
+
+ joints_homogen = F.pad(joints, [0, 0, 0, 1])
+
+ rel_transforms = transforms - F.pad(
+ torch.matmul(transforms, joints_homogen), [3, 0, 0, 0, 0, 0, 0, 0])
+
+ return posed_joints, rel_transforms
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/smplx/utils.py b/worldbench/third_party/drivestudio/third_party/smplx/smplx/utils.py
new file mode 100644
index 0000000..753c8e6
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/smplx/utils.py
@@ -0,0 +1,129 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from typing import NewType, Union, Optional
+from dataclasses import dataclass, asdict, fields
+import numpy as np
+import torch
+
+Tensor = NewType('Tensor', torch.Tensor)
+Array = NewType('Array', np.ndarray)
+
+
+@dataclass
+class ModelOutput:
+ vertices: Optional[Tensor] = None
+ joints: Optional[Tensor] = None
+ full_pose: Optional[Tensor] = None
+ global_orient: Optional[Tensor] = None
+ transl: Optional[Tensor] = None
+ v_posed: Optional[Tensor] = None
+
+ def __getitem__(self, key):
+ return getattr(self, key)
+
+ def get(self, key, default=None):
+ return getattr(self, key, default)
+
+ def __iter__(self):
+ return self.keys()
+
+ def keys(self):
+ keys = [t.name for t in fields(self)]
+ return iter(keys)
+
+ def values(self):
+ values = [getattr(self, t.name) for t in fields(self)]
+ return iter(values)
+
+ def items(self):
+ data = [(t.name, getattr(self, t.name)) for t in fields(self)]
+ return iter(data)
+
+
+@dataclass
+class SMPLOutput(ModelOutput):
+ betas: Optional[Tensor] = None
+ body_pose: Optional[Tensor] = None
+ T: Optional[Tensor] = None
+ A: Optional[Tensor] = None
+ J: Optional[Tensor] = None
+
+
+@dataclass
+class SMPLHOutput(SMPLOutput):
+ left_hand_pose: Optional[Tensor] = None
+ right_hand_pose: Optional[Tensor] = None
+ transl: Optional[Tensor] = None
+
+
+@dataclass
+class SMPLXOutput(SMPLHOutput):
+ expression: Optional[Tensor] = None
+ jaw_pose: Optional[Tensor] = None
+
+
+@dataclass
+class MANOOutput(ModelOutput):
+ betas: Optional[Tensor] = None
+ hand_pose: Optional[Tensor] = None
+
+
+@dataclass
+class FLAMEOutput(ModelOutput):
+ betas: Optional[Tensor] = None
+ expression: Optional[Tensor] = None
+ jaw_pose: Optional[Tensor] = None
+ neck_pose: Optional[Tensor] = None
+
+
+def find_joint_kin_chain(joint_id, kinematic_tree):
+ kin_chain = []
+ curr_idx = joint_id
+ while curr_idx != -1:
+ kin_chain.append(curr_idx)
+ curr_idx = kinematic_tree[curr_idx]
+ return kin_chain
+
+
+def to_tensor(
+ array: Union[Array, Tensor], dtype=torch.float32
+) -> Tensor:
+ if torch.is_tensor(array):
+ return array
+ else:
+ return torch.tensor(array, dtype=dtype)
+
+
+class Struct(object):
+ def __init__(self, **kwargs):
+ for key, val in kwargs.items():
+ setattr(self, key, val)
+
+
+def to_np(array, dtype=np.float32):
+ if 'scipy.sparse' in str(type(array)):
+ array = array.todense()
+ return np.array(array, dtype=dtype)
+
+
+def rot_mat_to_euler(rot_mats):
+ # Calculates rotation matrix to euler angles
+ # Careful for extreme cases of eular angles like [0.0, pi, 0.0]
+
+ sy = torch.sqrt(rot_mats[:, 0, 0] * rot_mats[:, 0, 0] +
+ rot_mats[:, 1, 0] * rot_mats[:, 1, 0])
+ return torch.atan2(-rot_mats[:, 2, 0], sy)
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/smplx/vertex_ids.py b/worldbench/third_party/drivestudio/third_party/smplx/smplx/vertex_ids.py
new file mode 100644
index 0000000..0e7a4c3
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/smplx/vertex_ids.py
@@ -0,0 +1,77 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from __future__ import print_function
+from __future__ import absolute_import
+from __future__ import division
+
+# Joint name to vertex mapping. SMPL/SMPL-H/SMPL-X vertices that correspond to
+# MSCOCO and OpenPose joints
+vertex_ids = {
+ 'smplh': {
+ 'nose': 332,
+ 'reye': 6260,
+ 'leye': 2800,
+ 'rear': 4071,
+ 'lear': 583,
+ 'rthumb': 6191,
+ 'rindex': 5782,
+ 'rmiddle': 5905,
+ 'rring': 6016,
+ 'rpinky': 6133,
+ 'lthumb': 2746,
+ 'lindex': 2319,
+ 'lmiddle': 2445,
+ 'lring': 2556,
+ 'lpinky': 2673,
+ 'LBigToe': 3216,
+ 'LSmallToe': 3226,
+ 'LHeel': 3387,
+ 'RBigToe': 6617,
+ 'RSmallToe': 6624,
+ 'RHeel': 6787
+ },
+ 'smplx': {
+ 'nose': 9120,
+ 'reye': 9929,
+ 'leye': 9448,
+ 'rear': 616,
+ 'lear': 6,
+ 'rthumb': 8079,
+ 'rindex': 7669,
+ 'rmiddle': 7794,
+ 'rring': 7905,
+ 'rpinky': 8022,
+ 'lthumb': 5361,
+ 'lindex': 4933,
+ 'lmiddle': 5058,
+ 'lring': 5169,
+ 'lpinky': 5286,
+ 'LBigToe': 5770,
+ 'LSmallToe': 5780,
+ 'LHeel': 8846,
+ 'RBigToe': 8463,
+ 'RSmallToe': 8474,
+ 'RHeel': 8635
+ },
+ 'mano': {
+ 'thumb': 744,
+ 'index': 320,
+ 'middle': 443,
+ 'ring': 554,
+ 'pinky': 671,
+ }
+}
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/smplx/vertex_joint_selector.py b/worldbench/third_party/drivestudio/third_party/smplx/smplx/vertex_joint_selector.py
new file mode 100644
index 0000000..1744972
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/smplx/vertex_joint_selector.py
@@ -0,0 +1,80 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+
+from .utils import to_tensor
+
+
+class VertexJointSelector(nn.Module):
+
+ def __init__(self, vertex_ids=None,
+ use_hands=True,
+ use_feet_keypoints=True, **kwargs):
+ super(VertexJointSelector, self).__init__()
+
+ extra_joints_idxs = []
+
+ face_keyp_idxs = np.array([
+ vertex_ids['nose'],
+ vertex_ids['reye'],
+ vertex_ids['leye'],
+ vertex_ids['rear'],
+ vertex_ids['lear']], dtype=np.int64)
+
+ extra_joints_idxs = np.concatenate([extra_joints_idxs,
+ face_keyp_idxs])
+
+ if use_feet_keypoints:
+ feet_keyp_idxs = np.array([vertex_ids['LBigToe'],
+ vertex_ids['LSmallToe'],
+ vertex_ids['LHeel'],
+ vertex_ids['RBigToe'],
+ vertex_ids['RSmallToe'],
+ vertex_ids['RHeel']], dtype=np.int32)
+
+ extra_joints_idxs = np.concatenate(
+ [extra_joints_idxs, feet_keyp_idxs])
+
+ if use_hands:
+ self.tip_names = ['thumb', 'index', 'middle', 'ring', 'pinky']
+
+ tips_idxs = []
+ for hand_id in ['l', 'r']:
+ for tip_name in self.tip_names:
+ tips_idxs.append(vertex_ids[hand_id + tip_name])
+
+ extra_joints_idxs = np.concatenate(
+ [extra_joints_idxs, tips_idxs])
+
+ self.register_buffer('extra_joints_idxs',
+ to_tensor(extra_joints_idxs, dtype=torch.long))
+
+ def forward(self, vertices, joints):
+ extra_joints = torch.index_select(vertices, 1, self.extra_joints_idxs.to(torch.long)) #The '.to(torch.long)'.
+ # added to make the trace work in c++,
+ # otherwise you get a runtime error in c++:
+ # 'index_select(): Expected dtype int32 or int64 for index'
+ joints = torch.cat([joints, extra_joints], dim=1)
+
+ return joints
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/README.md b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/README.md
new file mode 100644
index 0000000..9520a18
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/README.md
@@ -0,0 +1,253 @@
+# Model parameter transfer
+
+## Table of Contents
+ * [License](#license)
+ * [Description](#description)
+ * [Using the code](#using-the-code)
+ * [Data](#data)
+ * [Steps](#steps)
+ * [SMPL to SMPL-X](#smpl-to-smpl-x)
+ * [SMPL-X to SMPL](#smpl-x-to-smpl)
+ * [SMPL+H to SMPL](#smpl%2Bh-to-smpl)
+ * [SMPL to SMPL+H](#smpl-to-smpl%2Bh)
+ * [SMPL+H to SMPL-X](#smpl%2Bh-to-smpl-x)
+ * [SMPL-X to SMPL+H](#smpl-x-to-smpl%2Bh)
+ * [Visualize correspondences](visualize-correspondences)
+ * [Citation](#citation)
+ * [Acknowledgments](#acknowledgments)
+ * [Contact](#contact)
+
+## License
+
+Software Copyright License for **non-commercial scientific research purposes**.
+Please read carefully the [terms and conditions](https://github.com/vchoutas/smplx/blob/master/LICENSE) and any accompanying documentation before you download and/or use the SMPL-X/SMPLify-X model, data and software, (the "Model & Software"), including 3D meshes, blend weights, blend shapes, textures, software, scripts, and animations. By downloading and/or using the Model & Software (including downloading, cloning, installing, and any other use of this github repository), you acknowledge that you have read these terms and conditions, understand them, and agree to be bound by them. If you do not agree with these terms and conditions, you must not download and/or use the Model & Software. Any infringement of the terms of this agreement will automatically terminate your rights under this [License](./LICENSE).
+
+## Description
+
+The repository contains code for converting model parameters of one model to
+another. **Never** copy parameters between the models. You will not get the
+same poses. SMPL, SMPL+H and SMPL-X shape spaces are **NOT** compatible, since
+each model is the result of a different training process.
+A more detailed explanation on how we extract correspondences
+between the models and the loss function used to estimate the parameters can be
+found [here](./docs/transfer.md).
+
+## Requirements
+
+1. Install [mesh](https://github.com/MPI-IS/mesh)
+2. Start by cloning the SMPL-X repo:
+```Shell
+git clone https://github.com/vchoutas/smplx.git
+```
+3. Run the following command to install all necessary requirements
+```Shell
+ pip install -r requirements.txt
+```
+4. Install the Torch Trust Region optimizer by following the instructions [here](https://github.com/vchoutas/torch-trust-ncg)
+5. Install loguru
+6. Install open3d
+7. Install omegaconf
+
+## Using the code
+
+### Data
+
+Register on the [SMPL-X website](http://smpl-x.is.tue.mpg.de/), go to the
+downloads section to get the correspondences and sample data,
+by clicking on the *Model correspondences* button.
+Create a folder
+named `transfer_data` and extract the downloaded zip there. You should have the
+following folder structure now:
+
+```bash
+transfer_data
+├── meshes
+│ ├── smpl
+│ ├── smplx
+├── smpl2smplh_def_transfer.pkl
+├── smpl2smplx_deftrafo_setup.pkl
+├── smplh2smpl_def_transfer.pkl
+├── smplh2smplx_deftrafo_setup.pkl
+├── smplx2smpl_deftrafo_setup.pkl
+├── smplx2smplh_deftrafo_setup.pkl
+├── smplx_mask_ids.npy
+```
+
+### Steps
+
+First, break the motion into a set of pose `.obj` files. Depending on how the
+SMPL-* parameters are stored this code will differ. For the example AMASS data
+in this repository you can use the example code here:
+
+```
+python write_obj.py --model-folder ../models/ --motion-file ../transfer_data/support_data/github_data/amass_sample.npz --output-folder ../transfer_data/meshes/amass_sample/
+```
+
+To run the `transfer_model` utility you will require a `.yaml` config file,
+which can point to the location the output `.obj` files have been saved. Use the
+templates in `config_files` in the root of this repository. To convert the
+sample AMASS code to SMPL-X:
+
+```
+python -m transfer_model --exp-cfg config_files/smplh2smplx_as.yaml
+```
+
+Finally, the output `.obj` files have to be merged into a single motion
+sequence. Example code to do this in a way that matches `SMPL-X` AMASS archives
+can be found in `merge_output.py` and run as follows:
+
+```
+python merge_output.py --gender neutral ../output
+```
+
+Debug notes describing common problems encountered during this can be found
+[here](https://github.com/gngdb/smplx/blob/debug/transfer_model/DEBUG_NOTES.md).
+Problems are also discussed in
+[two](https://github.com/vchoutas/smplx/issues/82)
+[issues](https://github.com/vchoutas/smplx/issues/75).
+
+### SMPL to SMPL-X
+
+To run the code to convert SMPL meshes to SMPL-X parameters use the following command:
+ ```Shell
+ python -m transfer_model --exp-cfg config_files/smpl2smplx.yaml
+ ```
+This should be run from the top directory of the repository.
+
+The file *smpl2smplx.yaml* contains a sample configuration that reads meshes from a folder,
+processes them and returns pkl files with SMPL-X parameters. To run on your own data create a folder
+with SMPL meshes, in either ply or obj format, change the path in the config file and run the code.
+
+### SMPL-X to SMPL
+
+To run the code to convert SMPL-X meshes to SMPL parameters use the following command:
+ ```Shell
+ python main.py --exp-cfg config_files/smplx2smpl.yaml
+ ```
+
+The file *smplx2smpl.yaml* contains a sample configuration that reads meshes from a folder,
+processes them and returns pkl files with SMPL parameters. To run on your own data create a folder
+with SMPL-X meshes, in either ply or obj format, change the path in the config file and run the code.
+When creating the SMPL-X meshes, do not use the hand and face parameters.
+Naturally, you will lose all hand and face information if you choose this, since
+SMPL cannot model them.
+
+
+### SMPL+H to SMPL
+
+To run the code to convert SMPL+H meshes to SMPL parameters use the following command from the root `smplx` directory:
+ ```Shell
+ python -m transfer_model --exp-cfg config_files/smplh2smpl.yaml
+ ```
+This should be run from the top directory of the repository.
+
+The file *smplh2smpl.yaml* contains a sample configuration that reads meshes from a folder,
+processes them and returns pkl files with SMPL parameters. To run on your own data create a folder
+with SMPL+H meshes, in either ply or obj format, change the path in the config file and run the code.
+Note that using this direction means that you will lose information on the
+hands.
+
+
+### SMPL to SMPL+H
+
+To run the code to convert SMPL meshes to SMPL+H parameters use the following command:
+ ```Shell
+ python -m transfer_model --exp-cfg config_files/smpl2smplh.yaml
+ ```
+This should be run from the top directory of the repository.
+
+The file *smpl2smplh.yaml* contains a sample configuration that reads meshes from a folder,
+processes them and returns pkl files with SMPL parameters. To run on your own data create a folder
+with SMPL meshes, in either ply or obj format, change the path in the config file and run the code.
+
+### SMPL+H to SMPL-X
+
+To run the code to convert SMPL+H meshes to SMPL-X parameters use the following command:
+ ```Shell
+ python -m transfer_model --exp-cfg config_files/smplh2smplx.yaml
+ ```
+This should be run from the top directory of the repository.
+
+The file *smplh2smplx.yaml* contains a sample configuration that reads meshes from a folder,
+processes them and returns pkl files with SMPL-X parameters. To run on your own data create a folder
+with SMPL+H meshes, in either ply or obj format, change the path in the config file and run the code.
+
+
+### SMPL-X to SMPL+H
+
+To run the code to convert SMPL-X meshes to SMPL+H parameters use the following command:
+ ```Shell
+ python -m transfer_model --exp-cfg config_files/smplx2smplh.yaml
+ ```
+This should be run from the top directory of the repository.
+
+The file *smplx2smpl.yaml* contains a sample configuration that reads meshes from a folder,
+processes them and returns pkl files with SMPL+H parameters. To run on your own data create a folder
+with SMPL-X meshes, in either ply or obj format, change the path in the config file and run the code.
+Make sure that you do not use the jaw pose and expression parameters to generate
+the meshes.
+
+
+## Visualize correspondences
+
+To visualize correspondences:
+```Shell
+python vis_correspondences.py --exp-cfg configs/smpl2smplx.yaml --exp-opts colors_path PATH_TO_SMPL_COLORS
+```
+You should then see the following image. Points with similar color are in
+correspondence.
+
+
+## Citation
+
+Depending on which model is loaded for your project, i.e. SMPL-X or SMPL+H or SMPL, please cite the most relevant work:
+
+```
+@article{SMPL:2015,
+ author = {Loper, Matthew and Mahmood, Naureen and Romero, Javier and Pons-Moll, Gerard and Black, Michael J.},
+ title = {{SMPL}: A Skinned Multi-Person Linear Model},
+ journal = {ACM Transactions on Graphics, (Proc. SIGGRAPH Asia)},
+ month = oct,
+ number = {6},
+ pages = {248:1--248:16},
+ publisher = {ACM},
+ volume = {34},
+ year = {2015}
+}
+```
+
+```
+@article{MANO:SIGGRAPHASIA:2017,
+ title = {Embodied Hands: Modeling and Capturing Hands and Bodies Together},
+ author = {Romero, Javier and Tzionas, Dimitrios and Black, Michael J.},
+ journal = {ACM Transactions on Graphics, (Proc. SIGGRAPH Asia)},
+ volume = {36},
+ number = {6},
+ pages = {245:1--245:17},
+ series = {245:1--245:17},
+ publisher = {ACM},
+ month = nov,
+ year = {2017},
+ url = {http://doi.acm.org/10.1145/3130800.3130883},
+ month_numeric = {11}
+ }
+```
+
+
+```
+@inproceedings{SMPL-X:2019,
+ title = {Expressive Body Capture: 3D Hands, Face, and Body from a Single Image},
+ author = {Pavlakos, Georgios and Choutas, Vasileios and Ghorbani, Nima and Bolkart, Timo and Osman, Ahmed A. A. and Tzionas, Dimitrios and Black, Michael J.},
+ booktitle = {Proceedings IEEE Conf. on Computer Vision and Pattern Recognition (CVPR)},
+ year = {2019}
+}
+```
+
+
+## Acknowledgments
+The code of this repository was implemented by [Vassilis Choutas](vassilis.choutas@tuebingen.mpg.de),
+based on a Chumpy implementation from [Timo Bolkart](timo.bolkart@tuebingen.mpg.de).
+
+## Contact
+
+For questions, please contact [smplx@tue.mpg.de](smplx@tue.mpg.de).
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/__init__.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/__init__.py
new file mode 100644
index 0000000..6253a11
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/__init__.py
@@ -0,0 +1,15 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/__main__.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/__main__.py
new file mode 100644
index 0000000..2ffc250
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/__main__.py
@@ -0,0 +1,104 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+import os
+import os.path as osp
+import sys
+import pickle
+
+import numpy as np
+import open3d as o3d
+import torch
+from loguru import logger
+from tqdm import tqdm
+
+from smplx import build_layer
+
+from .config import parse_args
+from .data import build_dataloader
+from .transfer_model import run_fitting
+from .utils import read_deformation_transfer, np_mesh_to_o3d
+
+
+def main() -> None:
+ exp_cfg = parse_args()
+
+ if torch.cuda.is_available() and exp_cfg["use_cuda"]:
+ device = torch.device('cuda')
+ else:
+ device = torch.device('cpu')
+ if exp_cfg["use_cuda"]:
+ if input("use_cuda=True and GPU is not available, using CPU instead,"
+ " would you like to continue? (y/n)") != "y":
+ sys.exit(3)
+
+ logger.remove()
+ logger.add(
+ lambda x: tqdm.write(x, end=''), level=exp_cfg.logger_level.upper(),
+ colorize=True)
+
+ output_folder = osp.expanduser(osp.expandvars(exp_cfg.output_folder))
+ logger.info(f'Saving output to: {output_folder}')
+ os.makedirs(output_folder, exist_ok=True)
+
+ model_path = exp_cfg.body_model.folder
+ body_model = build_layer(model_path, **exp_cfg.body_model)
+ logger.info(body_model)
+ body_model = body_model.to(device=device)
+
+ deformation_transfer_path = exp_cfg.get('deformation_transfer_path', '')
+ def_matrix = read_deformation_transfer(
+ deformation_transfer_path, device=device)
+
+ # Read mask for valid vertex ids
+ mask_ids_fname = osp.expandvars(exp_cfg.mask_ids_fname)
+ mask_ids = None
+ if osp.exists(mask_ids_fname):
+ logger.info(f'Loading mask ids from: {mask_ids_fname}')
+ mask_ids = np.load(mask_ids_fname)
+ mask_ids = torch.from_numpy(mask_ids).to(device=device)
+ else:
+ logger.warning(f'Mask ids fname not found: {mask_ids_fname}')
+
+ data_obj_dict = build_dataloader(exp_cfg)
+
+ dataloader = data_obj_dict['dataloader']
+
+ for ii, batch in enumerate(tqdm(dataloader)):
+ for key in batch:
+ if torch.is_tensor(batch[key]):
+ batch[key] = batch[key].to(device=device)
+ var_dict = run_fitting(
+ exp_cfg, batch, body_model, def_matrix, mask_ids)
+ paths = batch['paths']
+
+ for ii, path in enumerate(paths):
+ _, fname = osp.split(path)
+
+ output_path = osp.join(
+ output_folder, f'{osp.splitext(fname)[0]}.pkl')
+ with open(output_path, 'wb') as f:
+ pickle.dump(var_dict, f)
+
+ output_path = osp.join(
+ output_folder, f'{osp.splitext(fname)[0]}.obj')
+ mesh = np_mesh_to_o3d(
+ var_dict['vertices'][ii], var_dict['faces'])
+ o3d.io.write_triangle_mesh(output_path, mesh)
+
+
+if __name__ == '__main__':
+ main()
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/__init__.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/__init__.py
new file mode 100644
index 0000000..4711a91
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/__init__.py
@@ -0,0 +1,17 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+from .cmd_parser import parse_args
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/body_model_defaults.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/body_model_defaults.py
new file mode 100644
index 0000000..658149c
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/body_model_defaults.py
@@ -0,0 +1,107 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+from omegaconf import OmegaConf
+from loguru import logger
+from dataclasses import dataclass
+from .utils_cfg import Variable, Pose
+
+
+@dataclass
+class PCA:
+ num_comps: int = 12
+ flat_hand_mean: bool = False
+
+
+@dataclass
+class PoseWithPCA(Pose):
+ pca: PCA = PCA()
+
+
+@dataclass
+class Shape(Variable):
+ num: int = 10
+
+
+@dataclass
+class Expression(Variable):
+ num: int = 10
+
+
+@dataclass
+class SMPL:
+ betas: Shape = Shape()
+ global_rot: Pose = Pose()
+ body_pose: Pose = Pose()
+ translation: Variable = Variable()
+
+
+@dataclass
+class SMPLH(SMPL):
+ left_hand_pose: PoseWithPCA = PoseWithPCA()
+ right_hand_pose: PoseWithPCA = PoseWithPCA()
+
+
+@dataclass
+class SMPLX(SMPLH):
+ expression: Expression = Expression()
+ jaw_pose: Pose = Pose()
+ leye_pose: Pose = Pose()
+ reye_pose: Pose = Pose()
+
+
+@dataclass
+class MANO:
+ betas: Shape = Shape()
+ wrist_pose: Pose = Pose()
+ hand_pose: PoseWithPCA = PoseWithPCA()
+ translation: Variable = Variable()
+
+
+@dataclass
+class FLAME:
+ betas: Shape = Shape()
+ expression: Expression = Expression()
+ global_rot: Pose = Pose()
+ neck_pose: Pose = Pose()
+ jaw_pose: Pose = Pose()
+ leye_pose: Pose = Pose()
+ reye_pose: Pose = Pose()
+
+
+@dataclass
+class BodyModelConfig:
+ model_type: str = 'smplx'
+ use_compressed: bool = True
+ folder: str = 'models'
+ gender: str = 'neutral'
+ extra_joint_path: str = ''
+ ext: str = 'npz'
+
+ num_expression_coeffs: int = 10
+
+ use_face_contour: bool = True
+ joint_regressor_path: str = ''
+
+ smpl: SMPL = SMPL()
+ star: SMPL = SMPL()
+ smplh: SMPLH = SMPLH()
+ smplx: SMPLX = SMPLX()
+ mano: MANO = MANO()
+ flame: FLAME = FLAME()
+
+
+conf = OmegaConf.structured(BodyModelConfig)
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/cmd_parser.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/cmd_parser.py
new file mode 100644
index 0000000..474903e
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/cmd_parser.py
@@ -0,0 +1,51 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+from __future__ import absolute_import
+from __future__ import division
+
+import sys
+import os
+
+import argparse
+from loguru import logger
+
+from omegaconf import OmegaConf
+from .defaults import conf as default_conf
+
+
+def parse_args(argv=None) -> OmegaConf:
+ arg_formatter = argparse.ArgumentDefaultsHelpFormatter
+
+ description = 'Model transfer script'
+ parser = argparse.ArgumentParser(formatter_class=arg_formatter,
+ description=description)
+
+ parser.add_argument('--exp-cfg', type=str, dest='exp_cfg',
+ help='The configuration of the experiment')
+ parser.add_argument('--exp-opts', default=[], dest='exp_opts',
+ nargs='*',
+ help='Command line arguments')
+
+ cmd_args = parser.parse_args()
+
+ cfg = default_conf.copy()
+ if cmd_args.exp_cfg:
+ cfg.merge_with(OmegaConf.load(cmd_args.exp_cfg))
+ if cmd_args.exp_opts:
+ cfg.merge_with(OmegaConf.from_cli(cmd_args.exp_opts))
+
+ return cfg
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/dataset_defaults.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/dataset_defaults.py
new file mode 100644
index 0000000..ce3e4f3
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/dataset_defaults.py
@@ -0,0 +1,33 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+from omegaconf import OmegaConf
+from dataclasses import dataclass
+
+
+@dataclass
+class MeshFolder:
+ data_folder: str = 'data/meshes'
+
+
+@dataclass
+class DatasetConfig:
+ num_workers: int = 0
+ name: str = 'mesh-folder'
+ mesh_folder: MeshFolder = MeshFolder()
+
+
+conf = OmegaConf.structured(DatasetConfig)
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/defaults.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/defaults.py
new file mode 100644
index 0000000..1656c21
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/defaults.py
@@ -0,0 +1,71 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+from typing import Tuple, Optional
+from copy import deepcopy
+# from yacs.config import CfgNode as CN
+from dataclasses import dataclass
+from omegaconf import OmegaConf
+
+from .loss_defaults import conf as loss_cfg, LossConfig
+from .dataset_defaults import conf as dataset_cfg, DatasetConfig
+from .optim_defaults import conf as optim_cfg, OptimConfig
+from .body_model_defaults import conf as body_model_cfg, BodyModelConfig
+
+
+@dataclass
+class EdgeFitting:
+ per_part: bool = False
+ reduction: str = 'mean'
+
+
+@dataclass
+class VertexFitting:
+ per_part: bool = False
+ reduction: str = 'mean'
+ type: str = 'l2'
+
+
+@dataclass
+class Config:
+ use_cuda: bool = True
+ log_file: str = '/tmp/logs'
+ output_folder: str = 'output'
+ save_verts: bool = True
+ save_joints: bool = True
+ save_mesh: bool = False
+ save_img_summaries: bool = True
+ summary_steps: int = 5
+ degrees: Tuple[float] = (90,)
+ float_type: str = 'float'
+ logger_level: str = 'INFO'
+ interactive: bool = True
+ batch_size: Optional[int] = 1
+ color_path: str = 'data/smpl_with_colors.ply'
+
+ optim: OptimConfig = optim_cfg
+ datasets: DatasetConfig = dataset_cfg
+ losses: LossConfig = loss_cfg
+ body_model: BodyModelConfig = body_model_cfg
+
+ deformation_transfer_path: str = ''
+ mask_ids_fname: str = ''
+
+ per_part: bool = True
+ edge_fitting: EdgeFitting = EdgeFitting()
+
+
+conf = OmegaConf.structured(Config)
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/loss_defaults.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/loss_defaults.py
new file mode 100644
index 0000000..14fa090
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/loss_defaults.py
@@ -0,0 +1,38 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+# from yacs.config import CfgNode as CN
+
+from typing import List, Tuple, Union
+from omegaconf import OmegaConf
+from loguru import logger
+from dataclasses import dataclass, make_dataclass
+
+
+@dataclass
+class LossTemplate:
+ type: str = 'l2'
+ active: bool = False
+ weight: Tuple[float] = (0.0,)
+ requires_grad: bool = True
+ enable: int = 0
+
+
+@dataclass
+class LossConfig:
+ type: str = 'smplify-x'
+
+
+conf = OmegaConf.structured(LossConfig)
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/optim_defaults.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/optim_defaults.py
new file mode 100644
index 0000000..6dfc6ac
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/optim_defaults.py
@@ -0,0 +1,68 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+from typing import Tuple
+from omegaconf import OmegaConf
+from dataclasses import dataclass
+
+
+@dataclass
+class LBFGS:
+ line_search_fn: str = 'strong_wolfe'
+ max_iter: int = 50
+
+
+@dataclass
+class SGD:
+ momentum: float = 0.9
+ nesterov: bool = True
+
+
+@dataclass
+class ADAM:
+ betas: Tuple[float, float] = (0.9, 0.999)
+ eps: float = 1e-08
+ amsgrad: bool = False
+
+
+@dataclass
+class RMSProp:
+ alpha: float = 0.99
+
+
+@dataclass
+class TrustRegionNewtonCG:
+ max_trust_radius: float = 1000
+ initial_trust_radius: float = 0.05
+ eta: float = 0.15
+ gtol: float = 1e-05
+
+
+@dataclass
+class OptimConfig:
+ type: str = 'trust-ncg'
+ lr: float = 1.0
+ gtol: float = 1e-8
+ ftol: float = -1.0
+ maxiters: int = 100
+
+ lbfgs: LBFGS = LBFGS()
+ sgd: SGD = SGD()
+ adam: ADAM = ADAM()
+ trust_ncg: TrustRegionNewtonCG = TrustRegionNewtonCG()
+
+
+conf = OmegaConf.structured(OptimConfig)
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/utils_cfg.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/utils_cfg.py
new file mode 100644
index 0000000..9ea2243
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/config/utils_cfg.py
@@ -0,0 +1,29 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+from typing import Tuple
+from dataclasses import dataclass
+
+
+@dataclass
+class Variable:
+ create: bool = True
+ requires_grad: bool = True
+
+
+@dataclass
+class Pose(Variable):
+ type: str = 'aa'
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/docs/images/smpl_smplx_correspondence.png b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/docs/images/smpl_smplx_correspondence.png
new file mode 100644
index 0000000..305580d
Binary files /dev/null and b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/docs/images/smpl_smplx_correspondence.png differ
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/docs/transfer.md b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/docs/transfer.md
new file mode 100644
index 0000000..38fd0d7
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/docs/transfer.md
@@ -0,0 +1,115 @@
+# Converting SMPL to SMPL-X
+
+
+
+
+
+
+
+
+
+
+
+
+The SMPL body model [1] is in wide use in computer vision and graphics for both
+research and industrial applications. While widely used, SMPL lacks details like
+articulated hands and an expressive face. The SMPL-X model [3] addresses this
+and includes both the face and hands.
+
+Many legacy applications and datasets are built on SMPL and people want to
+"upgrade" them to SMPL-X. While SMPL-X is based on the SMPL technology, they are
+not completely interchangeable.
+
+Importantly the shape and pose parameters of SMPL and SMPL-X seem tantalizingly
+similar. Sadly, you can't just take them from one model and use them with the
+other. In particular, the joint locations in SMPL-X differ from those in SMPL,
+meaning that the pose (theta) parameters are not interchangeable.
+
+Here we describe a tool to convert back and forth between the models. This
+involves fitting one model to the other to recover the right parameters.
+
+The first step in this process is to establish a mapping between SMPL and
+SMPL-X, since their topologies differ. For this, we assume we have a SMPL-X
+template mesh registered to the SMPL template. Now that the two surfaces match,
+we compute and store the following quantities:
+
+* For each SMPL-X vertex find the nearest point on the SMPL mesh and store:
+ * The index $t_i$ of the triangle where the nearest point is located.
+ * Store the barycentric coordinates of the nearest point with respect to
+ the SMPL triangle $\left[a_i, b_i, c_i\right]$.
+
+
+SMPL-X and SMPL share the same topology up to the neck, therefore the Barycentric coordinates of
+these points are a permutation of `[1.0, 0.0, 0.0]`. We also store a mask of
+valid vertices, to remove points that have no match between the two meshes,
+such as the eyeballs or the inner mouth. If we color-code the correspondences
+we end up with the following image, where the left mesh is SMPL and the right
+one is SMPL-X:
+
+
+
+Now that we have established the correspondences between the models, we can fit
+SMPL-X to the SMPL annotations.
+1. The first step is to build a mesh with the SMPL-X topology from the posed
+ SMPL annotations.
+
+ 1. If $t_i$ is the index of the corresponding SMPL triangle for the i-th SMPL-X
+ vertex, then let $f_i \in \mathbb{N}^3$ be the 3 indices of the SMPL vertices that
+ form the triangle.
+ 2. Let $m_i$ be the binary mask value for the validity of this vertex.
+ 2. The i-th vertex is computed using the barycentrics $\left[a_i, b_i, c_i\right]$ as:
+
+ $v_i^{SMPL-X} = a_i * v_{f_i^0}^{SMPL} + b_i * v_{f_i^1}^{SMPL} + c_i * v_{f_i^2}^{SMPL}$
+
+ 2. Now that we have a mesh in SMPL-X topology, we need to find the SMPL-X
+ parameters, i.e. pose $\theta$, shape $\beta$, expression $\psi$ and translation $\gamma$, that best explain it.
+ We use an iterative optimization scheme to
+ recover the parameters:
+
+ 1. Optimize over the pose with a 3D edge term. Make sure that we only use
+ the valid edges, i.e. those whose both end points are found on both
+ meshes:
+
+ $L_1\left(\theta\right) = \sum_{(i, j) \in \mathcal{E}} m_i m_j \left\lVert(v_i - v_j) - (\hat{v}_i - \hat{v}_j) \right\rVert_2^2$
+
+ 2. Optimize over the translation vector $\gamma$ to align the two models:
+
+ $L_2\left({\gamma}\right) = \sum_{i} m_i \left\lVert v_i - \hat{v}_i \right\rVert$
+
+ 3. Optimize over all parameters, to get the tightest possible fit:
+
+ $L_3\left((\theta, \beta, \psi, \gamma)\right) = \sum_{i} m_i \left\lVert v_i - \hat{v}_i \right\rVert_2^2$
+
+
+So now, if you have data in SMPL format, you can convert it to SMPL-X. This
+should allow you to use it for training.
+
+For the inverse mapping, from SMPL-X to
+SMPL, we follow a similar process to generate the correspondences and then optimize
+over the SMPL parameters that best fit the
+transferred mesh. Of course, if you choose to do this, you will lose all
+information about the hands and the face, since SMPL is not able to model this.
+
+For SMPL and SMPL+H [2], the process is easier, since they share the same
+topology. We can therefore skip the first step, since we already know the
+correspondences, compute a SMPL or SMPL+H mesh and estimate the parameters of
+the other model. If we wish to transfer SMPL+H annotations, such as the AMASS
+motion capture data [4], to SMPL-X, then we can use the correspondences of the
+SMPL to SMPL-X mapping.
+
+## Bibliography
+
+[1]: Loper, M., Mahmood, N., Romero, J., Pons-Moll, G., Black, M.J.: SMPL: A
+skinned multi-person linear model. ACM Transactions on Graphics (TOG) - Proceedings of ACM SIGGRAPH Asia 34(6), 248:1–248:16 (2015)
+
+[2]: Romero, J., Tzionas, D., Black, M.J.: Embodied hands: Modeling and capturing
+hands and bodies together. ACM Transactions on Graphics (TOG) - Proceedings
+of ACM SIGGRAPH Asia 36(6), 245:1–245:17 (2017)
+
+[3]: Pavlakos, G., Choutas, V., Ghorbani, N., Bolkart, T., Osman, A.A.A., Tzionas,
+D., Black, M.J.: Expressive body capture: 3D hands, face, and body from a single
+image. In: Proceedings of the IEEE Conference on Computer Vision and Pattern
+Recognition (CVPR). pp. 10967–10977 (2019)
+
+[4]: Mahmood, N., Ghorbani, N., Troje, N.F., Pons-Moll, G., Black, M.J.: Amass:
+Archive of motion capture as surface shapes. ICCV (2019)
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/losses/__init__.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/losses/__init__.py
new file mode 100644
index 0000000..d991ba0
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/losses/__init__.py
@@ -0,0 +1,17 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+from .losses import *
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/losses/losses.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/losses/losses.py
new file mode 100644
index 0000000..57d96fd
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/losses/losses.py
@@ -0,0 +1,149 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+from __future__ import print_function
+from __future__ import absolute_import
+from __future__ import division
+
+import sys
+import time
+from typing import Callable, Iterator, Union, Optional, List
+
+import os.path as osp
+import yaml
+from loguru import logger
+
+import pickle
+
+import numpy as np
+
+import torch
+import torch.autograd as autograd
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .utils import get_reduction_method
+
+__all__ = [
+ 'VertexEdgeLoss',
+ 'build_loss',
+]
+
+
+def build_loss(type='l2', reduction='mean', **kwargs) -> nn.Module:
+ logger.debug(f'Building loss: {type}')
+ if type == 'l2':
+ return WeightedMSELoss(reduction=reduction, **kwargs)
+ elif type == 'vertex-edge':
+ return VertexEdgeLoss(reduction=reduction, **kwargs)
+ elif type == 'l1':
+ return nn.L1Loss()
+ else:
+ raise ValueError(f'Unknown loss type: {type}')
+
+
+class WeightedMSELoss(nn.Module):
+ def __init__(self, reduction='mean', **kwargs):
+ super(WeightedMSELoss, self).__init__()
+ self.reduce_str = reduction
+ self.reduce = get_reduction_method(reduction)
+
+ def forward(self, input, target, weights=None):
+ diff = input - target
+ if weights is None:
+ return diff.pow(2).sum() / diff.shape[0]
+ else:
+ return (
+ weights.unsqueeze(dim=-1) * diff.pow(2)).sum() / diff.shape[0]
+
+
+class VertexEdgeLoss(nn.Module):
+ def __init__(self, norm_type='l2',
+ gt_edges=None,
+ gt_edge_path='',
+ est_edges=None,
+ est_edge_path='',
+ robustifier=None,
+ edge_thresh=0.0, epsilon=1e-8,
+ reduction='sum',
+ **kwargs):
+ super(VertexEdgeLoss, self).__init__()
+
+ assert norm_type in ['l1', 'l2'], 'Norm type must be [l1, l2]'
+ self.norm_type = norm_type
+ self.epsilon = epsilon
+ self.reduction = reduction
+ assert self.reduction in ['sum', 'mean']
+ logger.info(f'Building edge loss with'
+ f' norm_type={norm_type},'
+ f' reduction={reduction},'
+ )
+
+ gt_edge_path = osp.expandvars(gt_edge_path)
+ est_edge_path = osp.expandvars(est_edge_path)
+ assert osp.exists(gt_edge_path) or gt_edges is not None, (
+ 'gt_edges must not be None or gt_edge_path must exist'
+ )
+ assert osp.exists(est_edge_path) or est_edges is not None, (
+ 'est_edges must not be None or est_edge_path must exist'
+ )
+ if osp.exists(gt_edge_path) and gt_edges is None:
+ gt_edges = np.load(gt_edge_path)
+ if osp.exists(est_edge_path) and est_edges is None:
+ est_edges = np.load(est_edge_path)
+
+ self.register_buffer(
+ 'gt_connections', torch.tensor(gt_edges, dtype=torch.long))
+ self.register_buffer(
+ 'est_connections', torch.tensor(est_edges, dtype=torch.long))
+
+ def extra_repr(self):
+ msg = [
+ f'Norm type: {self.norm_type}',
+ ]
+ if self.has_connections:
+ msg.append(
+ f'GT Connections shape: {self.gt_connections.shape}'
+ )
+ msg.append(
+ f'Est Connections shape: {self.est_connections.shape}'
+ )
+ return '\n'.join(msg)
+
+ def compute_edges(self, points, connections):
+ edge_points = torch.index_select(
+ points, 1, connections.view(-1)).reshape(points.shape[0], -1, 2, 3)
+ return edge_points[:, :, 1] - edge_points[:, :, 0]
+
+ def forward(self, gt_vertices, est_vertices, weights=None):
+ gt_edges = self.compute_edges(
+ gt_vertices, connections=self.gt_connections)
+ est_edges = self.compute_edges(
+ est_vertices, connections=self.est_connections)
+
+ raw_edge_diff = (gt_edges - est_edges)
+
+ batch_size = gt_vertices.shape[0]
+ if self.norm_type == 'l2':
+ edge_diff = raw_edge_diff.pow(2)
+ elif self.norm_type == 'l1':
+ edge_diff = raw_edge_diff.abs()
+ else:
+ raise NotImplementedError(
+ f'Loss type not implemented: {self.loss_type}')
+ if self.reduction == 'sum':
+ return edge_diff.sum()
+ elif self.reduction == 'mean':
+ return edge_diff.sum() / batch_size
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/losses/utils.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/losses/utils.py
new file mode 100644
index 0000000..99d9fee
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/losses/utils.py
@@ -0,0 +1,27 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+import torch
+
+def get_reduction_method(reduction='mean'):
+ if reduction == 'mean':
+ return torch.mean
+ elif reduction == 'sum':
+ return torch.sum
+ elif reduction == 'none':
+ return lambda x: x
+ else:
+ raise ValueError('Unknown reduction method: {}'.format(reduction))
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/merge_output.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/merge_output.py
new file mode 100644
index 0000000..f7e4d52
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/merge_output.py
@@ -0,0 +1,88 @@
+# merges the output of the main transfer_model script
+
+import torch
+from pathlib import Path
+import pickle
+from scipy.spatial.transform import Rotation as R
+
+KEYS = [
+"transl",
+"global_orient",
+"body_pose",
+"betas",
+"left_hand_pose",
+"right_hand_pose",
+"jaw_pose",
+"leye_pose",
+"reye_pose",
+"expression",
+"vertices",
+"joints",
+"full_pose",
+"v_shaped",
+"faces"
+]
+
+IGNORED_KEYS = [
+"vertices",
+"faces",
+"v_shaped"
+]
+
+def aggregate_rotmats(x):
+ x = torch.cat(x, dim=0).detach().numpy()
+ s = x.shape[:-2]
+ x = R.from_matrix(x.reshape(-1, 3, 3)).as_rotvec()
+ x = x.reshape(s[0], -1)
+ return x
+
+aggregate_function = {k: lambda x: torch.cat(x, 0).detach().numpy() for k in KEYS}
+aggregate_function["betas"] = lambda x: torch.cat(x, 0).mean(0).detach().numpy()
+
+for k in ["global_orient", "body_pose", "left_hand_pose", "right_hand_pose", "jaw_pose", "full_pose"]:
+ aggregate_function[k] = aggregate_rotmats
+
+def merge(output_dir, gender):
+ output_dir = Path(output_dir)
+ assert output_dir.exists()
+ assert output_dir.is_dir()
+
+ # get list of all pkl files in output_dir with fixed length numeral names
+ pkl_files = [f for f in output_dir.glob("*.pkl") if f.stem != "merged"]
+ pkl_files = [f for f in sorted(pkl_files, key=lambda x: int(x.stem))]
+ assert "merged.pkl" not in [f.name for f in pkl_files]
+
+ merged = {}
+ # iterate over keys and put all values in lists
+ keys = set(KEYS) - set(IGNORED_KEYS)
+ for k in keys:
+ merged[k] = []
+ for pkl_file in pkl_files:
+ with open(pkl_file, "rb") as f:
+ data = pickle.load(f)
+ for k in keys:
+ if k in data:
+ merged[k].append(data[k])
+ b = torch.cat(merged["betas"], 0)
+ print("betas:")
+ for mu, sigma in zip(b.mean(0), b.std(0)):
+ print(" {:.3f} +/- {:.3f}".format(mu, sigma))
+
+ # aggregate all values
+ for k in keys:
+ merged[k] = aggregate_function[k](merged[k])
+
+ # add gender
+ merged["gender"] = gender
+
+ # save merged data to same output_dir
+ with open(output_dir / "merged.pkl", "wb") as f:
+ pickle.dump(merged, f)
+
+if __name__ == '__main__':
+ import argparse
+ parser = argparse.ArgumentParser(description='Merge output of transfer_model script')
+ parser.add_argument('output_dir', type=str, help='output directory of transfer_model script')
+ parser.add_argument('--gender', type=str, choices=['male', 'female', 'neutral'], help='gender of actor in motion sequence')
+ args = parser.parse_args()
+ merge(args.output_dir, args.gender)
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/optimizers/__init__.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/optimizers/__init__.py
new file mode 100644
index 0000000..6adf379
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/optimizers/__init__.py
@@ -0,0 +1,18 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+from .optim_factory import build_optimizer
+from .minimize import minimize
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/optimizers/minimize.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/optimizers/minimize.py
new file mode 100644
index 0000000..dd45c41
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/optimizers/minimize.py
@@ -0,0 +1,86 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+from typing import List, Union, Callable, Optional, Dict
+import torch
+from loguru import logger
+from tqdm import tqdm
+
+from transfer_model.utils import (
+ from_torch, Tensor, Array, rel_change)
+
+
+def minimize(
+ optimizer: torch.optim,
+ closure,
+ params: List[Tensor],
+ summary_closure: Optional[Callable[[], Dict[str, float]]] = None,
+ maxiters=100,
+ ftol=-1.0,
+ gtol=1e-9,
+ interactive=True,
+ summary_steps=10,
+ **kwargs
+):
+ ''' Helper function for running an optimization process
+ Args:
+ - optimizer: The PyTorch optimizer object
+ - closure: The function used to calculate the gradients
+ - params: a list containing the parameters that will be optimized
+ Keyword arguments:
+ - maxiters (100): The maximum number of iterations for the
+ optimizer
+ - ftol: The tolerance for the relative change in the loss
+ function.
+ If it is lower than this value, then the process stops
+ - gtol: The tolerance for the maximum change in the gradient.
+ If the maximum absolute values of the all gradient tensors
+ are less than this, then the process will stop.
+ '''
+ prev_loss = None
+ for n in tqdm(range(maxiters), desc='Fitting iterations'):
+ loss = optimizer.step(closure)
+
+ if n > 0 and prev_loss is not None and ftol > 0:
+ loss_rel_change = rel_change(prev_loss, loss.item())
+
+ if loss_rel_change <= ftol:
+ prev_loss = loss.item()
+ break
+
+ if (all([var.grad.view(-1).abs().max().item() < gtol
+ for var in params if var.grad is not None]) and gtol > 0):
+ prev_loss = loss.item()
+ break
+
+ if interactive and n % summary_steps == 0:
+ logger.info(f'[{n:05d}] Loss: {loss.item():.4f}')
+ if summary_closure is not None:
+ summaries = summary_closure()
+ for key, val in summaries.items():
+ logger.info(f'[{n:05d}] {key}: {val:.4f}')
+
+ prev_loss = loss.item()
+
+ # Save the final step
+ if interactive:
+ logger.info(f'[{n + 1:05d}] Loss: {loss.item():.4f}')
+ if summary_closure is not None:
+ summaries = summary_closure()
+ for key, val in summaries.items():
+ logger.info(f'[{n + 1:05d}] {key}: {val:.4f}')
+
+ return prev_loss
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/optimizers/optim_factory.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/optimizers/optim_factory.py
new file mode 100644
index 0000000..43b161e
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/optimizers/optim_factory.py
@@ -0,0 +1,72 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+import sys
+
+from typing import NewType, List, Dict
+
+import torch
+import torch.optim as optim
+from loguru import logger
+from torchtrustncg import TrustRegion
+
+Tensor = NewType('Tensor', torch.Tensor)
+
+
+def build_optimizer(parameters: List[Tensor],
+ optim_cfg: Dict
+ ) -> Dict:
+ ''' Creates the optimizer
+ '''
+ optim_type = optim_cfg.get('type', 'sgd')
+ logger.info(f'Building: {optim_type.title()}')
+
+ num_params = len(parameters)
+ parameters = list(filter(lambda x: x.requires_grad, parameters))
+ if num_params != len(parameters):
+ logger.info(f'Some parameters have requires_grad off')
+
+ if optim_type == 'adam':
+ optimizer = optim.Adam(parameters, **optim_cfg.get('adam', {}))
+ create_graph = False
+ elif optim_type == 'lbfgs' or optim_type == 'lbfgsls':
+ optimizer = optim.LBFGS(parameters, **optim_cfg.get('lbfgs', {}))
+ create_graph = False
+ elif optim_type == 'trust_ncg' or optim_type == 'trust-ncg':
+ optimizer = TrustRegion(
+ parameters, **optim_cfg.get('trust_ncg', {}))
+ create_graph = True
+ elif optim_type == 'rmsprop':
+ optimizer = optim.RMSprop(parameters, **optim_cfg.get('rmsprop', {}))
+ create_graph = False
+ elif optim_type == 'sgd':
+ optimizer = optim.SGD(parameters, **optim_cfg.get('sgd', {}))
+ create_graph = False
+ else:
+ raise ValueError(f'Optimizer {optim_type} not supported!')
+ return {'optimizer': optimizer, 'create_graph': create_graph}
+
+
+def build_scheduler(optimizer, sched_type='exp',
+ lr_lambda=0.1, **kwargs):
+ if lr_lambda <= 0.0:
+ return None
+
+ if sched_type == 'exp':
+ return optim.lr_scheduler.ExponentialLR(optimizer, lr_lambda)
+ else:
+ raise ValueError('Unknown learning rate' +
+ ' scheduler: '.format(sched_type))
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/requirements.txt b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/requirements.txt
new file mode 100644
index 0000000..068b658
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/requirements.txt
@@ -0,0 +1,10 @@
+numpy>=1.16.2
+torch>=1.0.1.post2
+dataclasses>=0.6
+pyrender>=0.1.23
+shapely
+trimesh>=2.37.6
+open3d
+smplx
+omegaconf
+loguru
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/transfer_model.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/transfer_model.py
new file mode 100644
index 0000000..22db753
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/transfer_model.py
@@ -0,0 +1,396 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+from typing import Optional, Dict, Callable
+import sys
+import numpy as np
+import torch
+import torch.nn as nn
+
+from tqdm import tqdm
+
+from loguru import logger
+from .utils import get_vertices_per_edge
+
+from .optimizers import build_optimizer, minimize
+from .utils import (
+ Tensor, batch_rodrigues, apply_deformation_transfer)
+from .losses import build_loss
+
+
+def summary_closure(gt_vertices, var_dict, body_model, mask_ids=None):
+ param_dict = {}
+ for key, var in var_dict.items():
+ # Decode the axis-angles
+ if 'pose' in key or 'orient' in key:
+ param_dict[key] = batch_rodrigues(
+ var.reshape(-1, 3)).reshape(len(var), -1, 3, 3)
+ else:
+ # Simply pass the variable
+ param_dict[key] = var
+ body_model_output = body_model(
+ return_full_pose=True, get_skin=True, **param_dict)
+ est_vertices = body_model_output['vertices']
+ if mask_ids is not None:
+ est_vertices = est_vertices[:, mask_ids]
+ gt_vertices = gt_vertices[:, mask_ids]
+
+ v2v = (est_vertices - gt_vertices).pow(2).sum(dim=-1).sqrt().mean()
+ return {
+ 'Vertex-to-Vertex': v2v * 1000}
+
+
+def build_model_forward_closure(
+ body_model: nn.Module,
+ var_dict: Dict[str, Tensor],
+ per_part: bool = True,
+ part_key: Optional[str] = None,
+ jidx: Optional[int] = None,
+ part: Optional[Tensor] = None
+) -> Callable:
+ if per_part:
+ cond = part is not None and part_key is not None and jidx is not None
+ assert cond, (
+ 'When per-part is True, "part", "part_key", "jidx" must not be'
+ ' None.'
+ )
+
+ def model_forward():
+ param_dict = {}
+ for key, var in var_dict.items():
+ if part_key == key:
+ param_dict[key] = batch_rodrigues(
+ var.reshape(-1, 3)).reshape(len(var), -1, 3, 3)
+ param_dict[key][:, jidx] = batch_rodrigues(
+ part.reshape(-1, 3)).reshape(-1, 3, 3)
+ else:
+ # Decode the axis-angles
+ if 'pose' in key or 'orient' in key:
+ param_dict[key] = batch_rodrigues(
+ var.reshape(-1, 3)).reshape(len(var), -1, 3, 3)
+ else:
+ # Simply pass the variable
+ param_dict[key] = var
+
+ return body_model(
+ return_full_pose=True, get_skin=True, **param_dict)
+ else:
+ def model_forward():
+ param_dict = {}
+ for key, var in var_dict.items():
+ # Decode the axis-angles
+ if 'pose' in key or 'orient' in key:
+ param_dict[key] = batch_rodrigues(
+ var.reshape(-1, 3)).reshape(len(var), -1, 3, 3)
+ else:
+ # Simply pass the variable
+ param_dict[key] = var
+
+ return body_model(return_full_pose=True, get_skin=True,
+ **param_dict)
+ return model_forward
+
+
+def build_edge_closure(
+ body_model: nn.Module,
+ var_dict: Dict[str, Tensor],
+ edge_loss: nn.Module,
+ optimizer_dict,
+ gt_vertices: Tensor,
+ per_part: bool = True,
+ part_key: Optional[str] = None,
+ jidx: Optional[int] = None,
+ part: Optional[Tensor] = None
+) -> Callable:
+ ''' Builds the closure for the edge objective
+ '''
+ optimizer = optimizer_dict['optimizer']
+ create_graph = optimizer_dict['create_graph']
+
+ if per_part:
+ params_to_opt = [part]
+ else:
+ params_to_opt = [p for key, p in var_dict.items() if 'pose' in key]
+
+ model_forward = build_model_forward_closure(
+ body_model, var_dict, per_part=per_part, part_key=part_key,
+ jidx=jidx, part=part)
+
+ def closure(backward=True):
+ if backward:
+ optimizer.zero_grad()
+
+ body_model_output = model_forward()
+ est_vertices = body_model_output['vertices']
+
+ loss = edge_loss(est_vertices, gt_vertices)
+ if backward:
+ if create_graph:
+ # Use this instead of .backward to avoid GPU memory leaks
+ grads = torch.autograd.grad(
+ loss, params_to_opt, create_graph=True)
+ torch.autograd.backward(
+ params_to_opt, grads, create_graph=True)
+ else:
+ loss.backward()
+
+ return loss
+ return closure
+
+
+def build_vertex_closure(
+ body_model: nn.Module,
+ var_dict: Dict[str, Tensor],
+ optimizer_dict,
+ gt_vertices: Tensor,
+ vertex_loss: nn.Module,
+ mask_ids=None,
+ per_part: bool = True,
+ part_key: Optional[str] = None,
+ jidx: Optional[int] = None,
+ part: Optional[Tensor] = None,
+ params_to_opt: Optional[Tensor] = None,
+) -> Callable:
+ ''' Builds the closure for the vertex objective
+ '''
+ optimizer = optimizer_dict['optimizer']
+ create_graph = optimizer_dict['create_graph']
+
+ model_forward = build_model_forward_closure(
+ body_model, var_dict, per_part=per_part, part_key=part_key,
+ jidx=jidx, part=part)
+
+ if params_to_opt is None:
+ params_to_opt = [p for key, p in var_dict.items()]
+
+ def closure(backward=True):
+ if backward:
+ optimizer.zero_grad()
+
+ body_model_output = model_forward()
+ est_vertices = body_model_output['vertices']
+
+ loss = vertex_loss(
+ est_vertices[:, mask_ids] if mask_ids is not None else
+ est_vertices,
+ gt_vertices[:, mask_ids] if mask_ids is not None else gt_vertices)
+ if backward:
+ if create_graph:
+ # Use this instead of .backward to avoid GPU memory leaks
+ grads = torch.autograd.grad(
+ loss, params_to_opt, create_graph=True)
+ torch.autograd.backward(
+ params_to_opt, grads, create_graph=True)
+ else:
+ loss.backward()
+
+ return loss
+ return closure
+
+
+def get_variables(
+ batch_size: int,
+ body_model: nn.Module,
+ dtype: torch.dtype = torch.float32
+) -> Dict[str, Tensor]:
+ var_dict = {}
+
+ device = next(body_model.buffers()).device
+
+ if (body_model.name() == 'SMPL' or body_model.name() == 'SMPL+H' or
+ body_model.name() == 'SMPL-X'):
+ var_dict.update({
+ 'transl': torch.zeros(
+ [batch_size, 3], device=device, dtype=dtype),
+ 'global_orient': torch.zeros(
+ [batch_size, 1, 3], device=device, dtype=dtype),
+ 'body_pose': torch.zeros(
+ [batch_size, body_model.NUM_BODY_JOINTS, 3],
+ device=device, dtype=dtype),
+ 'betas': torch.zeros([batch_size, body_model.num_betas],
+ dtype=dtype, device=device),
+ })
+
+ if body_model.name() == 'SMPL+H' or body_model.name() == 'SMPL-X':
+ var_dict.update(
+ left_hand_pose=torch.zeros(
+ [batch_size, body_model.NUM_HAND_JOINTS, 3], device=device,
+ dtype=dtype),
+ right_hand_pose=torch.zeros(
+ [batch_size, body_model.NUM_HAND_JOINTS, 3], device=device,
+ dtype=dtype),
+ )
+
+ if body_model.name() == 'SMPL-X':
+ var_dict.update(
+ jaw_pose=torch.zeros([batch_size, 1, 3],
+ device=device, dtype=dtype),
+ leye_pose=torch.zeros([batch_size, 1, 3],
+ device=device, dtype=dtype),
+ reye_pose=torch.zeros([batch_size, 1, 3],
+ device=device, dtype=dtype),
+ expression=torch.zeros(
+ [batch_size, body_model.num_expression_coeffs],
+ device=device, dtype=dtype),
+ )
+
+ # Toggle gradients to True
+ for key, val in var_dict.items():
+ val.requires_grad_(True)
+
+ return var_dict
+
+
+def run_fitting(
+ exp_cfg,
+ batch: Dict[str, Tensor],
+ body_model: nn.Module,
+ def_matrix: Tensor,
+ mask_ids: Optional = None
+) -> Dict[str, Tensor]:
+ ''' Runs fitting
+ '''
+ vertices = batch['vertices']
+ faces = batch['faces']
+
+ batch_size = len(vertices)
+ dtype, device = vertices.dtype, vertices.device
+ summary_steps = exp_cfg.get('summary_steps')
+ interactive = exp_cfg.get('interactive')
+
+ # Get the parameters from the model
+ var_dict = get_variables(batch_size, body_model)
+
+ # Build the optimizer object for the current batch
+ optim_cfg = exp_cfg.get('optim', {})
+
+ def_vertices = apply_deformation_transfer(def_matrix, vertices, faces)
+
+ if mask_ids is None:
+ f_sel = np.ones_like(body_model.faces[:, 0], dtype=np.bool_)
+ else:
+ f_per_v = [[] for _ in range(body_model.get_num_verts())]
+ [f_per_v[vv].append(iff) for iff, ff in enumerate(body_model.faces)
+ for vv in ff]
+ f_sel = list(set(tuple(sum([f_per_v[vv] for vv in mask_ids], []))))
+ vpe = get_vertices_per_edge(
+ body_model.v_template.detach().cpu().numpy(), body_model.faces[f_sel])
+
+ def log_closure():
+ return summary_closure(def_vertices, var_dict, body_model,
+ mask_ids=mask_ids)
+
+ edge_fitting_cfg = exp_cfg.get('edge_fitting', {})
+ edge_loss = build_loss(type='vertex-edge', gt_edges=vpe, est_edges=vpe,
+ **edge_fitting_cfg)
+ edge_loss = edge_loss.to(device=device)
+
+ vertex_fitting_cfg = exp_cfg.get('vertex_fitting', {})
+ vertex_loss = build_loss(**vertex_fitting_cfg)
+ vertex_loss = vertex_loss.to(device=device)
+
+ per_part = edge_fitting_cfg.get('per_part', True)
+ logger.info(f'Per-part: {per_part}')
+ # Optimize edge-based loss to initialize pose
+ if per_part:
+ for key, var in tqdm(var_dict.items(), desc='Parts'):
+ if 'pose' not in key:
+ continue
+
+ for jidx in tqdm(range(var.shape[1]), desc='Joints'):
+ part = torch.zeros(
+ [batch_size, 3], dtype=dtype, device=device,
+ requires_grad=True)
+ # Build the optimizer for the current part
+ optimizer_dict = build_optimizer([part], optim_cfg)
+ closure = build_edge_closure(
+ body_model, var_dict, edge_loss, optimizer_dict,
+ def_vertices, per_part=per_part, part_key=key, jidx=jidx,
+ part=part)
+
+ minimize(optimizer_dict['optimizer'], closure,
+ params=[part],
+ summary_closure=log_closure,
+ summary_steps=summary_steps,
+ interactive=interactive,
+ **optim_cfg)
+ with torch.no_grad():
+ var[:, jidx] = part
+ else:
+ optimizer_dict = build_optimizer(list(var_dict.values()), optim_cfg)
+ closure = build_edge_closure(
+ body_model, var_dict, edge_loss, optimizer_dict,
+ def_vertices, per_part=per_part)
+
+ minimize(optimizer_dict['optimizer'], closure,
+ params=var_dict.values(),
+ summary_closure=log_closure,
+ summary_steps=summary_steps,
+ interactive=interactive,
+ **optim_cfg)
+
+ if 'translation' in var_dict:
+ optimizer_dict = build_optimizer([var_dict['translation']], optim_cfg)
+ closure = build_vertex_closure(
+ body_model, var_dict,
+ optimizer_dict,
+ def_vertices,
+ vertex_loss=vertex_loss,
+ mask_ids=mask_ids,
+ per_part=False,
+ params_to_opt=[var_dict['translation']],
+ )
+ # Optimize translation
+ minimize(optimizer_dict['optimizer'],
+ closure,
+ params=[var_dict['translation']],
+ summary_closure=log_closure,
+ summary_steps=summary_steps,
+ interactive=interactive,
+ **optim_cfg)
+
+ # Optimize all model parameters with vertex-based loss
+ optimizer_dict = build_optimizer(list(var_dict.values()), optim_cfg)
+ closure = build_vertex_closure(
+ body_model, var_dict,
+ optimizer_dict,
+ def_vertices,
+ vertex_loss=vertex_loss,
+ per_part=False,
+ mask_ids=mask_ids)
+ minimize(optimizer_dict['optimizer'], closure,
+ params=list(var_dict.values()),
+ summary_closure=log_closure,
+ summary_steps=summary_steps,
+ interactive=interactive,
+ **optim_cfg)
+
+ param_dict = {}
+ for key, var in var_dict.items():
+ # Decode the axis-angles
+ if 'pose' in key or 'orient' in key:
+ param_dict[key] = batch_rodrigues(
+ var.reshape(-1, 3)).reshape(len(var), -1, 3, 3)
+ else:
+ # Simply pass the variable
+ param_dict[key] = var
+
+ body_model_output = body_model(
+ return_full_pose=True, get_skin=True, **param_dict)
+ var_dict.update(body_model_output)
+ var_dict['faces'] = body_model.faces
+
+ return var_dict
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/__init__.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/__init__.py
new file mode 100644
index 0000000..78c1bfe
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/__init__.py
@@ -0,0 +1,25 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+from .np_utils import to_np, rel_change
+from .torch_utils import from_torch
+from .timer import Timer, timer_decorator
+from .typing import *
+from .pose_utils import batch_rodrigues, batch_rot2aa
+from .metrics import v2v
+from .def_transfer import read_deformation_transfer, apply_deformation_transfer
+from .mesh_utils import get_vertices_per_edge
+from .o3d_utils import np_mesh_to_o3d
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/def_transfer.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/def_transfer.py
new file mode 100644
index 0000000..30d2f62
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/def_transfer.py
@@ -0,0 +1,75 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+import os
+import os.path as osp
+import pickle
+
+import numpy as np
+import torch
+from loguru import logger
+
+from .typing import Tensor
+
+
+def read_deformation_transfer(
+ deformation_transfer_path: str,
+ device=None,
+ use_normal: bool = False,
+) -> Tensor:
+ ''' Reads a deformation transfer
+ '''
+ if device is None:
+ device = torch.device('cpu')
+ assert osp.exists(deformation_transfer_path), (
+ 'Deformation transfer path does not exist:'
+ f' {deformation_transfer_path}')
+ logger.info(
+ f'Loading deformation transfer from: {deformation_transfer_path}')
+ # Read the deformation transfer matrix
+ with open(deformation_transfer_path, 'rb') as f:
+ def_transfer_setup = pickle.load(f, encoding='latin1')
+ if 'mtx' in def_transfer_setup:
+ def_matrix = def_transfer_setup['mtx']
+ if hasattr(def_matrix, 'todense'):
+ def_matrix = def_matrix.todense()
+ def_matrix = np.array(def_matrix, dtype=np.float32)
+ if not use_normal:
+ num_verts = def_matrix.shape[1] // 2
+ def_matrix = def_matrix[:, :num_verts]
+ elif 'matrix' in def_transfer_setup:
+ def_matrix = def_transfer_setup['matrix']
+ else:
+ valid_keys = ['mtx', 'matrix']
+ raise KeyError(f'Deformation transfer setup must contain {valid_keys}')
+
+ def_matrix = torch.tensor(def_matrix, device=device, dtype=torch.float32)
+ return def_matrix
+
+
+def apply_deformation_transfer(
+ def_matrix: Tensor,
+ vertices: Tensor,
+ faces: Tensor,
+ use_normals=False
+) -> Tensor:
+ ''' Applies the deformation transfer on the given meshes
+ '''
+ if use_normals:
+ raise NotImplementedError
+ else:
+ def_vertices = torch.einsum('mn,bni->bmi', [def_matrix, vertices])
+ return def_vertices
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/mesh_utils.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/mesh_utils.py
new file mode 100644
index 0000000..4d681b2
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/mesh_utils.py
@@ -0,0 +1,59 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Code from Chumpy and OpenDR. Placed here to avoid chumpy dependency
+# The original code can be found in https://github.com/MPI-IS/mesh
+import numpy as np
+import scipy.sparse as sp
+
+
+def row(A):
+ return A.reshape((1, -1))
+
+
+def col(A):
+ return A.reshape((-1, 1))
+
+
+def get_vert_connectivity(mesh_v, mesh_f):
+ """Returns a sparse matrix (of size #verts x #verts) where each nonzero
+ element indicates a neighborhood relation. For example, if there is a
+ nonzero element in position (15,12), that means vertex 15 is connected
+ by an edge to vertex 12."""
+
+ vpv = sp.csc_matrix((len(mesh_v), len(mesh_v)))
+
+ # for each column in the faces...
+ for i in range(3):
+ IS = mesh_f[:, i]
+ JS = mesh_f[:, (i + 1) % 3]
+ data = np.ones(len(IS))
+ ij = np.vstack((row(IS.flatten()), row(JS.flatten())))
+ mtx = sp.csc_matrix((data, ij), shape=vpv.shape)
+ vpv = vpv + mtx + mtx.T
+
+ return vpv
+
+
+def get_vertices_per_edge(mesh_v, mesh_f):
+ """Returns an Ex2 array of adjacencies between vertices, where
+ each element in the array is a vertex index. Each edge is included
+ only once. If output of get_faces_per_edge is provided, this is used to
+ avoid call to get_vert_connectivity()"""
+
+ vc = sp.coo_matrix(get_vert_connectivity(mesh_v, mesh_f))
+ result = np.hstack((col(vc.row), col(vc.col)))
+ result = result[result[:, 0] < result[:, 1]] # for uniqueness
+
+ return result
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/metrics.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/metrics.py
new file mode 100644
index 0000000..d7f8209
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/metrics.py
@@ -0,0 +1,25 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+import numpy as np
+import torch
+
+
+def v2v(x, y):
+ if torch.is_tensor(x):
+ return (x - y).pow(2).sum(dim=-1).sqrt().mean()
+ else:
+ return np.sqrt(np.power(x - y, 2)).sum(axis=-1).mean()
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/np_utils.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/np_utils.py
new file mode 100644
index 0000000..d0cfa56
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/np_utils.py
@@ -0,0 +1,34 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import numpy as np
+
+
+def rel_change(prev_val, curr_val):
+ return (prev_val - curr_val) / max([np.abs(prev_val), np.abs(curr_val), 1])
+
+
+def max_grad_change(grad_arr):
+ return grad_arr.abs().max()
+
+
+def to_np(array, dtype=np.float32):
+ if hasattr(array, 'todense'):
+ array = array.todense()
+ return np.array(array, dtype=dtype)
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/o3d_utils.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/o3d_utils.py
new file mode 100644
index 0000000..758f856
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/o3d_utils.py
@@ -0,0 +1,34 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+import open3d as o3d
+import torch
+
+Vector3d = o3d.utility.Vector3dVector
+Vector3i = o3d.utility.Vector3iVector
+
+Mesh = o3d.geometry.TriangleMesh
+
+
+def np_mesh_to_o3d(vertices, faces):
+ if torch.is_tensor(vertices):
+ vertices = vertices.detach().cpu().numpy()
+ if torch.is_tensor(faces):
+ faces = faces.detach().cpu().numpy()
+ mesh = Mesh()
+ mesh.vertices = Vector3d(vertices)
+ mesh.triangles = Vector3i(faces)
+ return mesh
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/pose_utils.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/pose_utils.py
new file mode 100644
index 0000000..a2dc90b
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/pose_utils.py
@@ -0,0 +1,147 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+import sys
+from typing import NewType, List, Dict, Optional
+import os
+import os.path as osp
+
+import pickle
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from omegaconf import OmegaConf
+from loguru import logger
+
+from .typing import Tensor
+
+
+def rotation_matrix_to_cont_repr(x: Tensor) -> Tensor:
+ assert len(x.shape) == 3, (
+ f'Expects an array of size Bx3x3, but received {x.shape}')
+ return x[:, :3, :2]
+
+
+def cont_repr_to_rotation_matrix(
+ x: Tensor
+) -> Tensor:
+ ''' Converts tensor in continous representation to rotation matrices
+ '''
+ batch_size = x.shape[0]
+ reshaped_input = x.view(-1, 3, 2)
+
+ # Normalize the first vector
+ b1 = F.normalize(reshaped_input[:, :, 0].clone(), dim=1)
+
+ dot_prod = torch.sum(
+ b1 * reshaped_input[:, :, 1].clone(), dim=1, keepdim=True)
+ # Compute the second vector by finding the orthogonal complement to it
+ b2 = F.normalize(reshaped_input[:, :, 1] - dot_prod * b1, dim=1)
+ # Finish building the basis by taking the cross product
+ b3 = torch.cross(b1, b2, dim=1)
+ rot_mats = torch.stack([b1, b2, b3], dim=-1)
+
+ return rot_mats.view(batch_size, -1, 3, 3)
+
+
+def batch_rodrigues(
+ rot_vecs: Tensor,
+ epsilon: float = 1e-8
+) -> Tensor:
+ ''' Calculates the rotation matrices for a batch of rotation vectors
+ Parameters
+ ----------
+ rot_vecs: torch.tensor Nx3
+ array of N axis-angle vectors
+ Returns
+ -------
+ R: torch.tensor Nx3x3
+ The rotation matrices for the given axis-angle parameters
+ '''
+ assert len(rot_vecs.shape) == 2, (
+ f'Expects an array of size Bx3, but received {rot_vecs.shape}')
+
+ batch_size = rot_vecs.shape[0]
+ device = rot_vecs.device
+ dtype = rot_vecs.dtype
+
+ angle = torch.norm(rot_vecs + epsilon, dim=1, keepdim=True, p=2)
+ rot_dir = rot_vecs / angle
+
+ cos = torch.unsqueeze(torch.cos(angle), dim=1)
+ sin = torch.unsqueeze(torch.sin(angle), dim=1)
+
+ # Bx1 arrays
+ rx, ry, rz = torch.split(rot_dir, 1, dim=1)
+ K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device)
+
+ zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device)
+ K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
+ .view((batch_size, 3, 3))
+
+ ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
+ rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
+ return rot_mat
+
+
+def batch_rot2aa(
+ Rs: Tensor, epsilon: float = 1e-7
+) -> Tensor:
+ """
+ Rs is B x 3 x 3
+ void cMathUtil::RotMatToAxisAngle(const tMatrix& mat, tVector& out_axis,
+ double& out_theta)
+ {
+ double c = 0.5 * (mat(0, 0) + mat(1, 1) + mat(2, 2) - 1);
+ c = cMathUtil::Clamp(c, -1.0, 1.0);
+
+ out_theta = std::acos(c);
+
+ if (std::abs(out_theta) < 0.00001)
+ {
+ out_axis = tVector(0, 0, 1, 0);
+ }
+ else
+ {
+ double m21 = mat(2, 1) - mat(1, 2);
+ double m02 = mat(0, 2) - mat(2, 0);
+ double m10 = mat(1, 0) - mat(0, 1);
+ double denom = std::sqrt(m21 * m21 + m02 * m02 + m10 * m10);
+ out_axis[0] = m21 / denom;
+ out_axis[1] = m02 / denom;
+ out_axis[2] = m10 / denom;
+ out_axis[3] = 0;
+ }
+ }
+ """
+
+ cos = 0.5 * (torch.einsum('bii->b', [Rs]) - 1)
+ cos = torch.clamp(cos, -1 + epsilon, 1 - epsilon)
+
+ theta = torch.acos(cos)
+
+ m21 = Rs[:, 2, 1] - Rs[:, 1, 2]
+ m02 = Rs[:, 0, 2] - Rs[:, 2, 0]
+ m10 = Rs[:, 1, 0] - Rs[:, 0, 1]
+ denom = torch.sqrt(m21 * m21 + m02 * m02 + m10 * m10 + epsilon)
+
+ axis0 = torch.where(torch.abs(theta) < 0.00001, m21, m21 / denom)
+ axis1 = torch.where(torch.abs(theta) < 0.00001, m02, m02 / denom)
+ axis2 = torch.where(torch.abs(theta) < 0.00001, m10, m10 / denom)
+
+ return theta.unsqueeze(1) * torch.stack([axis0, axis1, axis2], 1)
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/timer.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/timer.py
new file mode 100644
index 0000000..027de5b
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/timer.py
@@ -0,0 +1,60 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+import time
+import numpy as np
+import torch
+
+from loguru import logger
+
+
+class Timer(object):
+ def __init__(self, name='', sync=False):
+ super(Timer, self).__init__()
+ self.elapsed = []
+ self.name = name
+ self.sync = sync
+
+ def __enter__(self):
+ if self.sync:
+ torch.cuda.synchronize()
+ self.start = time.perf_counter()
+
+ def __exit__(self, type, value, traceback):
+ if self.sync:
+ torch.cuda.synchronize()
+ elapsed = time.perf_counter() - self.start
+ self.elapsed.append(elapsed)
+ logger.info(f'[{self.name}]: {np.mean(self.elapsed):.3f}')
+
+
+def timer_decorator(sync=False, name=''):
+ def wrapper(method):
+ elapsed = []
+
+ def timed(*args, **kw):
+ if sync:
+ torch.cuda.synchronize()
+ ts = time.perf_counter()
+ result = method(*args, **kw)
+ if sync:
+ torch.cuda.synchronize()
+ te = time.perf_counter()
+ elapsed.append(te - ts)
+ logger.info(f'[{name}]: {np.mean(elapsed):.3f}')
+ return result
+ return timed
+ return wrapper
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/torch_utils.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/torch_utils.py
new file mode 100644
index 0000000..ffa820d
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/torch_utils.py
@@ -0,0 +1,24 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+import numpy as np
+import torch
+
+
+def from_torch(x, dtype=np.float32):
+ if torch.is_tensor(x):
+ x = x.detach().cpu().numpy()
+ return x.astype(dtype)
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/typing.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/typing.py
new file mode 100644
index 0000000..c6ba310
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/utils/typing.py
@@ -0,0 +1,27 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: Vassilis Choutas, vassilis.choutas@tuebingen.mpg.de
+
+from typing import NewType, List, Union
+import numpy as np
+import torch
+
+__all__ = [
+ 'Tensor',
+ 'Array',
+]
+
+Tensor = NewType('Tensor', torch.Tensor)
+Array = NewType('Array', np.ndarray)
diff --git a/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/view_pkl.py b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/view_pkl.py
new file mode 100644
index 0000000..0cafdef
--- /dev/null
+++ b/worldbench/third_party/drivestudio/third_party/smplx/transfer_model/view_pkl.py
@@ -0,0 +1,140 @@
+import os.path as osp
+import argparse
+
+import numpy as np
+import torch
+
+import pyrender
+import trimesh
+
+import smplx
+
+from tqdm.auto import tqdm, trange
+
+from pathlib import Path
+
+def main(model_folder,
+ motion_file,
+ model_type='smplx',
+ ext='npz',
+ gender='neutral',
+ plot_joints=False,
+ num_betas=10,
+ sample_expression=True,
+ num_expression_coeffs=10,
+ use_face_contour=False):
+
+ # open motion file
+ motion = np.load(motion_file, allow_pickle=True)
+ _motion = {}
+ for k,v in motion.items():
+ if isinstance(v, np.ndarray):
+ print(k, motion[k].shape, motion[k].dtype)
+ if motion[k].dtype in ("/ at 12 fps
+ 2) Renders an extra left-shifted (−1 m) lateral trajectory (front camera) to the same folder
+ 3) Exports original-view depth video aligned with full_set__rgbs.mp4
+ 4) Saves raw depth (.npy) for the original view
+
+Requirements:
+ - PYTHONPATH points to repo root when running this script.
+
+Usage:
+ python tools/batch_eval_nus.py \
+ --root work_dirs \
+ --pattern "drivestudio-nus-*" \
+ --fps 12 \
+ --raw-depth-dir raw_depth
+
+The script overwrites existing files with the same names.
+"""
+from __future__ import annotations
+
+import argparse
+import glob
+import os
+import re
+import sys
+from typing import Dict, List, Optional
+
+import numpy as np
+
+import torch
+from omegaconf import OmegaConf
+
+from datasets.driving_dataset import DrivingDataset
+from utils.misc import import_str
+from utils.camera import lateral_offset_trajectory
+from models.video_utils import render_novel_views, render_images, save_videos
+
+
+def find_checkpoint(run_dir: str) -> Optional[str]:
+ final = os.path.join(run_dir, "checkpoint_final.pth")
+ if os.path.isfile(final):
+ return final
+ cands = sorted(glob.glob(os.path.join(run_dir, "checkpoint_*.pth")))
+ return cands[-1] if cands else None
+
+
+def load_cfg(log_dir: str) -> OmegaConf:
+ cfg_path = os.path.join(log_dir, "config.yaml")
+ if not os.path.isfile(cfg_path):
+ raise FileNotFoundError(f"config.yaml not found in {log_dir}")
+ return OmegaConf.load(cfg_path)
+
+
+def build_dataset_trainer(cfg: OmegaConf, device: torch.device):
+ dataset = DrivingDataset(data_cfg=cfg.data)
+ trainer = import_str(cfg.trainer.type)(
+ **cfg.trainer,
+ num_timesteps=dataset.num_img_timesteps,
+ model_config=cfg.model,
+ num_train_images=len(dataset.train_image_set),
+ num_full_images=len(dataset.full_image_set),
+ test_set_indices=dataset.test_timesteps,
+ scene_aabb=dataset.get_aabb().reshape(2, 3),
+ device=device,
+ )
+ return dataset, trainer
+
+
+def ensure_dir(p: str):
+ os.makedirs(p, exist_ok=True)
+
+
+def remove_if_exists(path: str):
+ try:
+ if os.path.isfile(path):
+ os.remove(path)
+ except Exception:
+ pass
+
+
+def render_three_builtins_novel(dataset: DrivingDataset, trainer, fps: int, out_dir: str, skip_existing: bool):
+ # three built-ins defined in eval: s_curve, front_center_interp, lateral_offset
+ trajs = dataset.get_novel_render_traj(
+ traj_types=["s_curve", "front_center_interp", "lateral_offset"],
+ target_frames=dataset.frame_num,
+ )
+ ensure_dir(out_dir)
+ for tname, traj in trajs.items():
+ render_data = dataset.prepare_novel_view_render_data(traj, cam_id=0)
+ save_path = os.path.join(out_dir, f"{tname}.mp4")
+ if skip_existing and os.path.isfile(save_path):
+ print(f"[Skip] novel {tname} exists: {save_path}")
+ continue
+ remove_if_exists(save_path)
+ render_novel_views(trainer, render_data, save_path, fps=fps)
+
+
+def render_leftshift_novel(dataset: DrivingDataset, trainer, fps: int, out_dir: str, skip_existing: bool):
+ # Build a left-shifted (−1 m) lateral trajectory using front camera as reference
+ per_cam_poses: Dict[int, torch.Tensor] = {
+ cam_id: dataset.pixel_source.camera_data[cam_id].cam_to_worlds
+ for cam_id in dataset.pixel_source.camera_list
+ }
+ traj_left = lateral_offset_trajectory(
+ dataset_type=dataset.type,
+ per_cam_poses=per_cam_poses,
+ original_frames=dataset.frame_num,
+ target_frames=dataset.frame_num,
+ offset_distance=-1.0,
+ )
+ render_data = dataset.prepare_novel_view_render_data(traj_left, cam_id=0)
+ ensure_dir(out_dir)
+ save_path = os.path.join(out_dir, "lateral_offset_left.mp4")
+ if skip_existing and os.path.isfile(save_path):
+ print(f"[Skip] novel lateral_offset_left exists: {save_path}")
+ return
+ remove_if_exists(save_path)
+ render_novel_views(trainer, render_data, save_path, fps=fps)
+
+
+def export_full_depth(cfg: OmegaConf, dataset: DrivingDataset, trainer, log_dir: str, step: int, raw_depth_dir: str, skip_existing: bool):
+ # Depth video aligned with full_set__rgbs.mp4 naming
+ videos_dir = os.path.join(log_dir, "videos")
+ ensure_dir(videos_dir)
+ base_mp4 = os.path.join(videos_dir, f"full_set_{step}.mp4")
+ depth_mp4 = base_mp4.replace(".mp4", "_depths.mp4")
+ out_base = os.path.join(log_dir, raw_depth_dir)
+ raw_dir = os.path.join(out_base, f"full_set_{step}")
+ # We export raw depth only for front camera frames in this script.
+ expected_n = dataset.num_img_timesteps # CAM_FRONT only
+ have_depth_mp4 = os.path.isfile(depth_mp4)
+ have_raw = os.path.isdir(raw_dir) and len(glob.glob(os.path.join(raw_dir, "*.npy"))) >= expected_n
+ if skip_existing and have_depth_mp4 and have_raw:
+ print(f"[Skip] full-set depth video and raw already exist for step {step}")
+ return
+ # Render if needed
+ results = render_images(
+ trainer=trainer,
+ dataset=dataset.full_image_set,
+ compute_metrics=False,
+ compute_error_map=False,
+ vis_indices=None,
+ save_image_pairs=False,
+ )
+ # Save depth video unless skipping due to existing
+ if not (skip_existing and have_depth_mp4):
+ save_videos(
+ results,
+ base_mp4,
+ layout=dataset.layout,
+ num_timestamps=dataset.num_img_timesteps,
+ keys=["depths"],
+ num_cams=dataset.pixel_source.num_cams,
+ save_seperate_video=True,
+ fps=(cfg.render.fps if hasattr(cfg, "render") and "fps" in cfg.render else 10),
+ verbose=True,
+ )
+ # Save raw depth for front camera unless skipping
+ if not (skip_existing and have_raw):
+ ensure_dir(raw_dir)
+ if "depths" in results:
+ for i, depth in enumerate(results["depths"]):
+ cam_name = results["cam_names"][i]
+ if cam_name == 'CAM_FRONT':
+ np.save(os.path.join(raw_dir, f"{i:05d}_{cam_name}.npy"), depth)
+
+
+def main():
+ parser = argparse.ArgumentParser("Batch eval for nuScenes 3DGS")
+ parser.add_argument("--root", type=str, default="work_dirs", help="Root directory of work dirs")
+ parser.add_argument(
+ "--pattern", type=str, default="drivestudio-nus-*", help="Glob pattern to select method dirs"
+ )
+ parser.add_argument("--fps", type=int, default=12, help="FPS for novel view videos")
+ parser.add_argument(
+ "--raw-depth-dir",
+ type=str,
+ default="raw_depth",
+ help="Directory under each run's log_dir to save raw depth npy",
+ )
+ parser.add_argument("--skip-existing", action="store_true", help="Skip rendering if target outputs already exist")
+ parser.add_argument("--verbose", action="store_true", help="Print progress information")
+ args = parser.parse_args()
+
+ root = args.root
+ method_dirs = sorted(glob.glob(os.path.join(root, args.pattern)))
+ if not method_dirs:
+ print(f"No method dirs matched under {root} with pattern {args.pattern}")
+ sys.exit(0)
+
+ for mdir in method_dirs:
+ for run in sorted(os.listdir(mdir)):
+ run_dir = os.path.join(mdir, run)
+ if not os.path.isdir(run_dir):
+ continue
+
+ ckpt = find_checkpoint(run_dir)
+ if ckpt is None:
+ continue
+ log_dir = os.path.dirname(ckpt)
+
+ step_str = ckpt.split("_")[-1].split(".")[0]
+ if step_str == 'final':
+ step_str = "30000"
+ novel_dir = os.path.join(log_dir, "videos", f"novel_{step_str}_benchmark")
+ depth_dir = os.path.join(log_dir, "videos", f"full_set_{step_str}_depths.mp4")
+ raw_depth_dir = os.path.join(log_dir, args.raw_depth_dir, f"full_set_{step_str}")
+
+ if args.skip_existing:
+ novel_views = [
+ "s_curve.mp4",
+ "front_center_interp.mp4",
+ "lateral_offset.mp4",
+ "lateral_offset_left.mp4",
+ ]
+ novel_done = all(os.path.isfile(os.path.join(novel_dir, view)) for view in novel_views)
+ depth_done = os.path.isfile(depth_dir)
+ raw_done = (not args.raw_depth_dir) or (
+ os.path.isdir(raw_depth_dir) and any(fname.endswith(".npy") for fname in os.listdir(raw_depth_dir))
+ )
+ if novel_done and depth_done and raw_done:
+ if args.verbose:
+ print(f"[Skip] {log_dir} already processed.")
+ continue
+
+ if args.verbose:
+ print(f"[Eval] {ckpt}")
+
+ # Load cfg and build dataset/trainer
+ cfg = load_cfg(log_dir)
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ dataset, trainer = build_dataset_trainer(cfg, device)
+ trainer.resume_from_checkpoint(ckpt, load_only_model=True)
+ step = trainer.step
+
+ # Novel views: three built-ins + left shift (front camera)
+ novel_dir = os.path.join(log_dir, "videos", f"novel_{step}_benchmark")
+ render_three_builtins_novel(dataset, trainer, args.fps, novel_dir, args.skip_existing)
+ render_leftshift_novel(dataset, trainer, args.fps, novel_dir, args.skip_existing)
+
+ # Full-view depth video + raw depth
+ export_full_depth(cfg, dataset, trainer, log_dir, step, args.raw_depth_dir, args.skip_existing)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/third_party/drivestudio/tools/collect_nus_benchmark.py b/worldbench/third_party/drivestudio/tools/collect_nus_benchmark.py
new file mode 100644
index 0000000..c76ca6b
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/collect_nus_benchmark.py
@@ -0,0 +1,172 @@
+#!/usr/bin/env python3
+"""
+Collect nuScenes novel-view benchmark videos into a unified delivery tree.
+
+Source layout assumption (per run):
+ work_dirs/
+ drivestudio-nus-/
+ /
+ noval_30000_benchmark/ # user-customized output folder
+ front_center_interp.mp4
+ s_curve.mp4
+ lateral_offset.mp4
+ lateral_offset_left.mp4
+
+Target layout:
+ nus_novel_view_benchmark/
+ /
+ /
+ .mp4
+
+Steps:
+ 1) Enumerate methods matching --pattern under --root (default: work_dirs/drivestudio-nus-*)
+ 2) For each method, find the set of clips (run subdirs) that contain all required view mp4s
+ specifically under: /videos//
+ 3) Compute intersection across methods
+ 4) For the common clips, copy/symlink/Hardlink each view mp4 to target tree
+
+Safety:
+ - Never deletes anything; only creates target directories as needed
+ - Overwrites existing target files if --overwrite
+
+Usage:
+ python tools/collect_nus_benchmark.py \
+ --root work_dirs \
+ --pattern "drivestudio-nus-*" \
+ --src-folder "noval_30000_benchmark" \
+ --dest nus_novel_view_benchmark \
+ --views front_center_interp,s_curve,lateral_offset,lateral_offset_left \
+ --mode copy --overwrite
+"""
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+from typing import Dict, Iterable, List, Optional, Sequence, Set, Tuple
+import shutil
+
+
+DEFAULT_VIEWS = [
+ "front_center_interp",
+ "s_curve",
+ "lateral_offset",
+ "lateral_offset_left",
+]
+
+
+def parse_args() -> argparse.Namespace:
+ ap = argparse.ArgumentParser("Collect nuScenes novel-view benchmark videos")
+ ap.add_argument("--root", type=str, default="work_dirs", help="Root containing method dirs")
+ ap.add_argument("--pattern", type=str, default="drivestudio-nus-*", help="Method dir glob pattern")
+ ap.add_argument("--src-folder", type=str, default="novel_30000_benchmark", help="Subfolder under run_dir/videos/ containing view mp4s")
+ ap.add_argument("--dest", type=str, default="data/nus_novel_view_benchmark", help="Destination root for collected videos")
+ ap.add_argument("--views", type=str, default=",".join(DEFAULT_VIEWS), help="Comma-separated view names to collect")
+ ap.add_argument("--mode", type=str, choices=["copy", "symlink", "hardlink"], default="copy", help="File materialization mode")
+ ap.add_argument("--overwrite", action="store_true", help="Overwrite existing files at destination")
+ ap.add_argument("--strip-prefix", type=str, default="drivestudio-nus-", help="Prefix to strip from method dir name when naming method in dest")
+ return ap.parse_args()
+
+
+def list_method_dirs(root: Path, pattern: str) -> List[Path]:
+ return sorted(root.glob(pattern))
+
+
+def find_src_view_file(run_dir: Path, src_folder: str, view: str) -> Optional[Path]:
+ """Return path to a view mp4 strictly under /videos//.
+
+ Success is defined only if the file exists in this directory; no other fallbacks are considered.
+ """
+ base = run_dir / "videos" / src_folder
+ if not base.exists():
+ return None
+ # prefer explicit cam0 if present, else plain view
+ for name in [f"{view}_cam0.mp4", f"{view}.mp4"]:
+ p = base / name
+ if p.is_file():
+ return p
+ return None
+
+
+def collect_runs_with_all_views(method_dir: Path, src_folder: str, views: Sequence[str]) -> Set[str]:
+ """Return set of clip names (run subdir names) that contain all requested view files."""
+ ok: Set[str] = set()
+ for run in sorted(d for d in method_dir.iterdir() if d.is_dir()):
+ has_all = True
+ for view in views:
+ if find_src_view_file(run, src_folder, view) is None:
+ has_all = False
+ break
+ if has_all:
+ ok.add(run.name)
+ return ok
+
+
+def materialize(src: Path, dst: Path, mode: str, overwrite: bool) -> None:
+ dst.parent.mkdir(parents=True, exist_ok=True)
+ if dst.exists():
+ if not overwrite:
+ return
+ # Overwrite by replacing file atomically
+ dst.unlink()
+ if mode == "copy":
+ shutil.copy2(src, dst)
+ elif mode == "symlink":
+ os.symlink(os.path.abspath(src), dst)
+ elif mode == "hardlink":
+ os.link(src, dst)
+ else:
+ raise ValueError(f"Unknown mode: {mode}")
+
+
+def main() -> None:
+ args = parse_args()
+ root = Path(args.root)
+ dest_root = Path(args.dest)
+ views = [v for v in (args.views.split(",") if args.views else []) if v]
+
+ method_dirs = list_method_dirs(root, args.pattern)
+ if not method_dirs:
+ print(f"No method dirs under {root} matching pattern {args.pattern}")
+ return
+
+ # Compute intersection of clips across methods (and that contain all views per method)
+ per_method_clips: Dict[str, Set[str]] = {}
+ for mdir in method_dirs:
+ clips = collect_runs_with_all_views(mdir, args.src_folder, views)
+ method_name = mdir.name
+ per_method_clips[method_name] = clips
+ print(f"Method {method_name}: {len(clips)} eligible clips")
+
+ if not per_method_clips:
+ print("No eligible clips found")
+ return
+
+ # Intersection
+ common: Optional[Set[str]] = None
+ for clips in per_method_clips.values():
+ common = clips if common is None else (common & clips)
+ common = common or set()
+ print(f"Common clips across methods: {len(common)}")
+
+ # Materialize into target tree
+ for mdir in method_dirs:
+ method_label = mdir.name
+ if args.strip_prefix and method_label.startswith(args.strip_prefix):
+ method_label = method_label[len(args.strip_prefix):]
+ for clip in sorted(common):
+ run_dir = mdir / clip
+ for view in views:
+ src = find_src_view_file(run_dir, args.src_folder, view)
+ if src is None:
+ # Should not happen given we filtered above; skip defensively
+ print(f"[WARN] Missing view {view} for {mdir.name}/{clip}")
+ continue
+ dst = dest_root / method_label / view / f"{clip}.mp4"
+ materialize(src, dst, args.mode, args.overwrite)
+
+ print(f"Done. Collected into: {dest_root}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/third_party/drivestudio/tools/compose_depth_grid.py b/worldbench/third_party/drivestudio/tools/compose_depth_grid.py
new file mode 100644
index 0000000..773a1dd
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/compose_depth_grid.py
@@ -0,0 +1,534 @@
+#!/usr/bin/env python3
+"""
+Compose a 2x2 grid depth visualization video across multiple methods for a given clip.
+
+Inputs (per method/run):
+ work_dirs/
+ /
+ /
+ videos/
+ full_set__depths.mp4
+
+Output:
+ A single MP4 that tiles the per-method depth videos in a 2x2 layout (for 4 methods).
+ If the number of methods differs, the script lays them out in a near-square grid.
+
+This follows the style of tools/compose_novel_view_grid.py, but specializes to depth videos.
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import shutil
+import subprocess
+from pathlib import Path
+from typing import List, Optional, Sequence, Tuple, Dict
+import csv
+
+
+def parse_args() -> argparse.Namespace:
+ p = argparse.ArgumentParser("Compose depth grid video for a given nuScenes clip")
+ p.add_argument("clip_id", type=str, help="Clip identifier (e.g., 084)")
+ p.add_argument("--root", type=Path, default=Path("work_dirs"), help="Root directory containing per-model runs")
+ p.add_argument("--models", type=str, nargs="*", help="Explicit list of model directory names; default: all that contain the clip")
+ p.add_argument("--step", type=int, default=30000, help="Training step used in file naming (default: 30000)")
+ p.add_argument("--output", type=Path, default=None, help="Output MP4 path (default: work_dirs/grids/clip__depth_grid.mp4)")
+ p.add_argument("--ffmpeg-bin", type=str, default="ffmpeg", help="FFmpeg executable (default: ffmpeg)")
+ p.add_argument("--codec", type=str, default=None, help="Preferred FFmpeg encoder (e.g., libx264)")
+ p.add_argument("--bitrate", type=str, default="10M", help="Bitrate for bitrate-based encoders (default: 10M)")
+ p.add_argument("--font", type=Path, default=None, help="Optional TrueType/OpenType font for labels")
+ p.add_argument("--fontsize", type=int, default=28, help="Label font size (default: 28)")
+ p.add_argument("--gap", type=int, default=8, help="Gap (pixels) between tiles (default: 8)")
+ p.add_argument("--tile-height", type=int, default=0, help="Optional per-tile fixed height (px); 0 disables scaling.")
+ p.add_argument("--tile-width", type=int, default=0, help="Optional per-tile fixed width (px); when height or width >0, tiles are resized to this WxH.")
+ p.add_argument("--fit", type=str, choices=["contain", "cover"], default="contain", help="Resize mode when both tile width/height are set: contain (scale+pad) or cover (scale+crop). Default: contain.")
+ # Optional manual crop to remove baked-in letterbox bars before resizing (pixels on source frames)
+ p.add_argument("--crop-top", type=int, default=0, help="Crop N pixels from top before resizing (default: 0)")
+ p.add_argument("--crop-bottom", type=int, default=0, help="Crop N pixels from bottom before resizing (default: 0)")
+ p.add_argument("--crop-left", type=int, default=0, help="Crop N pixels from left before resizing (default: 0)")
+ p.add_argument("--crop-right", type=int, default=0, help="Crop N pixels from right before resizing (default: 0)")
+ p.add_argument("--verbose", action="store_true", help="Print ffmpeg command")
+ p.add_argument("--fps", type=int, default=0, help="Target output FPS override; 0 = keep source FPS (default)")
+ p.add_argument("--frames", type=int, default=0, help="Trim output to this many frames after stacking (0=disable)")
+ # Metric overlay (AbsRel for Depth by default)
+ p.add_argument("--depth-per-seq", type=Path, nargs="*", default=None, help="Per-sequence depth CSV(s) (per_seq.csv) for AbsRel overlay. If omitted, tries depth_tables_3/per_seq.csv then depth_tables_2/per_seq.csv.")
+ p.add_argument("--metric-label", type=str, default="AbsRel", help="Display label for the metric (default: AbsRel)")
+ p.add_argument("--no-metric", action="store_true", help="Disable metric overlay on labels")
+ p.add_argument("--debug", action="store_true", help="Print detailed alignment info (models, inputs, probed FPS/frames, metric values)")
+ return p.parse_args()
+
+
+def resolve_ffmpeg(binary: str) -> str:
+ for cand in [binary, os.environ.get("FFMPEG_BIN"), "ffmpeg"]:
+ if cand and shutil.which(cand):
+ return cand
+ raise SystemExit("FFmpeg not found; install ffmpeg or set --ffmpeg-bin/FFMPEG_BIN")
+
+
+def resolve_ffprobe(ffmpeg_bin: str) -> str:
+ p = Path(ffmpeg_bin)
+ cand = p.parent / "ffprobe" if p.parent else None
+ if cand and cand.exists():
+ return str(cand)
+ probed = shutil.which("ffprobe")
+ if probed:
+ return probed
+ raise SystemExit("ffprobe not found; install ffmpeg (includes ffprobe) or add to PATH")
+
+
+def probe_fps(ffprobe_bin: str, video: Path) -> Optional[float]:
+ try:
+ out = subprocess.check_output([
+ ffprobe_bin,
+ "-v", "error",
+ "-select_streams", "v:0",
+ "-show_entries", "stream=avg_frame_rate",
+ "-of", "default=nw=1:nk=1",
+ str(video),
+ ], text=True)
+ fr = out.strip()
+ if not fr:
+ return None
+ if "/" in fr:
+ num, den = fr.split("/", 1)
+ num = float(num); den = float(den)
+ return num / den if den else None
+ return float(fr)
+ except Exception:
+ return None
+
+
+def probe_nb_frames(ffprobe_bin: str, video: Path) -> Optional[int]:
+ try:
+ out = subprocess.check_output([
+ ffprobe_bin,
+ "-v", "error",
+ "-select_streams", "v:0",
+ "-count_frames",
+ "-show_entries", "stream=nb_read_frames",
+ "-of", "csv=p=0",
+ str(video),
+ ], text=True)
+ s = out.strip()
+ return int(s) if s.isdigit() else None
+ except Exception:
+ return None
+
+
+CANDIDATE_FONT_NAMES = [
+ "DejaVuSans.ttf",
+ "DejaVuSans-Regular.ttf",
+ "LiberationSans-Regular.ttf",
+ "NotoSans-Regular.ttf",
+]
+
+
+def resolve_font_path(user_font: Optional[Path]) -> Optional[Path]:
+ """Try best-effort to locate a usable TrueType/OpenType font for drawtext.
+
+ Returns a Path if found; otherwise None (caller may choose to disable labels).
+ """
+ if user_font:
+ if user_font.is_file():
+ return user_font
+ raise SystemExit(f"Specified font file '{user_font}' does not exist.")
+
+ env_font = os.environ.get("FFMPEG_FONT")
+ if env_font and Path(env_font).is_file():
+ return Path(env_font)
+
+ candidates: List[Path] = []
+ system_dirs = [
+ Path("/usr/share/fonts"),
+ Path("/usr/local/share/fonts"),
+ Path.home() / ".fonts",
+ ]
+ for base in system_dirs:
+ for name in CANDIDATE_FONT_NAMES:
+ candidates.append(base / "truetype" / name)
+ candidates.append(base / "truetype" / "dejavu" / name)
+ candidates.append(base / name)
+
+ conda_prefix = os.environ.get("CONDA_PREFIX")
+ if conda_prefix:
+ cp = Path(conda_prefix)
+ for name in CANDIDATE_FONT_NAMES:
+ candidates.append(cp / "lib" / "python3.10" / "site-packages" / name)
+ candidates.append(cp / "lib" / "python3.10" / "site-packages" / "matplotlib" / "mpl-data" / "fonts" / "ttf" / name)
+ pkgs_dir = cp / "pkgs"
+ if pkgs_dir.is_dir():
+ for name in CANDIDATE_FONT_NAMES:
+ candidates.extend(pkgs_dir.glob(f"pillow*/info/test/Tests/fonts/{name}"))
+
+ for cand in candidates:
+ if cand.is_file() and "Symbols" not in cand.name and "Color" not in cand.name:
+ return cand
+
+ return None
+
+
+def discover_models(root: Path, clip_id: str, explicit: Optional[Sequence[str]]) -> List[str]:
+ if explicit:
+ return list(explicit)
+ models = []
+ for d in sorted(root.iterdir()):
+ if (d / clip_id).is_dir():
+ models.append(d.name)
+ if not models:
+ raise SystemExit(f"No models containing clip '{clip_id}' found under {root}")
+ models.sort()
+ if "drivestudio-nus-gt" in models:
+ models.remove("drivestudio-nus-gt")
+ models.insert(0, "drivestudio-nus-gt")
+ return models
+
+
+def method_label(name: str) -> str:
+ return name.split("drivestudio-nus-", 1)[-1].replace("_", " ").replace("-", " ")
+
+
+def find_depth_video(root: Path, model: str, clip_id: str, step: int) -> Path:
+ path = root / model / clip_id / "videos" / f"full_set_{step}_depths.mp4"
+ if not path.is_file():
+ raise SystemExit(f"Missing depth video: {path}")
+ return path
+
+
+def grid_shape(n: int) -> Tuple[int, int]:
+ # Prefer fewer rows to reduce perceived row spacing. Use r=floor(sqrt(n)), c=ceil(n/r).
+ if n <= 0:
+ return 0, 0
+ r = int(n ** 0.5)
+ if r < 1:
+ r = 1
+ from math import ceil
+ c = int(ceil(n / r))
+ return r, c
+
+
+def drawtext_filter(font: Optional[Path], fontsize: int, text: str) -> str:
+ if not font:
+ return "" # no labeling when no font is available
+ return (
+ f"drawtext=fontfile={font.as_posix()}:text='{text}':fontcolor=white:fontsize={fontsize}"
+ ":box=1:boxcolor=black@0.6:boxborderw=10:x=20:y=20"
+ )
+
+
+def build_filter_complex(
+ inputs: List[Path],
+ labels: List[str],
+ rows: int,
+ cols: int,
+ font: Optional[Path],
+ fontsize: int,
+ gap: int,
+ tile_height: int,
+ tile_width: int,
+ fps: int = 30,
+ frames: int = 0,
+ fit: str = "contain",
+ crop: Tuple[int, int, int, int] = (0, 0, 0, 0),
+ placeholder_height: int = 0,
+ row_target_widths: Optional[List[int]] = None,
+) -> Tuple[str, str]:
+ if len(inputs) != len(labels):
+ raise ValueError("inputs and labels length mismatch")
+ n = len(inputs)
+ cmds: List[str] = []
+ out_labels: List[str] = []
+ th = tile_height if tile_height and tile_height > 0 else 0
+ # Height-only mode (tw==0) will use scale=-2:th (no internal vertical padding). If tw>0, use contain/cover path.
+ tw = tile_width if tile_width and tile_width > 0 else 0
+
+ for i in range(n):
+ label_text = labels[i].replace(":", "\\:").replace("'", "\\'")
+ filt = drawtext_filter(font, fontsize, label_text)
+ out_name = f"v{i}"
+ chain = f"[{i}:v]"
+ if filt:
+ chain += filt
+ else:
+ chain += "null"
+ ct, cb, cl, cr = crop
+ if any(v > 0 for v in (ct, cb, cl, cr)):
+ chain += f",crop=iw-{cl}-{cr}:ih-{ct}-{cb}:{cl}:{ct}"
+ if th:
+ if fit == "cover":
+ chain += f",scale={tw if tw else -2}:{th}:force_original_aspect_ratio=increase"
+ if tw:
+ chain += f",crop={tw}:{th}:(iw-{tw})/2:(ih-{th})/2"
+ else:
+ chain += f",scale={tw if tw else -2}:{th}:force_original_aspect_ratio=decrease"
+ if tw:
+ chain += f",pad={tw}:{th}:(ow-iw)/2:(oh-ih)/2:color=black"
+ cmds.append(f"{chain}[{out_name}]")
+ out_labels.append(out_name)
+
+ # pad with placeholder frames if necessary
+ pads = 0
+ total = rows * cols
+ while len(out_labels) < total:
+ pad_label = f"pad{pads}"
+ if th:
+ size_str = f"size={(tw if tw else 64)}x{th}"
+ else:
+ ph = placeholder_height if placeholder_height and placeholder_height > 0 else 16
+ size_str = f"size=16x{ph}"
+ if fps and fps > 0:
+ cmds.append(f"color=c=black:{size_str}:rate={fps},format=rgb24[{pad_label}]")
+ else:
+ cmds.append(f"color=c=black:{size_str},format=rgb24[{pad_label}]")
+ out_labels.append(pad_label)
+ pads += 1
+
+ # stack per row with hstack, then vstack rows
+ row_outs: List[str] = []
+ for r in range(rows):
+ seg = out_labels[r * cols : (r + 1) * cols]
+ # pad right gap for all but last column to create horizontal spacing
+ padded_cols: List[str] = []
+ for c_idx, s in enumerate(seg):
+ cur = s
+ if gap > 0 and c_idx < cols - 1:
+ pad_label = f"{s}_pad"
+ cmds.append(f"[{s}]pad=iw+{gap}:ih:0:0:color=black[{pad_label}]")
+ cur = pad_label
+ padded_cols.append(cur)
+ tmp_row = f"row{r}_tmp"
+ if len(padded_cols) == 1:
+ cmds.append(f"[{padded_cols[0]}]null[{tmp_row}]")
+ else:
+ inputs_concat = "".join(f"[{s}]" for s in padded_cols)
+ cmds.append(f"{inputs_concat}hstack=inputs={len(padded_cols)}:shortest=1[{tmp_row}]")
+ # pad to target row width if in height-only mode (tw==0) and provided
+ row_label = f"row{r}"
+ if tw == 0 and row_target_widths is not None and r < len(row_target_widths) and row_target_widths[r] > 0:
+ target_w = row_target_widths[r]
+ cmds.append(f"[{tmp_row}]pad={target_w}:ih:0:0:color=black[{row_label}_padw]")
+ tmp_row = f"{row_label}_padw"
+ # pad bottom gap for all but last row to create vertical spacing
+ if gap > 0 and r < rows - 1:
+ cmds.append(f"[{tmp_row}]pad=iw:ih+{gap}:0:0:color=black[{row_label}]")
+ else:
+ cmds.append(f"[{tmp_row}]null[{row_label}]")
+ row_outs.append(row_label)
+
+ if len(row_outs) == 1:
+ final = row_outs[0]
+ else:
+ final = "grid"
+ inputs_concat = "".join(f"[{s}]" for s in row_outs)
+ cmds.append(f"{inputs_concat}vstack=inputs={len(row_outs)}:shortest=1[{final}]")
+
+ tail_parts = []
+ if fps and fps > 0:
+ tail_parts.append(f"fps={fps}")
+ if frames and frames > 0:
+ tail_parts.append(f"trim=end_frame={frames}")
+ tail_parts.append("setsar=1")
+ tail = ",".join(tail_parts)
+ cmds.append(f"[{final}]{tail}[outv]")
+ return ";".join(cmds), "outv"
+
+
+def build_ffmpeg_cmd(binary: str, inputs: List[Path], filter_complex: str, final_label: str, output: Path, codec: Optional[str], bitrate: str) -> List[str]:
+ cmd: List[str] = [binary, "-y"]
+ for p in inputs:
+ cmd.extend(["-i", str(p)])
+ cmd.extend(["-filter_complex", filter_complex, "-map", f"[{final_label}]", "-pix_fmt", "yuv420p", "-vsync", "2", "-shortest", "-an"])
+ if codec:
+ cmd.extend(["-c:v", codec])
+ if codec == "libx264":
+ cmd.extend(["-preset", "medium", "-crf", "18"])
+ elif codec in {"libopenh264", "mpeg4"}:
+ cmd.extend(["-b:v", bitrate])
+ cmd.append(str(output))
+ return cmd
+
+
+def main() -> None:
+ args = parse_args()
+ ffmpeg = resolve_ffmpeg(args.ffmpeg_bin)
+ models = discover_models(args.root, args.clip_id, args.models)
+
+ inputs: List[Path] = []
+ labels: List[str] = []
+ # Prepare metric overlay map: method -> value for this clip
+ metric_map: Dict[str, Optional[float]] = {m: None for m in models}
+ if not args.no_metric:
+ csv_paths: List[Path] = []
+ if args.depth_per_seq:
+ csv_paths = [p for p in args.depth_per_seq if p and p.is_file()]
+ else:
+ for cand in [Path("depth_tables_3/per_seq.csv"), Path("depth_tables_2/per_seq.csv"), Path("depth_tables/per_seq.csv")]:
+ if cand.is_file():
+ csv_paths.append(cand)
+ # Load rows and build lookup
+ lookup: Dict[Tuple[str, str], float] = {}
+ for path in csv_paths:
+ try:
+ with path.open("r", newline="") as f:
+ reader = csv.DictReader(f)
+ if "Method/Seq" in reader.fieldnames and "AbsRel" in reader.fieldnames:
+ for row in reader:
+ label = row.get("Method/Seq", "").strip()
+ absrel = row.get("AbsRel")
+ if "/" in label and absrel:
+ method, seq = label.split("/", 1)
+ try:
+ v = float(absrel)
+ except ValueError:
+ continue
+ lookup[(method.strip(), seq.strip())] = v
+ except Exception:
+ continue
+ for m in models:
+ meth = method_label(m).replace(" ", "") # our labels may have spaces removed elsewhere
+ key1 = (meth, args.clip_id)
+ key2 = (meth.lower(), args.clip_id)
+ if key1 in lookup:
+ metric_map[m] = lookup[key1]
+ elif key2 in lookup:
+ metric_map[m] = lookup[key2]
+
+ for m in models:
+ vid = find_depth_video(args.root, m, args.clip_id, args.step)
+ inputs.append(vid)
+ base = method_label(m)
+ v = metric_map.get(m)
+ if v is not None:
+ labels.append(f"{base} | {args.metric_label} {v:.3f}")
+ else:
+ labels.append(base)
+ rows, cols = grid_shape(len(inputs))
+
+ out_path = args.output
+ if out_path is None:
+ out_dir = args.root / "grids"
+ out_dir.mkdir(parents=True, exist_ok=True)
+ out_path = out_dir / f"clip_{args.clip_id}_depth_grid.mp4"
+ else:
+ out_path.parent.mkdir(parents=True, exist_ok=True)
+
+ font_path = resolve_font_path(args.font)
+ if args.font and not font_path:
+ raise SystemExit(f"Font file not found: {args.font}")
+
+ # Determine placeholder height when not resizing tiles
+ ph = 0
+ if (not args.tile_height or args.tile_height == 0) and inputs:
+ try:
+ ffprobe_bin = resolve_ffprobe(ffmpeg)
+ info = subprocess.check_output([
+ ffprobe_bin,
+ "-v","error",
+ "-select_streams","v:0",
+ "-show_entries","stream=height",
+ "-of","csv=p=0",
+ str(inputs[0])
+ ], text=True).strip()
+ ph = int(info) if info.isdigit() else 0
+ except Exception:
+ ph = 0
+ # Height-only mode: compute per-row target widths so rows align and avoid vertical bars
+ tile_width_eff = args.tile_width or 0
+ row_targets: Optional[List[int]] = None
+ if (args.tile_height and args.tile_height > 0) and (not args.tile_width or args.tile_width == 0) and inputs:
+ try:
+ ffprobe_bin = resolve_ffprobe(ffmpeg)
+ except Exception:
+ ffprobe_bin = None
+ wh_list: List[Tuple[int,int]] = []
+ for pth in inputs:
+ w = h = 0
+ if ffprobe_bin:
+ try:
+ s = subprocess.check_output([
+ ffprobe_bin,
+ "-v","error",
+ "-select_streams","v:0",
+ "-show_entries","stream=width,height",
+ "-of","csv=p=0",
+ str(pth)
+ ], text=True).strip()
+ parts = s.split(',')
+ if len(parts) >= 2:
+ w = int(parts[0]); h = int(parts[1])
+ except Exception:
+ w = h = 0
+ wh_list.append((w, h))
+ sw: List[int] = []
+ th = args.tile_height
+ for (w, h) in wh_list:
+ if w > 0 and h > 0:
+ ww = int(round(w * (th / float(h))))
+ if ww % 2 != 0:
+ ww += 1
+ else:
+ ww = 1280
+ sw.append(ww)
+ row_widths: List[int] = []
+ for r in range(rows):
+ seg = sw[r * cols : (r + 1) * cols]
+ while len(seg) < cols:
+ seg.append(64)
+ row_widths.append(sum(seg) + (args.gap * (cols - 1 if args.gap > 0 else 0)))
+ max_w = max(row_widths) if row_widths else 0
+ row_targets = [max_w for _ in range(rows)]
+
+ filter_complex, final_label = build_filter_complex(
+ inputs, labels, rows, cols, font_path, args.fontsize, args.gap,
+ args.tile_height, tile_width_eff, fps=args.fps, frames=args.frames, fit=args.fit,
+ crop=(args.crop_top, args.crop_bottom, args.crop_left, args.crop_right), placeholder_height=ph,
+ row_target_widths=row_targets)
+ # Debug alignment info
+ if args.debug:
+ print("[DEBUG] clip:", args.clip_id, "step:", args.step)
+ print("[DEBUG] grid:", f"rows={rows} cols={cols} gap={args.gap} tile=({('auto' if tile_width_eff==0 else tile_width_eff)}x{args.tile_height or 'auto'})")
+ print("[DEBUG] models order:", ", ".join(models))
+ try:
+ ffprobe_bin = resolve_ffprobe(ffmpeg)
+ except Exception:
+ ffprobe_bin = None
+ for m, path, lab in zip(models, inputs, labels):
+ fps_val = probe_fps(ffprobe_bin, path) if ffprobe_bin else None
+ frames_val = probe_nb_frames(ffprobe_bin, path) if ffprobe_bin else None
+ wh = None
+ sar = dar = None
+ if ffprobe_bin:
+ try:
+ whs = subprocess.check_output([
+ ffprobe_bin,
+ "-v","error",
+ "-select_streams","v:0",
+ "-show_entries","stream=width,height,sample_aspect_ratio,display_aspect_ratio",
+ "-of","csv=p=0",
+ str(path)
+ ], text=True).strip()
+ parts = whs.split('\n')[0].split(',')
+ if len(parts) >= 2:
+ wh = f"{parts[0]}x{parts[1]}"
+ if len(parts) >= 3:
+ sar = parts[2]
+ if len(parts) >= 4:
+ dar = parts[3]
+ except Exception:
+ pass
+ print(f"[DEBUG] {m}: video={path} size={wh} sar={sar} dar={dar} fps={fps_val} frames={frames_val} label=\"{lab}\"")
+ # Metric source for Depth
+ print("[DEBUG] metric source candidates:", end=" ")
+ if args.depth_per_seq:
+ print(", ".join(str(p) for p in args.depth_per_seq))
+ else:
+ print("depth_tables_3/per_seq.csv, depth_tables_2/per_seq.csv (fallback)")
+ cmd = build_ffmpeg_cmd(ffmpeg, inputs, filter_complex, final_label, out_path, args.codec, args.bitrate)
+ if args.verbose:
+ print("FFmpeg command:")
+ print(" ".join(cmd))
+ subprocess.run(cmd, check=True)
+ print(f"Saved depth grid video to {out_path}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/third_party/drivestudio/tools/compose_novel_view_grid.py b/worldbench/third_party/drivestudio/tools/compose_novel_view_grid.py
new file mode 100644
index 0000000..977ecd6
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/compose_novel_view_grid.py
@@ -0,0 +1,429 @@
+#!/usr/bin/env python3
+"""
+Compose a grid video of novel-view renderings across multiple models.
+Each row corresponds to a novel trajectory (e.g., s_curve) and
+each column corresponds to a model (including the GT reference).
+"""
+from __future__ import annotations
+
+import argparse
+import shutil
+import subprocess
+from pathlib import Path
+import os
+from typing import Dict, List, Optional, Sequence, Tuple
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(
+ description="Create a model-by-view grid video for a given clip."
+ )
+ parser.add_argument("clip_id", type=str, help="Clip identifier (e.g., 084).")
+ parser.add_argument(
+ "--root",
+ type=Path,
+ default=Path("work_dirs"),
+ help="Root directory containing per-model results (default: work_dirs).",
+ )
+ parser.add_argument(
+ "--models",
+ type=str,
+ nargs="*",
+ help="Explicit list of model directory names. Defaults to all models that contain the clip.",
+ )
+ parser.add_argument(
+ "--novel-dir",
+ type=str,
+ default=None,
+ help="Name of the novel-view directory (e.g., novel_30000). Defaults to the latest 'novel_*' folder found per model.",
+ )
+ parser.add_argument(
+ "--output",
+ type=Path,
+ default=None,
+ help="Output video path. Defaults to work_dirs/grids/clip__novel_grid.mp4.",
+ )
+ parser.add_argument(
+ "--font",
+ type=Path,
+ default=None,
+ help="Font file for drawtext. If omitted, the script will attempt to auto-detect a suitable TrueType/OpenType font.",
+ )
+ parser.add_argument(
+ "--fontsize", type=int, default=28, help="Font size for labels (default: 28)."
+ )
+ parser.add_argument(
+ "--codec",
+ type=str,
+ default=None,
+ help="Preferred FFmpeg video encoder (e.g., libx264, libopenh264). Defaults to auto-detect.",
+ )
+ parser.add_argument(
+ "--bitrate",
+ type=str,
+ default="10M",
+ help="Fallback bitrate when CRF is unavailable (default: 10M).",
+ )
+ parser.add_argument(
+ "--ffmpeg-bin",
+ type=str,
+ default="ffmpeg",
+ help="FFmpeg executable to use (default: ffmpeg).",
+ )
+ parser.add_argument(
+ "--verbose", action="store_true", help="Print detailed ffmpeg command."
+ )
+ return parser.parse_args()
+
+
+CANDIDATE_FONT_NAMES = [
+ "DejaVuSans.ttf",
+ "DejaVuSans-Regular.ttf",
+ "LiberationSans-Regular.ttf",
+ "NotoSans-Regular.ttf",
+]
+
+MODEL_NAME_MAP = {
+ "drivestudio-nus-gt": "GT",
+ "drivestudio-nus-dreamforge": "dreamforge",
+ "drivestudio-nus-drivedreamer2": "drivedreamer2",
+ "drivestudio-nus-magicdrive": "magicdrive",
+}
+
+
+def resolve_ffmpeg_binary(preferred: str) -> str:
+ candidates: List[str] = []
+ if preferred:
+ candidates.append(preferred)
+ env_bin = os.environ.get("FFMPEG_BIN")
+ if env_bin:
+ candidates.append(env_bin)
+ conda_prefix = os.environ.get("CONDA_PREFIX")
+ if conda_prefix:
+ candidates.append(str(Path(conda_prefix) / "bin" / "ffmpeg"))
+ candidates.append("ffmpeg")
+
+ for cand in candidates:
+ if shutil.which(cand):
+ return cand
+ raise SystemExit(
+ "Could not find an FFmpeg executable. Consider installing ffmpeg or set FFMPEG_BIN/path via --ffmpeg-bin."
+ )
+
+
+def detect_available_codecs(binary: str) -> List[str]:
+ try:
+ result = subprocess.run(
+ [binary, "-hide_banner", "-encoders"],
+ stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE,
+ text=True,
+ check=True,
+ )
+ except subprocess.CalledProcessError:
+ return []
+ encoders = []
+ for line in result.stdout.splitlines():
+ if not line or line.startswith(" "):
+ continue
+ encoders.append(line.strip())
+ # actual encoder names are at fixed position (last token)
+ available = []
+ for line in result.stdout.splitlines():
+ parts = line.split()
+ if not parts:
+ continue
+ name = parts[-1]
+ if name.startswith("lib") or name in {"h264", "mpeg4"}:
+ available.append(parts[-1])
+ return available
+
+
+def resolve_font_path(user_font: Optional[Path]) -> Path:
+ if user_font:
+ if user_font.is_file():
+ return user_font
+ raise SystemExit(f"Specified font file '{user_font}' does not exist.")
+
+ env_font = os.environ.get("FFMPEG_FONT")
+ if env_font and Path(env_font).is_file():
+ return Path(env_font)
+
+ candidates: List[Path] = []
+ system_dirs = [
+ Path("/usr/share/fonts"),
+ Path("/usr/local/share/fonts"),
+ Path.home() / ".fonts",
+ ]
+ for base in system_dirs:
+ for name in CANDIDATE_FONT_NAMES:
+ candidates.append(base / "truetype" / name)
+ candidates.append(base / "truetype" / "dejavu" / name)
+ candidates.append(base / name)
+
+ conda_prefix = os.environ.get("CONDA_PREFIX")
+ if conda_prefix:
+ cp = Path(conda_prefix)
+ for name in CANDIDATE_FONT_NAMES:
+ candidates.append(cp / "lib" / "python3.10" / "site-packages" / name)
+ candidates.append(cp / "lib" / "python3.10" / "site-packages" / "matplotlib" / "mpl-data" / "fonts" / "ttf" / name)
+ pkgs_dir = cp / "pkgs"
+ if pkgs_dir.is_dir():
+ for name in CANDIDATE_FONT_NAMES:
+ candidates.extend(pkgs_dir.glob(f"pillow*/info/test/Tests/fonts/{name}"))
+
+ for cand in candidates:
+ if cand.is_file() and "Symbols" not in cand.name and "Color" not in cand.name:
+ return cand
+
+ raise SystemExit(
+ "No suitable font file found. Please install a TrueType/OpenType font or provide --font."
+ )
+
+
+def choose_codec(binary: str, user_codec: Optional[str]) -> str:
+ if user_codec:
+ return user_codec
+ preferred_order = ["libx264", "libopenh264", "libx265", "mpeg4"]
+ available = detect_available_codecs(binary)
+ for codec in preferred_order:
+ if codec in available:
+ return codec
+ # fallback to ffmpeg default (no codec flag)
+ return ""
+
+
+def discover_models(root: Path, clip_id: str) -> List[str]:
+ models = []
+ for model_dir in sorted(root.iterdir()):
+ if not model_dir.is_dir():
+ continue
+ clip_dir = model_dir / clip_id
+ if clip_dir.is_dir():
+ models.append(model_dir.name)
+ if not models:
+ raise SystemExit(f"No model directories containing clip '{clip_id}' were found under {root}.")
+ # Prefer GT first if present
+ models.sort()
+ if "drivestudio-nus-gt" in models:
+ models.remove("drivestudio-nus-gt")
+ models.insert(0, "drivestudio-nus-gt")
+ return models
+
+
+def select_novel_dir(clip_dir: Path, explicit: str | None) -> Path:
+ videos_dir = clip_dir / "videos"
+ if not videos_dir.is_dir():
+ raise SystemExit(f"No videos directory found at {videos_dir}")
+ if explicit:
+ candidate = videos_dir / explicit
+ if not candidate.is_dir():
+ raise SystemExit(f"Specified novel directory '{explicit}' not found at {candidate}")
+ return candidate
+ candidates = sorted(p for p in videos_dir.iterdir() if p.is_dir() and p.name.startswith("novel_"))
+ if not candidates:
+ raise SystemExit(f"No novel_* directory found inside {videos_dir}")
+ return candidates[-1]
+
+
+def load_novel_videos(novel_dir: Path) -> Dict[str, Path]:
+ videos = {}
+ for vid in sorted(novel_dir.glob("*.mp4")):
+ if vid.is_file():
+ videos[vid.stem] = vid
+ if not videos:
+ raise SystemExit(f"No MP4 files found in {novel_dir}")
+ return videos
+
+
+def prettify_model_name(name: str) -> str:
+ label = MODEL_NAME_MAP.get(name, name)
+ if label == name and name.startswith("drivestudio-nus-"):
+ label = name.split("drivestudio-nus-", 1)[1]
+ return label.replace("_", " ").replace("-", " ")
+
+
+def prettify_view_name(name: str) -> str:
+ return name.replace("_", " ").replace("-", " ")
+
+
+def sanitize_text(text: str) -> str:
+ return (
+ text.replace("\\", r"\\")
+ .replace("'", r"\'")
+ .replace(":", r"\:")
+ )
+
+
+def build_filter_complex(
+ inputs: List[Path],
+ models: Sequence[str],
+ views: Sequence[str],
+ font: Path | None,
+ fontsize: int,
+) -> Tuple[str, str]:
+ num_models = len(models)
+ filter_cmds: List[str] = []
+ stream_labels: List[str] = []
+
+ if font is None:
+ raise SystemExit(
+ "No font file available for drawtext. Please provide --font pointing to a .ttf or .otf file."
+ )
+ font_opt = f"fontfile={font.as_posix()}:"
+
+ for idx, (view, model) in enumerate([(v, m) for v in views for m in models]):
+ col = idx % num_models
+ row = idx // num_models
+ in_label = f"{idx}:v"
+ current = f"{in_label}"
+ filters: List[str] = []
+
+ model_label = prettify_model_name(model)
+ model_text = sanitize_text(model_label)
+ filters.append(
+ f"drawtext={font_opt}text='{model_text}':fontcolor=white:fontsize={fontsize}"
+ f":box=1:boxcolor=black@0.6:boxborderw=10:x=(w-tw)/2:y=20"
+ )
+
+ if col == 0:
+ view_label = prettify_view_name(view)
+ view_text = sanitize_text(view_label)
+ filters.append(
+ f"drawtext={font_opt}text='{view_text}':fontcolor=white:fontsize={fontsize}"
+ f":box=1:boxcolor=black@0.6:boxborderw=10:x=20:y=h-th-40"
+ )
+
+ prev = f"[{in_label}]"
+ final_label = f"v{idx}"
+ if filters:
+ for f_idx, filt in enumerate(filters):
+ out_label = final_label if f_idx == len(filters) - 1 else f"v{idx}_{f_idx}"
+ filter_cmds.append(f"{prev}{filt}[{out_label}]")
+ prev = f"[{out_label}]"
+ else:
+ filter_cmds.append(f"{prev}null[{final_label}]")
+ stream_labels.append(final_label)
+
+ row_labels: List[str] = []
+ for r, view in enumerate(views):
+ cols = stream_labels[r * num_models : (r + 1) * num_models]
+ if len(cols) == 1:
+ out_label = f"row{r}"
+ filter_cmds.append(f"[{cols[0]}]null[{out_label}]")
+ row_labels.append(out_label)
+ else:
+ out_label = f"row{r}"
+ inputs_concat = "".join(f"[{label}]" for label in cols)
+ filter_cmds.append(f"{inputs_concat}hstack=inputs={len(cols)}[{out_label}]")
+ row_labels.append(out_label)
+
+ if len(row_labels) == 1:
+ final_label = row_labels[0]
+ else:
+ final_label = "grid"
+ inputs_concat = "".join(f"[{label}]" for label in row_labels)
+ filter_cmds.append(f"{inputs_concat}vstack=inputs={len(row_labels)}[{final_label}]")
+
+ filter_complex = ";".join(filter_cmds)
+ return filter_complex, final_label
+
+
+def build_ffmpeg_command(
+ binary: str,
+ inputs: List[Path],
+ filter_complex: str,
+ final_label: str,
+ output: Path,
+ codec: str,
+ bitrate: str,
+) -> List[str]:
+ cmd = [binary, "-y"]
+ for path in inputs:
+ cmd.extend(["-i", str(path)])
+ cmd.extend(
+ [
+ "-filter_complex",
+ filter_complex,
+ "-map",
+ f"[{final_label}]",
+ ]
+ )
+ if codec:
+ cmd.extend(["-c:v", codec])
+ if codec == "libx264":
+ cmd.extend(["-preset", "medium", "-crf", "18"])
+ elif codec.startswith("libx26"):
+ cmd.extend(["-preset", "medium", "-crf", "20"])
+ elif codec in {"libopenh264", "mpeg4"}:
+ cmd.extend(["-b:v", bitrate])
+ cmd.extend(["-pix_fmt", "yuv420p", "-an", str(output)])
+ return cmd
+
+
+def main() -> None:
+ args = parse_args()
+ ffmpeg_bin = resolve_ffmpeg_binary(args.ffmpeg_bin)
+
+ clip_id = args.clip_id.strip()
+ models = args.models if args.models else discover_models(args.root, clip_id)
+
+ per_model_videos: Dict[str, Dict[str, Path]] = {}
+ per_model_paths: Dict[str, Path] = {}
+ for model in models:
+ clip_dir = args.root / model / clip_id
+ if not clip_dir.is_dir():
+ raise SystemExit(f"Clip '{clip_id}' not found under {clip_dir.parent}")
+ novel_dir = select_novel_dir(clip_dir, args.novel_dir)
+ per_model_paths[model] = novel_dir
+ per_model_videos[model] = load_novel_videos(novel_dir)
+
+ common_views = set.intersection(*(set(v.keys()) for v in per_model_videos.values()))
+ if not common_views:
+ raise SystemExit("No common novel-view videos across the selected models.")
+ view_list = sorted(common_views)
+
+ inputs: List[Path] = []
+ for view in view_list:
+ for model in models:
+ inputs.append(per_model_videos[model][view])
+
+ output = args.output
+ if output is None:
+ output_dir = args.root / "grids"
+ output_dir.mkdir(parents=True, exist_ok=True)
+ output = output_dir / f"clip_{clip_id}_novel_grid.mp4"
+ else:
+ output.parent.mkdir(parents=True, exist_ok=True)
+
+ font_path = resolve_font_path(args.font)
+
+ if args.verbose:
+ print(f"Using font: {font_path}")
+
+ filter_complex, final_label = build_filter_complex(
+ inputs=inputs,
+ models=models,
+ views=view_list,
+ font=font_path,
+ fontsize=args.fontsize,
+ )
+
+ codec = choose_codec(ffmpeg_bin, args.codec)
+ if codec == "":
+ if args.verbose:
+ print("No specific encoder selected; FFmpeg default will be used.")
+ cmd = build_ffmpeg_command(ffmpeg_bin, inputs, filter_complex, final_label, output, codec, args.bitrate)
+
+ if args.verbose:
+ print("Running ffmpeg command:")
+ print(" ".join(cmd))
+
+ subprocess.run(cmd, check=True)
+ print(f"Saved grid video to {output}")
+ print("Models:", ", ".join(models))
+ print("Views:", ", ".join(view_list))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/third_party/drivestudio/tools/compose_sam_mask_video.py b/worldbench/third_party/drivestudio/tools/compose_sam_mask_video.py
new file mode 100644
index 0000000..9c4c8aa
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/compose_sam_mask_video.py
@@ -0,0 +1,250 @@
+#!/usr/bin/env python3
+"""Compose tiled SAM mask visualization videos for nuScenes clips.
+
+This script mirrors the interface style of ``tools/compose_depth_grid.py`` but
+targets semantic masks generated by SAM for the GT data only. It gathers
+per-frame PNG masks from
+
+ data/nuscenes_trainval/processed_gt/advanced_12Hz_trainval/{clip}/sam_mask/
+
+and builds a single MP4 video that tiles the requested cameras (typically all
+six nuScenes cameras) without any gaps between tiles.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import shutil
+import subprocess
+from pathlib import Path
+from typing import List
+
+
+CAMERA_NAME_TO_ID = {
+ "CAM_FRONT": 0,
+ "CAM_FRONT_LEFT": 1,
+ "CAM_FRONT_RIGHT": 2,
+ "CAM_BACK_LEFT": 3,
+ "CAM_BACK_RIGHT": 4,
+ "CAM_BACK": 5,
+}
+
+CAMERA_ID_TO_NAME = {idx: name for name, idx in CAMERA_NAME_TO_ID.items()}
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser("Compose SAM mask video for a nuScenes clip")
+ parser.add_argument("clip_id", type=str, help="Clip identifier (e.g., 002)")
+ parser.add_argument(
+ "--mask-root",
+ type=Path,
+ default=Path("data/nuscenes_trainval/processed_gt/advanced_12Hz_trainval"),
+ help=(
+ "Root directory that contains per-clip SAM masks. "
+ "Default matches processed nuScenes GT layout."
+ ),
+ )
+ parser.add_argument(
+ "--cameras",
+ type=str,
+ default="all",
+ help=(
+ "Comma-separated camera names (e.g. CAM_FRONT,CAM_BACK). "
+ "Use 'all' (default) to include every camera."
+ ),
+ )
+ parser.add_argument(
+ "--fps",
+ type=float,
+ default=12.0,
+ help="Frame rate for the output video (default: 12.0)",
+ )
+ parser.add_argument(
+ "--output",
+ type=Path,
+ default=None,
+ help=(
+ "Output MP4 path (default: work_dirs/sam_mask_videos/clip_{clip}_sam_masks.mp4)"
+ ),
+ )
+ parser.add_argument(
+ "--ffmpeg-bin",
+ type=str,
+ default="ffmpeg",
+ help="FFmpeg executable to use (default: ffmpeg)",
+ )
+ parser.add_argument(
+ "--codec",
+ type=str,
+ default=None,
+ help="Optional FFmpeg encoder (e.g. libx264)",
+ )
+ parser.add_argument(
+ "--bitrate",
+ type=str,
+ default="10M",
+ help="Bitrate to use with bitrate-based encoders (default: 10M)",
+ )
+ parser.add_argument(
+ "--verbose",
+ action="store_true",
+ help="Print FFmpeg commands before execution",
+ )
+ return parser.parse_args()
+
+
+def resolve_ffmpeg(binary: str) -> str:
+ for cand in [binary, os.environ.get("FFMPEG_BIN"), "ffmpeg"]:
+ if cand and shutil.which(cand):
+ return cand
+ raise SystemExit("FFmpeg not found; install ffmpeg or set --ffmpeg-bin/FFMPEG_BIN")
+
+
+def normalize_camera_name(name: str) -> str:
+ token = name.strip().upper()
+ if not token:
+ raise ValueError("Empty camera name provided")
+ if not token.startswith("CAM_"):
+ token = "CAM_" + token
+ if token not in CAMERA_NAME_TO_ID:
+ raise ValueError(f"Unknown camera name '{name}'")
+ return token
+
+
+def parse_camera_list(raw: str) -> List[str]:
+ if raw.lower() == "all":
+ return [CAMERA_ID_TO_NAME[idx] for idx in sorted(CAMERA_ID_TO_NAME)]
+ cameras: List[str] = []
+ for part in raw.split(","):
+ if not part.strip():
+ continue
+ cam = normalize_camera_name(part)
+ if cam not in cameras:
+ cameras.append(cam)
+ if not cameras:
+ raise SystemExit("No valid cameras specified via --cameras")
+ return cameras
+
+
+def collect_frame_paths(mask_dir: Path, cam_id: int) -> List[Path]:
+ return sorted(mask_dir.glob(f"*_{cam_id}.png"))
+
+
+def grid_shape(n: int) -> tuple[int, int]:
+ if n <= 0:
+ raise ValueError("grid_shape requires at least one input")
+ cols = max(1, int(round(n ** 0.5)))
+ while cols * cols < n:
+ cols += 1
+ rows = (n + cols - 1) // cols
+ return rows, cols
+
+
+def build_filter_complex(num_inputs: int) -> tuple[str, str]:
+ rows, cols = grid_shape(num_inputs)
+ cmds: List[str] = []
+ row_labels: List[str] = []
+ idx = 0
+ row_idx = 0
+ while idx < num_inputs:
+ count = min(cols, num_inputs - idx)
+ inputs_concat = "".join(f"[{i}:v]" for i in range(idx, idx + count))
+ row_label = f"row{row_idx}"
+ if count == 1:
+ cmds.append(f"{inputs_concat}null[{row_label}]")
+ else:
+ cmds.append(f"{inputs_concat}hstack=inputs={count}[{row_label}]")
+ row_labels.append(row_label)
+ idx += count
+ row_idx += 1
+ if len(row_labels) == 1:
+ final_label = row_labels[0]
+ else:
+ final_label = "stacked"
+ inputs_concat = "".join(f"[{label}]" for label in row_labels)
+ cmds.append(f"{inputs_concat}vstack=inputs={len(row_labels)}[{final_label}]")
+ return ";".join(cmds), final_label
+
+
+def build_ffmpeg_cmd(
+ ffmpeg: str,
+ patterns: List[Path],
+ fps: float,
+ output: Path,
+ codec: str | None,
+ bitrate: str,
+ filter_complex: str,
+ final_label: str,
+) -> List[str]:
+ cmd: List[str] = [ffmpeg, "-y"]
+ for pattern in patterns:
+ cmd.extend(["-framerate", str(fps), "-pattern_type", "glob", "-i", str(pattern)])
+ cmd.extend(["-filter_complex", filter_complex, "-map", f"[{final_label}]", "-pix_fmt", "yuv420p", "-an"])
+ if codec:
+ cmd.extend(["-c:v", codec])
+ if codec == "libx264":
+ cmd.extend(["-preset", "medium", "-crf", "18"])
+ elif codec in {"libopenh264", "mpeg4"}:
+ cmd.extend(["-b:v", bitrate])
+ cmd.append(str(output))
+ return cmd
+
+
+def main() -> None:
+ args = parse_args()
+ ffmpeg = resolve_ffmpeg(args.ffmpeg_bin)
+
+ clip_dir = args.mask_root / args.clip_id / "sam_mask"
+ if not clip_dir.is_dir():
+ raise SystemExit(f"Mask directory does not exist: {clip_dir}")
+
+ cameras = parse_camera_list(args.cameras)
+ frame_counts: List[int] = []
+ patterns: List[Path] = []
+
+ for camera in cameras:
+ cam_id = CAMERA_NAME_TO_ID[camera]
+ frames = collect_frame_paths(clip_dir, cam_id)
+ if not frames:
+ raise SystemExit(f"No SAM mask frames found for camera {camera} in {clip_dir}")
+ frame_counts.append(len(frames))
+ patterns.append(clip_dir / f"*_{cam_id}.png")
+
+ if len(set(frame_counts)) != 1:
+ raise SystemExit(
+ "Frame count mismatch across cameras: "
+ + ", ".join(f"{cam}={cnt}" for cam, cnt in zip(cameras, frame_counts))
+ )
+
+ filter_complex, final_label = build_filter_complex(len(cameras))
+
+ if args.output is None:
+ out_dir = Path("work_dirs") / "sam_mask_videos"
+ out_dir.mkdir(parents=True, exist_ok=True)
+ out_path = out_dir / f"clip_{args.clip_id}_sam_masks.mp4"
+ else:
+ out_path = args.output
+ out_path.parent.mkdir(parents=True, exist_ok=True)
+
+ cmd = build_ffmpeg_cmd(
+ ffmpeg,
+ patterns,
+ args.fps,
+ out_path,
+ args.codec,
+ args.bitrate,
+ filter_complex,
+ final_label,
+ )
+ if args.verbose:
+ print("FFmpeg command:")
+ print(" ".join(cmd))
+ subprocess.run(cmd, check=True)
+ print(
+ f"Saved tiled SAM mask video ({frame_counts[0]} frames @ {args.fps} fps, {len(cameras)} cameras) to {out_path}"
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/third_party/drivestudio/tools/compose_train_view_grid.py b/worldbench/third_party/drivestudio/tools/compose_train_view_grid.py
new file mode 100644
index 0000000..3e70797
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/compose_train_view_grid.py
@@ -0,0 +1,522 @@
+#!/usr/bin/env python3
+"""
+Compose a tiled RGB (training view) grid video across multiple methods for a given clip.
+
+Source layout (per method/run):
+ work_dirs/
+ /
+ /
+ videos/
+ full_set_.mp4
+
+The script mirrors tools/compose_depth_grid.py but uses the RGB videos instead of depth.
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import shutil
+import subprocess
+from pathlib import Path
+from typing import List, Optional, Sequence, Tuple, Dict
+import json
+
+
+def parse_args() -> argparse.Namespace:
+ p = argparse.ArgumentParser("Compose RGB training-view grid video for a given nuScenes clip")
+ p.add_argument("clip_id", type=str, help="Clip identifier (e.g., 084)")
+ p.add_argument("--root", type=Path, default=Path("work_dirs"), help="Root directory containing per-model runs")
+ p.add_argument("--models", type=str, nargs="*", help="Explicit list of model directory names; default: all containing the clip")
+ p.add_argument("--step", type=int, default=30000, help="Training step used in file naming (default: 30000)")
+ p.add_argument("--output", type=Path, default=None, help="Output MP4 path (default: work_dirs/grids/clip__train_grid.mp4)")
+ p.add_argument("--ffmpeg-bin", type=str, default="ffmpeg", help="FFmpeg executable (default: ffmpeg)")
+ p.add_argument("--codec", type=str, default=None, help="Preferred FFmpeg encoder (e.g., libx264)")
+ p.add_argument("--bitrate", type=str, default="10M", help="Bitrate for bitrate-based encoders (default: 10M)")
+ p.add_argument("--font", type=Path, default=None, help="Optional TrueType/OpenType font for labels")
+ p.add_argument("--fontsize", type=int, default=28, help="Label font size (default: 28)")
+ p.add_argument("--gap", type=int, default=8, help="Gap (pixels) between tiles (default: 8)")
+ p.add_argument("--tile-height", type=int, default=0, help="Optional per-tile fixed height (px); 0 disables scaling.")
+ p.add_argument("--tile-width", type=int, default=0, help="Optional per-tile fixed width (px); when height or width >0, tiles are resized to this WxH.")
+ p.add_argument("--fit", type=str, choices=["contain", "cover"], default="contain", help="Resize mode when both tile width/height are set: contain (scale+pad) or cover (scale+crop). Default: contain.")
+ # Optional manual crop to remove baked-in letterbox bars before resizing (pixels on source frames)
+ p.add_argument("--crop-top", type=int, default=0, help="Crop N pixels from top before resizing (default: 0)")
+ p.add_argument("--crop-bottom", type=int, default=0, help="Crop N pixels from bottom before resizing (default: 0)")
+ p.add_argument("--crop-left", type=int, default=0, help="Crop N pixels from left before resizing (default: 0)")
+ p.add_argument("--crop-right", type=int, default=0, help="Crop N pixels from right before resizing (default: 0)")
+ p.add_argument("--verbose", action="store_true", help="Print ffmpeg command")
+ p.add_argument("--fps", type=int, default=0, help="Target output FPS override; 0 = keep source FPS (default)")
+ p.add_argument("--frames", type=int, default=0, help="Trim output to this many frames after stacking (0=disable)")
+ p.add_argument("--debug", action="store_true", help="Print detailed alignment info (models, inputs, probed FPS/frames, metric values)")
+ # Metric overlay options (LPIPS for RGB by default)
+ p.add_argument("--metrics-json", type=Path, default=Path("work_dirs_analysis_metrics.json"), help="Path to metrics JSON for LPIPS overlay (default: work_dirs_analysis_metrics.json)")
+ p.add_argument("--metric-key", type=str, default="image_metrics/full/lpips", help="Metric key in JSON to overlay (default: image_metrics/full/lpips)")
+ p.add_argument("--metric-label", type=str, default="LPIPS", help="Display label for the metric (default: LPIPS)")
+ p.add_argument("--no-metric", action="store_true", help="Disable metric overlay on labels")
+ return p.parse_args()
+
+
+def resolve_ffmpeg(binary: str) -> str:
+ for cand in [binary, os.environ.get("FFMPEG_BIN"), "ffmpeg"]:
+ if cand and shutil.which(cand):
+ return cand
+ raise SystemExit("FFmpeg not found; install ffmpeg or set --ffmpeg-bin/FFMPEG_BIN")
+
+
+def resolve_ffprobe(ffmpeg_bin: str) -> str:
+ p = Path(ffmpeg_bin)
+ cand = p.parent / "ffprobe" if p.parent else None
+ if cand and cand.exists():
+ return str(cand)
+ probed = shutil.which("ffprobe")
+ if probed:
+ return probed
+ raise SystemExit("ffprobe not found; install ffmpeg (includes ffprobe) or add to PATH")
+
+
+def probe_fps(ffprobe_bin: str, video: Path) -> Optional[float]:
+ try:
+ out = subprocess.check_output([
+ ffprobe_bin,
+ "-v", "error",
+ "-select_streams", "v:0",
+ "-show_entries", "stream=avg_frame_rate",
+ "-of", "default=nw=1:nk=1",
+ str(video),
+ ], text=True)
+ fr = out.strip()
+ if not fr:
+ return None
+ if "/" in fr:
+ num, den = fr.split("/", 1)
+ num = float(num); den = float(den)
+ return num / den if den else None
+ return float(fr)
+ except Exception:
+ return None
+
+def probe_nb_frames(ffprobe_bin: str, video: Path) -> Optional[int]:
+ try:
+ out = subprocess.check_output([
+ ffprobe_bin,
+ "-v", "error",
+ "-select_streams", "v:0",
+ "-count_frames",
+ "-show_entries", "stream=nb_read_frames",
+ "-of", "csv=p=0",
+ str(video),
+ ], text=True)
+ s = out.strip()
+ return int(s) if s.isdigit() else None
+ except Exception:
+ return None
+
+
+CANDIDATE_FONT_NAMES = [
+ "DejaVuSans.ttf",
+ "DejaVuSans-Regular.ttf",
+ "LiberationSans-Regular.ttf",
+ "NotoSans-Regular.ttf",
+]
+
+
+def resolve_font_path(user_font: Optional[Path]) -> Optional[Path]:
+ if user_font:
+ if user_font.is_file():
+ return user_font
+ raise SystemExit(f"Specified font file '{user_font}' does not exist.")
+
+ env_font = os.environ.get("FFMPEG_FONT")
+ if env_font and Path(env_font).is_file():
+ return Path(env_font)
+
+ candidates: List[Path] = []
+ system_dirs = [
+ Path("/usr/share/fonts"),
+ Path("/usr/local/share/fonts"),
+ Path.home() / ".fonts",
+ ]
+ for base in system_dirs:
+ for name in CANDIDATE_FONT_NAMES:
+ candidates.append(base / "truetype" / name)
+ candidates.append(base / "truetype" / "dejavu" / name)
+ candidates.append(base / name)
+
+ conda_prefix = os.environ.get("CONDA_PREFIX")
+ if conda_prefix:
+ cp = Path(conda_prefix)
+ for name in CANDIDATE_FONT_NAMES:
+ candidates.append(cp / "lib" / "python3.10" / "site-packages" / name)
+ candidates.append(cp / "lib" / "python3.10" / "site-packages" / "matplotlib" / "mpl-data" / "fonts" / "ttf" / name)
+ pkgs_dir = cp / "pkgs"
+ if pkgs_dir.is_dir():
+ for name in CANDIDATE_FONT_NAMES:
+ candidates.extend(pkgs_dir.glob(f"pillow*/info/test/Tests/fonts/{name}"))
+
+ for cand in candidates:
+ if cand.is_file() and "Symbols" not in cand.name and "Color" not in cand.name:
+ return cand
+
+ return None
+
+
+def discover_models(root: Path, clip_id: str, explicit: Optional[Sequence[str]]) -> List[str]:
+ if explicit:
+ return list(explicit)
+ models = []
+ for d in sorted(root.iterdir()):
+ if (d / clip_id).is_dir():
+ models.append(d.name)
+ if not models:
+ raise SystemExit(f"No models containing clip '{clip_id}' found under {root}")
+ models.sort()
+ if "drivestudio-nus-gt" in models:
+ models.remove("drivestudio-nus-gt")
+ models.insert(0, "drivestudio-nus-gt")
+ return models
+
+
+def method_label(name: str) -> str:
+ return name.split("drivestudio-nus-", 1)[-1].replace("_", " ").replace("-", " ")
+
+
+def find_rgb_video(root: Path, model: str, clip_id: str, step: int) -> Path:
+ base = root / model / clip_id / "videos"
+ target = base / f"full_set_{step}_rgbs.mp4"
+ if target.is_file():
+ return target
+ raise SystemExit(f"Missing RGB training video: {target}")
+
+
+def grid_shape(n: int) -> Tuple[int, int]:
+ # Prefer fewer rows to reduce perceived row spacing and placeholders at the last row.
+ # Choose r=floor(sqrt(n)), c=ceil(n/r). Ensures r*c>=n and r<=c.
+ if n <= 0:
+ return 0, 0
+ r = int(n ** 0.5)
+ if r < 1:
+ r = 1
+ from math import ceil
+ c = int(ceil(n / r))
+ return r, c
+
+
+def drawtext_filter(font: Optional[Path], fontsize: int, text: str) -> str:
+ if not font:
+ return ""
+ safe_text = text.replace(":", "\\:").replace("'", "\\'")
+ return (
+ f"drawtext=fontfile={font.as_posix()}:text='{safe_text}':fontcolor=white:fontsize={fontsize}"
+ ":box=1:boxcolor=black@0.6:boxborderw=10:x=20:y=20"
+ )
+
+
+def build_filter_complex(
+ inputs: List[Path],
+ labels: List[str],
+ rows: int,
+ cols: int,
+ font: Optional[Path],
+ fontsize: int,
+ gap: int,
+ tile_height: int,
+ tile_width: int,
+ fps: int = 0,
+ frames: int = 0,
+ fit: str = "contain",
+ crop: Tuple[int, int, int, int] = (0, 0, 0, 0),
+ placeholder_height: int = 0,
+ row_target_widths: Optional[List[int]] = None,
+) -> Tuple[str, str]:
+ if len(inputs) != len(labels):
+ raise ValueError("inputs and labels length mismatch")
+ cmds: List[str] = []
+ out_labels: List[str] = []
+ # Determine target tile size (letterbox to exact WxH if provided)
+ th = tile_height if tile_height and tile_height > 0 else 0
+ # Height-only mode (tw==0) will use scale=-2:th (no internal vertical padding). If tw>0, use contain/cover path.
+ tw = tile_width if tile_width and tile_width > 0 else 0
+
+ for i, lbl in enumerate(labels):
+ filt = drawtext_filter(font, fontsize, lbl)
+ out_name = f"v{i}"
+ chain = f"[{i}:v]"
+ if filt:
+ chain += filt
+ else:
+ chain += "null"
+ # Optional manual crop (remove baked-in bars): crop=iw-left-right:ih-top-bottom:left:top
+ ct, cb, cl, cr = crop
+ if any(v > 0 for v in (ct, cb, cl, cr)):
+ chain += f",crop=iw-{cl}-{cr}:ih-{ct}-{cb}:{cl}:{ct}"
+ if th:
+ if fit == "cover":
+ # Fill the tile: scale up to cover, then center-crop to exact WxH (no internal bars)
+ chain += f",scale={tw if tw else -2}:{th}:force_original_aspect_ratio=increase"
+ if tw:
+ chain += f",crop={tw}:{th}:(iw-{tw})/2:(ih-{th})/2"
+ else:
+ # contain: fit inside tile, then pad to exact WxH (may have bars)
+ chain += f",scale={tw if tw else -2}:{th}:force_original_aspect_ratio=decrease"
+ if tw:
+ chain += f",pad={tw}:{th}:(ow-iw)/2:(oh-ih)/2:color=black"
+ cmds.append(f"{chain}[{out_name}]")
+ out_labels.append(out_name)
+
+ total = rows * cols
+ pads = 0
+ while len(out_labels) < total:
+ pad_label = f"pad{pads}"
+ if th:
+ # In height-only mode (tw==0) use a small placeholder width; rows will be padded to target width later.
+ size_str = f"size={(tw if tw else 64)}x{th}"
+ else:
+ # match row video heights by using a reasonable placeholder height
+ ph = placeholder_height if placeholder_height and placeholder_height > 0 else 16
+ size_str = f"size=16x{ph}"
+ if fps and fps > 0:
+ cmds.append(f"color=c=black:{size_str}:rate={fps},format=rgb24[{pad_label}]")
+ else:
+ cmds.append(f"color=c=black:{size_str},format=rgb24[{pad_label}]")
+ out_labels.append(pad_label)
+ pads += 1
+
+ row_outs: List[str] = []
+ for r in range(rows):
+ seg = out_labels[r * cols : (r + 1) * cols]
+ padded: List[str] = []
+ for c_idx, s in enumerate(seg):
+ cur = s
+ if gap > 0 and c_idx < cols - 1:
+ pad_label = f"{s}_pad"
+ cmds.append(f"[{s}]pad=iw+{gap}:ih:0:0:color=black[{pad_label}]")
+ cur = pad_label
+ padded.append(cur)
+ tmp_row = f"row{r}_tmp"
+ if len(padded) == 1:
+ cmds.append(f"[{padded[0]}]null[{tmp_row}]")
+ else:
+ inputs_concat = "".join(f"[{s}]" for s in padded)
+ cmds.append(f"{inputs_concat}hstack=inputs={len(padded)}:shortest=1[{tmp_row}]")
+ # If using height-only mode (tw==0) and we computed target widths, right-pad the row to the exact width so vstack matches.
+ row_label = f"row{r}"
+ if tw == 0 and row_target_widths is not None and r < len(row_target_widths) and row_target_widths[r] > 0:
+ target_w = row_target_widths[r]
+ cmds.append(f"[{tmp_row}]pad={target_w}:ih:0:0:color=black[{row_label}_padw]")
+ tmp_row = f"{row_label}_padw"
+ if gap > 0 and r < rows - 1:
+ cmds.append(f"[{tmp_row}]pad=iw:ih+{gap}:0:0:color=black[{row_label}]")
+ else:
+ cmds.append(f"[{tmp_row}]null[{row_label}]")
+ row_outs.append(row_label)
+
+ if len(row_outs) == 1:
+ final = row_outs[0]
+ else:
+ final = "grid"
+ inputs_concat = "".join(f"[{s}]" for s in row_outs)
+ cmds.append(f"{inputs_concat}vstack=inputs={len(row_outs)}:shortest=1[{final}]")
+
+ tail_parts = []
+ if fps and fps > 0:
+ tail_parts.append(f"fps={fps}")
+ if frames and frames > 0:
+ tail_parts.append(f"trim=end_frame={frames}")
+ tail_parts.append("setsar=1")
+ tail = ",".join(tail_parts)
+ cmds.append(f"[{final}]{tail}[outv]")
+ return ";".join(cmds), "outv"
+
+
+def build_ffmpeg_cmd(binary: str, inputs: List[Path], filter_complex: str, final_label: str, output: Path, codec: Optional[str], bitrate: str) -> List[str]:
+ cmd: List[str] = [binary, "-y"]
+ for p in inputs:
+ cmd.extend(["-i", str(p)])
+ cmd.extend(["-filter_complex", filter_complex, "-map", f"[{final_label}]", "-pix_fmt", "yuv420p", "-vsync", "2", "-shortest", "-an"])
+ if codec:
+ cmd.extend(["-c:v", codec])
+ if codec == "libx264":
+ cmd.extend(["-preset", "medium", "-crf", "18"])
+ elif codec in {"libopenh264", "mpeg4"}:
+ cmd.extend(["-b:v", bitrate])
+ cmd.append(str(output))
+ return cmd
+
+
+def main() -> None:
+ args = parse_args()
+ root = args.root
+ clip_id = args.clip_id
+ models = discover_models(root, clip_id, args.models)
+ inputs = [find_rgb_video(root, m, clip_id, args.step) for m in models]
+
+ # Optional metric overlay per method from JSON
+ metrics: Dict[str, Optional[float]] = {}
+ if (not args.no_metric) and args.metrics_json and args.metrics_json.is_file():
+ try:
+ data = json.loads(args.metrics_json.read_text())
+ for m in models:
+ seq_map = data.get(m, {})
+ val = None
+ if isinstance(seq_map, dict):
+ clip_metrics = seq_map.get(clip_id, {})
+ if isinstance(clip_metrics, dict):
+ v = clip_metrics.get(args.metric_key)
+ if isinstance(v, (int, float)):
+ val = float(v)
+ metrics[m] = val
+ except Exception:
+ metrics = {m: None for m in models}
+ else:
+ metrics = {m: None for m in models}
+
+ labels = []
+ for m in models:
+ base = method_label(m)
+ v = metrics.get(m)
+ if v is not None:
+ labels.append(f"{base} | {args.metric_label} {v:.3f}")
+ else:
+ labels.append(base)
+
+ rows, cols = grid_shape(len(inputs))
+ font = resolve_font_path(args.font)
+ ffmpeg_bin = resolve_ffmpeg(args.ffmpeg_bin)
+ # detect source fps from the first input if user didn't override
+ if args.fps > 0:
+ fps_for_graph = args.fps
+ else:
+ try:
+ ffprobe_bin = resolve_ffprobe(ffmpeg_bin)
+ fps_detected = probe_fps(ffprobe_bin, inputs[0]) if inputs else None
+ fps_for_graph = int(fps_detected) if fps_detected and fps_detected > 0 else 0
+ except Exception:
+ fps_for_graph = 0
+ # Determine placeholder height when not resizing tiles
+ ph = 0
+ if (not args.tile_height or args.tile_height == 0) and inputs:
+ try:
+ ffprobe_bin = resolve_ffprobe(ffmpeg_bin)
+ # probe dimensions via ffprobe
+ info = subprocess.check_output([
+ ffprobe_bin,
+ "-v","error",
+ "-select_streams","v:0",
+ "-show_entries","stream=height",
+ "-of","csv=p=0",
+ str(inputs[0])
+ ], text=True).strip()
+ ph = int(info) if info.isdigit() else 0
+ except Exception:
+ ph = 0
+
+ # If only --tile-height is provided, switch to height-only mode (scale=-2:th) and pad each row to max width.
+ tile_width_eff = args.tile_width or 0
+ row_targets: Optional[List[int]] = None
+ if (args.tile_height and args.tile_height > 0) and (not args.tile_width or args.tile_width == 0) and inputs:
+ try:
+ ffprobe_bin = resolve_ffprobe(ffmpeg_bin)
+ except Exception:
+ ffprobe_bin = None
+ # probe widths/heights to predict scaled widths
+ wh_list: List[Tuple[int,int]] = []
+ for pth in inputs:
+ w = h = 0
+ if ffprobe_bin:
+ try:
+ s = subprocess.check_output([
+ ffprobe_bin,
+ "-v","error",
+ "-select_streams","v:0",
+ "-show_entries","stream=width,height",
+ "-of","csv=p=0",
+ str(pth)
+ ], text=True).strip()
+ parts = s.split(',')
+ if len(parts) >= 2:
+ w = int(parts[0]); h = int(parts[1])
+ except Exception:
+ w = h = 0
+ wh_list.append((w, h))
+ # compute scaled widths per input for height-only scaling; -2 makes width even
+ sw: List[int] = []
+ th = args.tile_height
+ for (w, h) in wh_list:
+ if w > 0 and h > 0:
+ ww = int(round(w * (th / float(h))))
+ if ww % 2 != 0:
+ ww += 1
+ else:
+ ww = 1280 # reasonable fallback
+ sw.append(ww)
+ # derive row target widths
+ row_widths: List[int] = []
+ for r in range(rows):
+ seg = sw[r * cols : (r + 1) * cols]
+ while len(seg) < cols:
+ seg.append(64)
+ row_widths.append(sum(seg) + (args.gap * (cols - 1 if args.gap > 0 else 0)))
+ max_w = max(row_widths) if row_widths else 0
+ row_targets = [max_w for _ in range(rows)]
+
+ filter_str, final_label = build_filter_complex(
+ inputs, labels, rows, cols, font, args.fontsize, args.gap,
+ args.tile_height, tile_width_eff, fps=fps_for_graph, frames=args.frames, fit=args.fit,
+ crop=(args.crop_top, args.crop_bottom, args.crop_left, args.crop_right), placeholder_height=ph,
+ row_target_widths=row_targets)
+
+ # ffmpeg_bin already resolved above
+ if args.output is None:
+ out_dir = root / "grids"
+ out_dir.mkdir(parents=True, exist_ok=True)
+ output = out_dir / f"clip_{clip_id}_train_grid.mp4"
+ else:
+ output = args.output
+ output.parent.mkdir(parents=True, exist_ok=True)
+
+ # Debug printout: alignment verification
+ if args.debug:
+ print("[DEBUG] clip:", clip_id, "step:", args.step)
+ print("[DEBUG] grid:", f"rows={rows} cols={cols} gap={args.gap} tile=({('auto' if tile_width_eff==0 else tile_width_eff)}x{args.tile_height or 'auto'})")
+ print("[DEBUG] models order:", ", ".join(models))
+ try:
+ ffprobe_bin = resolve_ffprobe(ffmpeg_bin)
+ except Exception:
+ ffprobe_bin = None
+ for m, path, lab in zip(models, inputs, labels):
+ fps = probe_fps(ffprobe_bin, path) if ffprobe_bin else None
+ nbf = probe_nb_frames(ffprobe_bin, path) if ffprobe_bin else None
+ wh = None
+ sar = dar = None
+ if ffprobe_bin:
+ try:
+ whs = subprocess.check_output([
+ ffprobe_bin,
+ "-v","error",
+ "-select_streams","v:0",
+ "-show_entries","stream=width,height,sample_aspect_ratio,display_aspect_ratio",
+ "-of","csv=p=0",
+ str(path)
+ ], text=True).strip()
+ parts = whs.split('\n')[0].split(',')
+ if len(parts) >= 2:
+ wh = f"{parts[0]}x{parts[1]}"
+ if len(parts) >= 3:
+ sar = parts[2]
+ if len(parts) >= 4:
+ dar = parts[3]
+ except Exception:
+ pass
+ print(f"[DEBUG] {m}: video={path} size={wh} sar={sar} dar={dar} fps={fps} frames={nbf} label=\"{lab}\"")
+ # Metric source for RGB
+ print("[DEBUG] metric source:", args.metrics_json, "key:", args.metric_key)
+
+ cmd = build_ffmpeg_cmd(ffmpeg_bin, inputs, filter_str, final_label, output, args.codec, args.bitrate)
+
+ if args.verbose:
+ print("Running:", " ".join(cmd))
+ subprocess.run(cmd, check=True)
+ print(f"Wrote {output}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/third_party/drivestudio/tools/eval.py b/worldbench/third_party/drivestudio/tools/eval.py
new file mode 100644
index 0000000..00f7b47
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/eval.py
@@ -0,0 +1,350 @@
+from typing import List, Optional
+from omegaconf import OmegaConf
+import os
+import numpy as np
+import time
+import json
+import wandb
+import logging
+import argparse
+
+import torch
+from datasets.driving_dataset import DrivingDataset
+from utils.misc import import_str
+from models.trainers import BasicTrainer
+from models.video_utils import (
+ render_images,
+ save_videos,
+ render_novel_views
+)
+from utils.visualization import depth_visualizer, to8b
+import imageio
+
+logger = logging.getLogger()
+current_time = time.strftime("%Y-%m-%d_%H-%M-%S", time.localtime())
+
+@torch.no_grad()
+def do_evaluation(
+ step: int = 0,
+ cfg: OmegaConf = None,
+ trainer: BasicTrainer = None,
+ dataset: DrivingDataset = None,
+ args: argparse.Namespace = None,
+ render_keys: Optional[List[str]] = None,
+ post_fix: str = "",
+ log_metrics: bool = True
+):
+ trainer.set_eval()
+
+ logger.info("Evaluating Pixels...")
+ if dataset.test_image_set is not None and cfg.render.render_test:
+ logger.info("Evaluating Test Set Pixels...")
+ render_results = render_images(
+ trainer=trainer,
+ dataset=dataset.test_image_set,
+ compute_metrics=True,
+ compute_error_map=cfg.render.vis_error,
+ )
+
+ if log_metrics:
+ eval_dict = {}
+ for k, v in render_results.items():
+ if k in [
+ "psnr",
+ "ssim",
+ "lpips",
+ "occupied_psnr",
+ "occupied_ssim",
+ "masked_psnr",
+ "masked_ssim",
+ "human_psnr",
+ "human_ssim",
+ "vehicle_psnr",
+ "vehicle_ssim",
+ ]:
+ eval_dict[f"image_metrics/test/{k}"] = v
+ if args.enable_wandb:
+ wandb.log(eval_dict)
+ test_metrics_file = f"{cfg.log_dir}/metrics{post_fix}/images_test_{current_time}.json"
+ with open(test_metrics_file, "w") as f:
+ json.dump(eval_dict, f)
+ logger.info(f"Image evaluation metrics saved to {test_metrics_file}")
+
+ if args.render_video_postfix is None:
+ video_output_pth = f"{cfg.log_dir}/videos{post_fix}/test_set_{step}.mp4"
+ else:
+ video_output_pth = (
+ f"{cfg.log_dir}/videos{post_fix}/test_set_{step}_{args.render_video_postfix}.mp4"
+ )
+ vis_frame_dict = save_videos(
+ render_results,
+ video_output_pth,
+ layout=dataset.layout,
+ num_timestamps=dataset.num_test_timesteps,
+ keys=render_keys,
+ num_cams=dataset.pixel_source.num_cams,
+ save_seperate_video=cfg.logging.save_seperate_video,
+ fps=2,
+ verbose=True,
+ save_images=False,
+ )
+ if args.enable_wandb:
+ for k, v in vis_frame_dict.items():
+ wandb.log({"image_rendering/test/" + k: wandb.Image(v)})
+ # optionally dump raw depth for test set
+ if getattr(args, "save_raw_depth", None) and getattr(args, "include_depth", False):
+ base_out = args.save_raw_depth
+ out_base = base_out if os.path.isabs(base_out) else os.path.join(cfg.log_dir, base_out)
+ raw_dir = os.path.join(out_base, f"test_set_{step}")
+ os.makedirs(raw_dir, exist_ok=True)
+ if "depths" in render_results:
+ for i, depth in enumerate(render_results["depths"]):
+ cam_name = render_results["cam_names"][i]
+ np.save(os.path.join(raw_dir, f"{i:05d}_{cam_name}.npy"), depth)
+ logger.info(f"Saved raw test depth to {raw_dir}")
+ del render_results, vis_frame_dict
+ torch.cuda.empty_cache()
+
+ if cfg.render.render_full:
+ logger.info("Evaluating Full Set...")
+ render_results = render_images(
+ trainer=trainer,
+ dataset=dataset.full_image_set,
+ compute_metrics=True,
+ compute_error_map=cfg.render.vis_error,
+ save_image_pairs=True
+ )
+
+ if log_metrics:
+ eval_dict = {}
+ for k, v in render_results.items():
+ if k in [
+ "psnr",
+ "ssim",
+ "lpips",
+ "occupied_psnr",
+ "occupied_ssim",
+ "masked_psnr",
+ "masked_ssim",
+ "human_psnr",
+ "human_ssim",
+ "vehicle_psnr",
+ "vehicle_ssim",
+ ]:
+ eval_dict[f"image_metrics/full/{k}"] = v
+ if args.enable_wandb:
+ wandb.log(eval_dict)
+ full_metrics_file = f"{cfg.log_dir}/metrics{post_fix}/images_full_{current_time}.json"
+ with open(full_metrics_file, "w") as f:
+ json.dump(eval_dict, f)
+ logger.info(f"Image evaluation metrics saved to {full_metrics_file}")
+
+ if args.render_video_postfix is None:
+ video_output_pth = f"{cfg.log_dir}/videos{post_fix}/full_set_{step}.mp4"
+ else:
+ video_output_pth = (
+ f"{cfg.log_dir}/videos{post_fix}/full_set_{step}_{args.render_video_postfix}.mp4"
+ )
+ vis_frame_dict = save_videos(
+ render_results,
+ video_output_pth,
+ layout=dataset.layout,
+ num_timestamps=dataset.num_img_timesteps,
+ keys=render_keys,
+ num_cams=dataset.pixel_source.num_cams,
+ save_seperate_video=cfg.logging.save_seperate_video,
+ fps=cfg.render.fps,
+ verbose=True,
+ )
+ if args.enable_wandb:
+ for k, v in vis_frame_dict.items():
+ wandb.log({"image_rendering/full/" + k: wandb.Image(v)})
+ # optionally dump raw depth for full set
+ if getattr(args, "save_raw_depth", None) and getattr(args, "include_depth", False):
+ base_out = args.save_raw_depth
+ out_base = base_out if os.path.isabs(base_out) else os.path.join(cfg.log_dir, base_out)
+ raw_dir = os.path.join(out_base, f"full_set_{step}")
+ os.makedirs(raw_dir, exist_ok=True)
+ if "depths" in render_results:
+ for i, depth in enumerate(render_results["depths"]):
+ cam_name = render_results["cam_names"][i]
+ np.save(os.path.join(raw_dir, f"{i:05d}_{cam_name}.npy"), depth)
+ logger.info(f"Saved raw full depth to {raw_dir}")
+ del render_results, vis_frame_dict
+ torch.cuda.empty_cache()
+
+ render_novel_cfg = cfg.render.get("render_novel", None)
+ if render_novel_cfg is not None:
+ logger.info("Rendering novel views...")
+ render_traj = dataset.get_novel_render_traj(
+ traj_types=['s_curve', 'front_center_interp', 'lateral_offset'], # render_novel_cfg.traj_types,
+ target_frames=render_novel_cfg.get("frames", dataset.frame_num),
+ )
+ video_output_dir = f"{cfg.log_dir}/videos{post_fix}/novel_{step}"
+ if not os.path.exists(video_output_dir):
+ os.makedirs(video_output_dir)
+
+ # determine which cameras to render for novel views
+ cam_selector = getattr(args, "novel_cam_ids", None)
+ if cam_selector is None:
+ cam_ids = [None] # keep default behavior
+ else:
+ if isinstance(cam_selector, str) and cam_selector.strip().lower() == "all":
+ cam_ids = dataset.pixel_source.camera_list
+ else:
+ try:
+ cam_ids = [int(x) for x in str(cam_selector).split(',') if x != '']
+ except Exception:
+ cam_ids = [None]
+
+ for traj_type, traj in render_traj.items():
+ for cid in cam_ids:
+ # Prepare rendering data for selected camera
+ render_data = dataset.prepare_novel_view_render_data(traj, cam_id=cid)
+ # Render and save video
+ suffix = "" if cid is None else f"_cam{cid}"
+ save_path = os.path.join(video_output_dir, f"{traj_type}{suffix}.mp4")
+ render_novel_views(
+ trainer, render_data, save_path,
+ fps=render_novel_cfg.get("fps", cfg.render.fps)
+ )
+ logger.info(
+ f"Saved novel view video for trajectory type: {traj_type}"
+ + ("" if cid is None else f", cam {cid}")
+ + f" to {save_path}"
+ )
+ # Optionally dump raw depth for novel views
+ if getattr(args, "save_raw_depth", None):
+ base_out = args.save_raw_depth
+ out_base = base_out if os.path.isabs(base_out) else os.path.join(cfg.log_dir, base_out)
+ raw_dir = os.path.join(out_base, f"novel_{step}", f"{traj_type}{suffix}")
+ os.makedirs(raw_dir, exist_ok=True)
+ # If also asked to include depth visualization, open a writer
+ writer = None
+ if getattr(args, "include_depth", False):
+ depth_mp4 = os.path.join(video_output_dir, f"{traj_type}{suffix}_depths.mp4")
+ writer = imageio.get_writer(depth_mp4, mode='I', fps=render_novel_cfg.get("fps", cfg.render.fps))
+ for i, frame_data in enumerate(render_data):
+ # move to device
+ for key, value in frame_data["cam_infos"].items():
+ frame_data["cam_infos"][key] = value.cuda(non_blocking=True)
+ for key, value in frame_data["image_infos"].items():
+ frame_data["image_infos"][key] = value.cuda(non_blocking=True)
+ outputs = trainer(
+ image_infos=frame_data["image_infos"],
+ camera_infos=frame_data["cam_infos"],
+ novel_view=True
+ )
+ depth = outputs["depth"].detach().cpu().numpy()
+ np.save(os.path.join(raw_dir, f"{i:05d}.npy"), depth)
+ if writer is not None:
+ opacity = outputs.get("opacity").detach().cpu().numpy() if "opacity" in outputs else None
+ colored = depth_visualizer(depth, opacity)
+ writer.append_data(to8b(colored))
+ if writer is not None:
+ writer.close()
+ logger.info(f"Saved novel depth video to {depth_mp4}")
+ logger.info(f"Saved raw novel depth to {raw_dir}")
+
+def main(args):
+ log_dir = os.path.dirname(args.resume_from)
+ cfg = OmegaConf.load(os.path.join(log_dir, "config.yaml"))
+ cfg = OmegaConf.merge(cfg, OmegaConf.from_cli(args.opts))
+ args.enable_wandb = False
+ for folder in ["videos_eval", "metrics_eval"]:
+ os.makedirs(os.path.join(log_dir, folder), exist_ok=True)
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+ # build dataset
+ dataset = DrivingDataset(data_cfg=cfg.data)
+
+ # setup trainer
+ trainer = import_str(cfg.trainer.type)(
+ **cfg.trainer,
+ num_timesteps=dataset.num_img_timesteps,
+ model_config=cfg.model,
+ num_train_images=len(dataset.train_image_set),
+ num_full_images=len(dataset.full_image_set),
+ test_set_indices=dataset.test_timesteps,
+ scene_aabb=dataset.get_aabb().reshape(2, 3),
+ device=device
+ )
+
+ # Resume from checkpoint
+ trainer.resume_from_checkpoint(
+ ckpt_path=args.resume_from,
+ load_only_model=True
+ )
+ logger.info(
+ f"Resuming training from {args.resume_from}, starting at step {trainer.step}"
+ )
+
+ if args.enable_viewer:
+ # a simple viewer for background visualization
+ trainer.init_viewer(port=args.viewer_port)
+
+ # define render keys
+ render_keys = [
+ "gt_rgbs",
+ "rgbs",
+ "Background_rgbs",
+ "RigidNodes_rgbs",
+ "DeformableNodes_rgbs",
+ "SMPLNodes_rgbs",
+ # "depths",
+ # "Background_depths",
+ # "RigidNodes_depths",
+ # "DeformableNodes_depths",
+ # "SMPLNodes_depths",
+ # "mask"
+ ]
+ if cfg.render.vis_lidar:
+ render_keys.insert(0, "lidar_on_images")
+ if cfg.render.vis_sky:
+ render_keys += ["rgb_sky_blend", "rgb_sky"]
+ if cfg.render.vis_error:
+ render_keys.insert(render_keys.index("rgbs") + 1, "rgb_error_maps")
+
+ if args.save_catted_videos:
+ cfg.logging.save_seperate_video = False
+ # optionally include depth renders in output videos
+ if getattr(args, "include_depth", False):
+ if "depths" not in render_keys:
+ render_keys.append("depths")
+
+ do_evaluation(
+ step=trainer.step,
+ cfg=cfg,
+ trainer=trainer,
+ dataset=dataset,
+ render_keys=render_keys,
+ args=args,
+ post_fix="_eval"
+ )
+
+ if args.enable_viewer:
+ print("Viewer running... Ctrl+C to exit.")
+ time.sleep(1000000)
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser("Train Gaussian Splatting for a single scene")
+ # eval
+ parser.add_argument("--resume_from", default=None, help="path to checkpoint to resume from", type=str, required=True)
+ parser.add_argument("--render_video_postfix", type=str, default=None, help="an optional postfix for video")
+ parser.add_argument("--save_catted_videos", type=bool, default=False, help="visualize lidar on image")
+ parser.add_argument("--novel_cam_ids", type=str, default=None,
+ help="Novel view cameras to render: 'all' or comma-separated ids (e.g., '0,1,2'). Default: dataset default (front cam).")
+ parser.add_argument("--include_depth", action="store_true",
+ help="Include depth renders in evaluation videos (adds 'depths' to render keys).")
+ parser.add_argument("--save_raw_depth", type=str, default=None,
+ help="Directory to save raw depth as .npy (meters). Relative paths are under log_dir.")
+
+ # viewer
+ parser.add_argument("--enable_viewer", action="store_true", help="enable viewer")
+ parser.add_argument("--viewer_port", type=int, default=8080, help="viewer port")
+
+ # misc
+ parser.add_argument("opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER)
+
+ args = parser.parse_args()
+ main(args)
diff --git a/worldbench/third_party/drivestudio/tools/eval_depth_diff.py b/worldbench/third_party/drivestudio/tools/eval_depth_diff.py
new file mode 100644
index 0000000..8113a70
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/eval_depth_diff.py
@@ -0,0 +1,831 @@
+#!/usr/bin/env python3
+"""Compute depth reconstruction error under semantic masks for nuScenes runs.
+
+For each method (e.g. gt, dreamforge, drivedreamer2, magicdrive), sequence, and
+rendered frame we compare the method's depth against the GT method's depth under
+road/vehicle masks stored in `sam_mask`.
+
+Assumptions:
+ - Method work directories follow `work_dirs/drivestudio-nus-{method}/{seq}`.
+ - Depth files are stored at `.../raw_depth/full_set_{step}/_CAM_*.npy`.
+ - GT depths are in the same structure with `{method}` = `gt` (default).
+ - Semantic masks reside at
+ data/nuscenes_trainval/processed_{method}/advanced_12Hz_trainval/{seq}/sam_mask/
+ falling back to `processed_{gt_method}` or `processed` if missing.
+ - Mask file naming convention: `{frame:03d}_{camera_id}.png`.
+
+The script outputs a CSV with per-image MAE/MSE under the mask, plus summary
+statistics per method/sequence.
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import os
+from pathlib import Path
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple
+
+import glob
+import math
+
+import numpy as np
+import imageio.v2 as imageio
+
+LOG_EPS = 1e-6
+
+try:
+ from tqdm import tqdm # type: ignore
+
+ def progress_iter(iterable, **kwargs):
+ return tqdm(iterable, **kwargs)
+
+except ImportError: # pragma: no cover - fallback when tqdm is unavailable
+
+ def progress_iter(iterable, **kwargs):
+ return iterable
+
+
+# Camera ordering for nuScenes 6-camera configuration
+CAMERA_NAME_TO_ID: Dict[str, int] = {
+ "CAM_FRONT": 0,
+ "CAM_FRONT_LEFT": 1,
+ "CAM_FRONT_RIGHT": 2,
+ "CAM_BACK_LEFT": 3,
+ "CAM_BACK_RIGHT": 4,
+ "CAM_BACK": 5,
+}
+NUM_CAMERAS = len(CAMERA_NAME_TO_ID)
+
+
+def parse_args() -> argparse.Namespace:
+ p = argparse.ArgumentParser("Evaluate depth differences under road/vehicle masks")
+ p.add_argument("--methods", nargs="*", default=None,
+ help="List of methods to compare against GT. Default: auto-discover all methods under work_dirs.")
+ p.add_argument("--gt-method", default="gt", help="Method name used as GT reference (default: gt)")
+ p.add_argument("--work-root", type=Path, default=Path("work_dirs"), help="Root of drivestudio run directories")
+ p.add_argument("--depth-subdir", default="raw_depth/full_set_{step}",
+ help="Relative depth subdirectory pattern (default: raw_depth/full_set_{step})")
+ p.add_argument("--step", type=int, default=30000, help="Training step used in depth filenames (default: 30000)")
+ p.add_argument("--mask-root-template", default="data/nuscenes_trainval/processed_{method}/advanced_12Hz_trainval",
+ help="Mask root template with {method} placeholder (default matches nuScenes)")
+ p.add_argument("--mask-method", default=None,
+ help="Force using masks from this method (overrides per-method masks). Use 'gt' to always use GT masks")
+ p.add_argument("--cameras", default="CAM_FRONT",
+ help="Comma-separated camera names to evaluate (e.g. CAM_FRONT,CAM_BACK). Use 'all' to include every camera. Default: CAM_FRONT")
+ p.add_argument("--table-csv-prefix", type=Path, default=None,
+ help="If set, save summary tables to CSV. Pass a directory (e.g. results/depth_tables) or a base file path.")
+ p.add_argument("--output", type=Path, default=Path("depth_mask_metrics.csv"),
+ help="CSV file to write per-image metrics (default: depth_mask_metrics.csv)")
+ p.add_argument("--summary", type=Path, default=Path("depth_mask_metrics_summary.csv"),
+ help="CSV file to write aggregated summaries (default: depth_mask_metrics_summary.csv)")
+ p.add_argument("--save-csv", action="store_true",
+ help="Explicitly save CSV outputs (--output/--summary). Default: do not write CSV files.")
+ p.add_argument("--verbose", action="store_true", help="Print progress information")
+ return p.parse_args()
+
+
+def depth_dir(work_root: Path, method: str, seq: str, depth_subdir: str, step: int) -> Path:
+ sub = depth_subdir.format(step=step)
+ return work_root / f"drivestudio-nus-{method}" / seq / sub
+
+
+def list_sequences(work_root: Path, method: str, depth_subdir: str, step: int) -> List[str]:
+ method_root = work_root / f"drivestudio-nus-{method}"
+ if not method_root.is_dir():
+ return []
+ seqs = []
+ for seq_dir in sorted(d for d in method_root.iterdir() if d.is_dir()):
+ if depth_dir(work_root, method, seq_dir.name, depth_subdir, step).is_dir():
+ seqs.append(seq_dir.name)
+ return seqs
+
+
+def resolve_mask_dir(mask_root_template: str, method: str, seq: str, fallback_methods: Sequence[str], force_method: Optional[str]) -> Optional[Tuple[Path, str]]:
+ candidates = []
+ if force_method:
+ candidates.append(force_method)
+ for token in [method, *fallback_methods]:
+ if token not in candidates:
+ candidates.append(token)
+ for token in candidates:
+ template = mask_root_template.format(method=token)
+ if token == "":
+ template = template.replace("processed_", "processed", 1)
+ template = template.replace("processed_//", "processed/")
+ mask_dir = Path(template) / seq / "sam_mask"
+ if mask_dir.is_dir():
+ return mask_dir, token or "processed"
+ return None
+
+
+def load_mask(mask_path: Path) -> np.ndarray:
+ mask_img = imageio.imread(mask_path)
+ if mask_img.ndim == 3:
+ mask_img = mask_img[..., 0]
+ return mask_img > 0
+
+
+def resize_mask_nearest(mask: np.ndarray, target_shape: Tuple[int, int]) -> np.ndarray:
+ """Resize a boolean mask to target shape using nearest-neighbor sampling."""
+ if mask.shape == target_shape:
+ return mask
+ src_h, src_w = mask.shape
+ tgt_h, tgt_w = target_shape
+ if tgt_h <= 0 or tgt_w <= 0:
+ raise ValueError(f"Invalid target shape {target_shape} when resizing mask.")
+ y_idx = np.linspace(0, src_h - 1, tgt_h).round().astype(int)
+ x_idx = np.linspace(0, src_w - 1, tgt_w).round().astype(int)
+ resized = mask[y_idx][:, x_idx]
+ return resized
+
+
+def camera_id_from_name(cam_name: str) -> Optional[int]:
+ return CAMERA_NAME_TO_ID.get(cam_name)
+
+
+def normalize_camera_name(name: str) -> Optional[str]:
+ if not name:
+ return None
+ name = name.strip()
+ if not name:
+ return None
+ name = name.upper()
+ if not name.startswith("CAM_"):
+ name = "CAM_" + name
+ return name if name in CAMERA_NAME_TO_ID else None
+
+
+def frame_index_from_depth_name(idx_str: str) -> int:
+ return int(idx_str) // NUM_CAMERAS
+
+
+def compute_metrics(diff: np.ndarray) -> Tuple[float, float]:
+ mae = float(np.mean(np.abs(diff)))
+ mse = float(np.mean(diff ** 2))
+ return mae, mse
+
+
+def collect_depth_files(depth_directory: Path) -> List[Path]:
+ return sorted(depth_directory.glob("*_CAM_*.npy"))
+
+
+def discover_methods(work_root: Path) -> List[str]:
+ methods = []
+ for path in glob.glob(str(work_root / "drivestudio-nus-*")):
+ name = Path(path).name
+ prefix = "drivestudio-nus-"
+ if name.startswith(prefix):
+ methods.append(name[len(prefix):])
+ return sorted(set(methods))
+
+
+def format_table(header: Sequence[str], rows: Sequence[Tuple[str, Sequence[str]]]) -> str:
+ widths = [len(col) for col in header]
+ for _, cells in rows:
+ for idx, cell in enumerate(cells):
+ if idx + 1 >= len(widths):
+ widths.append(len(cell))
+ else:
+ widths[idx + 1] = max(widths[idx + 1], len(cell))
+ line_parts = ["{:<{w}}".format(header[0], w=widths[0])]
+ for idx, col in enumerate(header[1:], start=1):
+ line_parts.append("{:>{w}}".format(col, w=widths[idx]))
+ output = [" ".join(line_parts)]
+ output.append(" ".join("-" * w for w in widths))
+ for label, cells in rows:
+ row_parts = ["{:<{w}}".format(label, w=widths[0])]
+ for idx, cell in enumerate(cells, start=1):
+ row_parts.append("{:>{w}}".format(cell, w=widths[idx]))
+ output.append(" ".join(row_parts))
+ return "\n".join(output)
+
+
+def make_table_csv_path(base: Path, kind: str) -> Path:
+ if base.suffix:
+ parent = base.parent if str(base.parent) != "" else Path(".")
+ return parent / f"{base.stem}_{kind}{base.suffix}"
+ directory = base if str(base) != "" else Path(".")
+ return directory / f"{kind}.csv"
+
+
+def write_table_csv(path: Path, header: Sequence[str], rows: Sequence[Tuple[str, Sequence[str]]]) -> None:
+ path.parent.mkdir(parents=True, exist_ok=True)
+ with path.open("w", newline="") as f:
+ writer = csv.writer(f)
+ writer.writerow(header)
+ for label, cells in rows:
+ writer.writerow([label] + list(cells))
+
+
+def main() -> None:
+ args = parse_args()
+ work_root: Path = args.work_root
+ methods = list(args.methods) if args.methods else discover_methods(work_root)
+ if not methods:
+ raise SystemExit(f"No methods found under {work_root} (looking for drivestudio-nus-*)")
+ gt_method = args.gt_method
+ if gt_method in methods:
+ methods = [m for m in methods if m != gt_method]
+ if args.verbose and args.methods is None:
+ print(f"Discovered methods: {', '.join(methods)}")
+
+ if args.cameras.lower() == "all":
+ allowed_cam_names = set(CAMERA_NAME_TO_ID.keys())
+ else:
+ allowed_cam_names = set()
+ for token in args.cameras.split(','):
+ norm = normalize_camera_name(token)
+ if norm is None:
+ if args.verbose:
+ print(f"[WARN] Unknown camera '{token.strip()}', ignoring")
+ continue
+ allowed_cam_names.add(norm)
+ if not allowed_cam_names:
+ allowed_cam_names = {"CAM_FRONT"}
+ allowed_cam_ids = {CAMERA_NAME_TO_ID[name] for name in allowed_cam_names}
+ if args.verbose:
+ cams_str = ', '.join(sorted(allowed_cam_names))
+ print(f"Evaluating cameras: {cams_str}")
+
+ depth_subdir = args.depth_subdir
+ fallback_mask_methods = (gt_method, "gt", "magicdrive", "")
+
+ gt_sequences = set(list_sequences(work_root, gt_method, depth_subdir, args.step))
+ if not gt_sequences:
+ raise SystemExit(f"No GT sequences found under {work_root}/drivestudio-nus-{gt_method}")
+
+ rows: List[Dict[str, object]] = []
+ summary: Dict[Tuple[str, str], Dict[str, float]] = {}
+
+ method_iter = progress_iter(methods, desc="Methods", total=len(methods))
+ for method in method_iter:
+ method_sequences = set(list_sequences(work_root, method, depth_subdir, args.step))
+ common_seqs = sorted(gt_sequences & method_sequences)
+ if args.verbose:
+ print(f"Method {method}: {len(common_seqs)} common sequences with GT")
+ seq_iter = progress_iter(common_seqs, desc=f"{method} seqs", leave=False, total=len(common_seqs))
+ for seq in seq_iter:
+ depth_dir_method = depth_dir(work_root, method, seq, depth_subdir, args.step)
+ depth_dir_gt = depth_dir(work_root, gt_method, seq, depth_subdir, args.step)
+ mask_info = resolve_mask_dir(args.mask_root_template, method, seq, fallback_mask_methods, args.mask_method)
+ if mask_info is None:
+ print(f"[WARN] No sam_mask found for method={method}, seq={seq}; skipping")
+ continue
+ mask_dir, mask_source = mask_info
+ depth_files = collect_depth_files(depth_dir_method)
+ if not depth_files:
+ print(f"[WARN] No depth files in {depth_dir_method}; skipping")
+ continue
+ depth_iter = progress_iter(depth_files, desc=f"{method}/{seq}", leave=False, total=len(depth_files))
+ for depth_path in depth_iter:
+ stem = depth_path.stem # e.g. 00036_CAM_FRONT
+ try:
+ idx_str, cam_name = stem.split("_CAM_")
+ cam_name = "CAM_" + cam_name if not cam_name.startswith("CAM_") else cam_name
+ except ValueError:
+ print(f"[WARN] Unexpected depth filename format: {depth_path.name}; skipping")
+ continue
+ cam_id = camera_id_from_name(cam_name)
+ if cam_id is None:
+ continue # skip cameras without mask mapping
+
+ if cam_name not in allowed_cam_names or cam_id not in allowed_cam_ids:
+ continue
+
+ frame_idx = frame_index_from_depth_name(idx_str)
+ mask_path = mask_dir / f"{frame_idx:03d}_{cam_id}.png"
+ if not mask_path.is_file():
+ print(f"[WARN] Missing mask {mask_path}; skipping frame")
+ continue
+
+ gt_path = depth_dir_gt / depth_path.name
+ if not gt_path.is_file():
+ print(f"[WARN] Missing GT depth {gt_path}; skipping frame")
+ continue
+
+ depth_pred = np.load(depth_path)
+ depth_gt = np.load(gt_path)
+ if depth_pred.shape != depth_gt.shape:
+ print(f"[WARN] Shape mismatch pred {depth_pred.shape} vs gt {depth_gt.shape} for {depth_path}")
+ continue
+
+ mask = load_mask(mask_path)
+ if mask.shape != depth_pred.shape:
+ if args.verbose:
+ print(f"[INFO] Resizing mask {mask.shape} -> {depth_pred.shape} for {mask_path}")
+ try:
+ mask = resize_mask_nearest(mask, depth_pred.shape)
+ except ValueError as exc:
+ print(f"[WARN] Failed to resize mask {mask_path}: {exc}; skipping frame")
+ continue
+ if not mask.any():
+ bucket = summary.setdefault(
+ (method, seq),
+ {
+ "abs_sum": 0.0,
+ "sq_sum": 0.0,
+ "pixels": 0,
+ "frames": 0,
+ "abs_rel_sum": 0.0,
+ "sq_rel_sum": 0.0,
+ "log_sq_sum": 0.0,
+ "log_pixels": 0,
+ "rel_pixels": 0,
+ "delta1_count": 0,
+ "delta2_count": 0,
+ "delta3_count": 0,
+ },
+ )
+ bucket["frames"] += 1
+ rows.append({
+ "method": method,
+ "sequence": seq,
+ "frame_index": frame_idx,
+ "camera": cam_name,
+ "depth_file": depth_path.name,
+ "mask_file": mask_path.name,
+ "mask_source": mask_source,
+ "pixels": 0,
+ "abs_sum": float("nan"),
+ "sq_sum": float("nan"),
+ "mae": float("nan"),
+ "mse": float("nan"),
+ "rmse": float("nan"),
+ "abs_rel": float("nan"),
+ "sq_rel": float("nan"),
+ "delta1": float("nan"),
+ "delta2": float("nan"),
+ "delta3": float("nan"),
+ "log_rmse": float("nan"),
+ })
+ continue
+
+ diff = depth_pred[mask] - depth_gt[mask]
+ abs_sum = float(np.abs(diff).sum())
+ sq_sum = float((diff ** 2).sum())
+ pixels = diff.size
+ if pixels == 0:
+ continue
+ mae = abs_sum / pixels
+ mse = sq_sum / pixels
+ rmse = math.sqrt(mse)
+
+ gt_vals = depth_gt[mask]
+ pred_vals = depth_pred[mask]
+ rel_mask = gt_vals > 1e-6
+ rel_pixels = int(rel_mask.sum())
+ abs_rel_sum = 0.0
+ sq_rel_sum = 0.0
+ delta1_cnt = delta2_cnt = delta3_cnt = 0
+ abs_rel = float("nan")
+ sq_rel = float("nan")
+ delta1 = float("nan")
+ delta2 = float("nan")
+ delta3 = float("nan")
+ log_rmse = float("nan")
+ if rel_pixels > 0:
+ gt_rel = gt_vals[rel_mask]
+ pred_rel = pred_vals[rel_mask]
+ rel_errors = np.abs(pred_rel - gt_rel) / gt_rel
+ abs_rel_sum = float(rel_errors.sum())
+ sq_rel_sum = float(((pred_rel - gt_rel) ** 2 / gt_rel).sum())
+ abs_rel = abs_rel_sum / rel_pixels
+ sq_rel = sq_rel_sum / rel_pixels
+ pred_safe = np.maximum(pred_rel, 1e-6)
+ ratios = np.maximum(pred_safe / gt_rel, gt_rel / pred_safe)
+ delta1_cnt = int((ratios < 1.25).sum())
+ delta2_cnt = int((ratios < 1.25 ** 2).sum())
+ delta3_cnt = int((ratios < 1.25 ** 3).sum())
+ delta1 = delta1_cnt / rel_pixels
+ delta2 = delta2_cnt / rel_pixels
+ delta3 = delta3_cnt / rel_pixels
+
+ bucket = summary.setdefault(
+ (method, seq),
+ {
+ "abs_sum": 0.0,
+ "sq_sum": 0.0,
+ "pixels": 0,
+ "frames": 0,
+ "abs_rel_sum": 0.0,
+ "sq_rel_sum": 0.0,
+ "log_sq_sum": 0.0,
+ "log_pixels": 0,
+ "rel_pixels": 0,
+ "delta1_count": 0,
+ "delta2_count": 0,
+ "delta3_count": 0,
+ },
+ )
+ bucket["abs_sum"] += abs_sum
+ bucket["sq_sum"] += sq_sum
+ bucket["pixels"] += pixels
+ bucket["frames"] += 1
+ if rel_pixels > 0:
+ bucket["abs_rel_sum"] += abs_rel_sum
+ bucket["sq_rel_sum"] += sq_rel_sum
+ bucket["rel_pixels"] += rel_pixels
+ bucket["delta1_count"] += delta1_cnt
+ bucket["delta2_count"] += delta2_cnt
+ bucket["delta3_count"] += delta3_cnt
+ diff_log = (np.log(np.maximum(pred_rel, LOG_EPS)) - np.log(np.maximum(gt_rel, LOG_EPS)))**2
+ log_rmse = math.sqrt(float(diff_log.mean())) if diff_log.size > 0 else float("nan")
+ bucket["log_sq_sum"] = bucket.get("log_sq_sum", 0.0) + float(diff_log.sum())
+ bucket["log_pixels"] = bucket.get("log_pixels", 0) + int(diff_log.size)
+
+ rows.append({
+ "method": method,
+ "sequence": seq,
+ "frame_index": frame_idx,
+ "camera": cam_name,
+ "depth_file": depth_path.name,
+ "mask_file": mask_path.name,
+ "mask_source": mask_source,
+ "pixels": pixels,
+ "abs_sum": abs_sum,
+ "sq_sum": sq_sum,
+ "mae": mae,
+ "mse": mse,
+ "rmse": rmse,
+ "abs_rel": abs_rel,
+ "sq_rel": sq_rel,
+ "delta1": delta1,
+ "delta2": delta2,
+ "delta3": delta3,
+ "log_rmse": log_rmse,
+ })
+
+ if not rows:
+ print("No metrics computed; check paths and inputs.")
+ return
+
+ # Write per-image metrics
+ fieldnames = [
+ "method",
+ "sequence",
+ "frame_index",
+ "camera",
+ "depth_file",
+ "mask_file",
+ "mask_source",
+ "pixels",
+ "abs_sum",
+ "sq_sum",
+ "mae",
+ "mse",
+ "rmse",
+ "log_rmse",
+ "abs_rel",
+ "sq_rel",
+ "delta1",
+ "delta2",
+ "delta3",
+ ]
+ if args.save_csv:
+ args.output.parent.mkdir(parents=True, exist_ok=True)
+ with args.output.open("w", newline="") as f:
+ writer = csv.DictWriter(f, fieldnames=fieldnames)
+ writer.writeheader()
+ for row in rows:
+ writer.writerow(row)
+
+ summary_fieldnames = [
+ "method",
+ "sequence",
+ "frames",
+ "mean_mae",
+ "mean_mse",
+ "mean_rmse",
+ "mean_log_rmse",
+ "mean_abs_rel",
+ "mean_sq_rel",
+ "delta1",
+ "delta2",
+ "delta3",
+ "total_pixels",
+ "rel_pixels",
+ "log_pixels",
+ ]
+ if args.save_csv:
+ with args.summary.open("w", newline="") as f:
+ writer = csv.DictWriter(f, fieldnames=summary_fieldnames)
+ writer.writeheader()
+ for (method, seq), stats in sorted(summary.items()):
+ frames = stats["frames"]
+ pixels = stats["pixels"]
+ rel_pixels = stats["rel_pixels"]
+ mean_mae = stats["abs_sum"] / pixels if pixels else float("nan")
+ mean_mse = stats["sq_sum"] / pixels if pixels else float("nan")
+ mean_rmse = math.sqrt(mean_mse) if math.isfinite(mean_mse) else float("nan")
+ log_pixels = stats["log_pixels"]
+ mean_log_rmse = math.sqrt(stats["log_sq_sum"] / log_pixels) if log_pixels else float("nan")
+ mean_abs_rel = stats["abs_rel_sum"] / rel_pixels if rel_pixels else float("nan")
+ mean_sq_rel = stats["sq_rel_sum"] / rel_pixels if rel_pixels else float("nan")
+ delta1 = stats["delta1_count"] / rel_pixels if rel_pixels else float("nan")
+ delta2 = stats["delta2_count"] / rel_pixels if rel_pixels else float("nan")
+ delta3 = stats["delta3_count"] / rel_pixels if rel_pixels else float("nan")
+ writer.writerow({
+ "method": method,
+ "sequence": seq,
+ "frames": frames,
+ "mean_mae": mean_mae,
+ "mean_mse": mean_mse,
+ "mean_rmse": mean_rmse,
+ "mean_log_rmse": mean_log_rmse,
+ "mean_abs_rel": mean_abs_rel,
+ "mean_sq_rel": mean_sq_rel,
+ "delta1": delta1,
+ "delta2": delta2,
+ "delta3": delta3,
+ "total_pixels": pixels,
+ "rel_pixels": rel_pixels,
+ "log_pixels": log_pixels,
+ })
+
+ if args.verbose and args.save_csv:
+ print(f"Wrote per-image metrics to {args.output}")
+ print(f"Wrote summaries to {args.summary}")
+
+
+ def weighted_mean(values: Sequence[float], weights: Sequence[int]) -> float:
+ if not values or not weights or len(values) != len(weights):
+ return float("nan")
+ acc = 0.0
+ total = 0
+ for v, w in zip(values, weights):
+ if not math.isfinite(v) or w <= 0:
+ continue
+ acc += v * w
+ total += w
+ return acc / total if total else float("nan")
+
+ def weighted_std(values: Sequence[float], weights: Sequence[int], mean_val: Optional[float] = None) -> float:
+ if not values or not weights or len(values) != len(weights):
+ return float("nan")
+ if mean_val is None:
+ mean_val = weighted_mean(values, weights)
+ if not math.isfinite(mean_val):
+ return float("nan")
+ acc = 0.0
+ total = 0
+ for v, w in zip(values, weights):
+ if not math.isfinite(v) or w <= 0:
+ continue
+ acc += w * ((v - mean_val) ** 2)
+ total += w
+ return math.sqrt(acc / total) if total else float("nan")
+
+# Console tables
+ def fmt_float(val: float) -> str:
+ return f"{val:.4f}" if math.isfinite(val) else "nan"
+
+ def fmt_mean_std(mean_val: float, std_val: float) -> str:
+ if not math.isfinite(mean_val):
+ return "nan"
+ if math.isfinite(std_val):
+ return f"{mean_val:.4f}±{std_val:.4f}"
+ return f"{mean_val:.4f}"
+
+ per_seq_rows: List[Tuple[str, Sequence[str]]] = []
+ aggregated: Dict[str, Dict[str, float]] = {}
+ method_metrics: Dict[str, Dict[str, float]] = {}
+
+ seq_mae: Dict[str, List[float]] = {}
+ seq_rmse: Dict[str, List[float]] = {}
+ seq_log_rmse: Dict[str, List[float]] = {}
+ seq_abs_rel: Dict[str, List[float]] = {}
+ seq_sq_rel: Dict[str, List[float]] = {}
+ seq_delta1: Dict[str, List[float]] = {}
+ seq_delta2: Dict[str, List[float]] = {}
+ seq_delta3: Dict[str, List[float]] = {}
+
+ pixel_weights: Dict[str, List[int]] = {}
+ rel_weights: Dict[str, List[int]] = {}
+ log_weights: Dict[str, List[int]] = {}
+
+ for (method, seq), stats in sorted(summary.items()):
+ frames = stats["frames"]
+ pixels = stats["pixels"]
+ rel_pixels = stats["rel_pixels"]
+ mae_seq = stats["abs_sum"] / pixels if pixels else float("nan")
+ mse_seq = stats["sq_sum"] / pixels if pixels else float("nan")
+ rmse_seq = math.sqrt(mse_seq) if math.isfinite(mse_seq) else float("nan")
+ log_rmse_seq = math.sqrt(stats["log_sq_sum"] / stats["log_pixels"]) if stats["log_pixels"] else float("nan")
+ abs_rel_seq = stats["abs_rel_sum"] / rel_pixels if rel_pixels else float("nan")
+ sq_rel_seq = stats["sq_rel_sum"] / rel_pixels if rel_pixels else float("nan")
+ delta1_seq = stats["delta1_count"] / rel_pixels if rel_pixels else float("nan")
+ delta2_seq = stats["delta2_count"] / rel_pixels if rel_pixels else float("nan")
+ delta3_seq = stats["delta3_count"] / rel_pixels if rel_pixels else float("nan")
+
+ per_seq_rows.append((
+ f"{method}/{seq}",
+ [
+ str(frames),
+ str(pixels),
+ fmt_float(mae_seq),
+ fmt_float(rmse_seq),
+ fmt_float(log_rmse_seq),
+ fmt_float(abs_rel_seq),
+ fmt_float(sq_rel_seq),
+ fmt_float(delta1_seq),
+ fmt_float(delta2_seq),
+ fmt_float(delta3_seq),
+ ],
+ ))
+
+ bucket = aggregated.setdefault(
+ method,
+ {
+ "abs_sum": 0.0,
+ "sq_sum": 0.0,
+ "pixels": 0,
+ "frames": 0,
+ "abs_rel_sum": 0.0,
+ "sq_rel_sum": 0.0,
+ "log_sq_sum": 0.0,
+ "log_pixels": 0,
+ "rel_pixels": 0,
+ "delta1_count": 0,
+ "delta2_count": 0,
+ "delta3_count": 0,
+ },
+ )
+ bucket["abs_sum"] += stats["abs_sum"]
+ bucket["sq_sum"] += stats["sq_sum"]
+ bucket["pixels"] += pixels
+ bucket["frames"] += frames
+ bucket["abs_rel_sum"] += stats["abs_rel_sum"]
+ bucket["sq_rel_sum"] += stats["sq_rel_sum"]
+ bucket["log_sq_sum"] += stats["log_sq_sum"]
+ bucket["log_pixels"] += stats["log_pixels"]
+ bucket["rel_pixels"] += rel_pixels
+ bucket["delta1_count"] += stats["delta1_count"]
+ bucket["delta2_count"] += stats["delta2_count"]
+ bucket["delta3_count"] += stats["delta3_count"]
+
+ seq_mae.setdefault(method, []).append(mae_seq)
+ seq_rmse.setdefault(method, []).append(rmse_seq)
+ seq_log_rmse.setdefault(method, []).append(log_rmse_seq)
+ seq_abs_rel.setdefault(method, []).append(abs_rel_seq)
+ seq_sq_rel.setdefault(method, []).append(sq_rel_seq)
+ seq_delta1.setdefault(method, []).append(delta1_seq)
+ seq_delta2.setdefault(method, []).append(delta2_seq)
+ seq_delta3.setdefault(method, []).append(delta3_seq)
+
+ pixel_weights.setdefault(method, []).append(pixels)
+ rel_weights.setdefault(method, []).append(rel_pixels)
+ log_weights.setdefault(method, []).append(stats["log_pixels"])
+
+ for method in sorted(aggregated):
+ stats = aggregated[method]
+ pixels = stats["pixels"]
+ rel_pixels = stats["rel_pixels"]
+ global_mae = stats["abs_sum"] / pixels if pixels else float("nan")
+ global_mse = stats["sq_sum"] / pixels if pixels else float("nan")
+ global_rmse = math.sqrt(global_mse) if math.isfinite(global_mse) else float("nan")
+ global_log_rmse = math.sqrt(stats["log_sq_sum"] / stats["log_pixels"]) if stats["log_pixels"] else float("nan")
+ global_abs_rel = stats["abs_rel_sum"] / rel_pixels if rel_pixels else float("nan")
+ global_sq_rel = stats["sq_rel_sum"] / rel_pixels if rel_pixels else float("nan")
+ global_delta1 = stats["delta1_count"] / rel_pixels if rel_pixels else float("nan")
+ global_delta2 = stats["delta2_count"] / rel_pixels if rel_pixels else float("nan")
+ global_delta3 = stats["delta3_count"] / rel_pixels if rel_pixels else float("nan")
+
+ mae_list = seq_mae.get(method, [])
+ rmse_list = seq_rmse.get(method, [])
+ log_list = seq_log_rmse.get(method, [])
+ abs_rel_list = seq_abs_rel.get(method, [])
+ sq_rel_list = seq_sq_rel.get(method, [])
+ delta1_list = seq_delta1.get(method, [])
+ delta2_list = seq_delta2.get(method, [])
+ delta3_list = seq_delta3.get(method, [])
+
+ pixel_w = pixel_weights.get(method, [])
+ rel_w = rel_weights.get(method, [])
+ log_w = log_weights.get(method, [])
+
+ seq_mae_mean = weighted_mean(mae_list, pixel_w)
+ seq_mae_std = weighted_std(mae_list, pixel_w, seq_mae_mean)
+ seq_rmse_mean = weighted_mean(rmse_list, pixel_w)
+ seq_rmse_std = weighted_std(rmse_list, pixel_w, seq_rmse_mean)
+ seq_log_rmse_mean = weighted_mean(log_list, log_w)
+ seq_log_rmse_std = weighted_std(log_list, log_w, seq_log_rmse_mean)
+ seq_abs_rel_mean = weighted_mean(abs_rel_list, rel_w)
+ seq_abs_rel_std = weighted_std(abs_rel_list, rel_w, seq_abs_rel_mean)
+ seq_sq_rel_mean = weighted_mean(sq_rel_list, rel_w)
+ seq_sq_rel_std = weighted_std(sq_rel_list, rel_w, seq_sq_rel_mean)
+ seq_delta1_mean = weighted_mean(delta1_list, rel_w)
+ seq_delta1_std = weighted_std(delta1_list, rel_w, seq_delta1_mean)
+ seq_delta2_mean = weighted_mean(delta2_list, rel_w)
+ seq_delta2_std = weighted_std(delta2_list, rel_w, seq_delta2_mean)
+ seq_delta3_mean = weighted_mean(delta3_list, rel_w)
+ seq_delta3_std = weighted_std(delta3_list, rel_w, seq_delta3_mean)
+
+ method_metrics[method] = {
+ "frames": stats["frames"],
+ "pixels": pixels,
+ "global_mae": global_mae,
+ "global_mse": global_mse,
+ "global_rmse": global_rmse,
+ "global_log_rmse": global_log_rmse,
+ "global_abs_rel": global_abs_rel,
+ "global_sq_rel": global_sq_rel,
+ "global_delta1": global_delta1,
+ "global_delta2": global_delta2,
+ "global_delta3": global_delta3,
+ "seq_mae_mean": seq_mae_mean,
+ "seq_mae_std": seq_mae_std,
+ "seq_rmse_mean": seq_rmse_mean,
+ "seq_rmse_std": seq_rmse_std,
+ "seq_log_rmse_mean": seq_log_rmse_mean,
+ "seq_log_rmse_std": seq_log_rmse_std,
+ "seq_abs_rel_mean": seq_abs_rel_mean,
+ "seq_abs_rel_std": seq_abs_rel_std,
+ "seq_sq_rel_mean": seq_sq_rel_mean,
+ "seq_sq_rel_std": seq_sq_rel_std,
+ "seq_delta1_mean": seq_delta1_mean,
+ "seq_delta1_std": seq_delta1_std,
+ "seq_delta2_mean": seq_delta2_mean,
+ "seq_delta2_std": seq_delta2_std,
+ "seq_delta3_mean": seq_delta3_mean,
+ "seq_delta3_std": seq_delta3_std,
+ }
+
+ if method_metrics:
+ method_order = sorted(method_metrics)
+ metric_specs = [
+ ("Frames", "frames", True),
+ ("Pixels", "pixels", True),
+ ("Global MAE", "global_mae", False),
+ ("Global MSE", "global_mse", False),
+ ("Global RMSE", "global_rmse", False),
+ ("Global Log RMSE", "global_log_rmse", False),
+ ("Global AbsRel", "global_abs_rel", False),
+ ("Global SqRel", "global_sq_rel", False),
+ ("Global δ1", "global_delta1", False),
+ ("Global δ2", "global_delta2", False),
+ ("Global δ3", "global_delta3", False),
+ ("Seq MAE", ("seq_mae_mean", "seq_mae_std"), False),
+ ("Seq RMSE", ("seq_rmse_mean", "seq_rmse_std"), False),
+ ("Seq Log RMSE", ("seq_log_rmse_mean", "seq_log_rmse_std"), False),
+ ("Seq AbsRel", ("seq_abs_rel_mean", "seq_abs_rel_std"), False),
+ ("Seq SqRel", ("seq_sq_rel_mean", "seq_sq_rel_std"), False),
+ ("Seq δ1", ("seq_delta1_mean", "seq_delta1_std"), False),
+ ("Seq δ2", ("seq_delta2_mean", "seq_delta2_std"), False),
+ ("Seq δ3", ("seq_delta3_mean", "seq_delta3_std"), False),
+ ]
+ method_rows: List[Tuple[str, Sequence[str]]] = []
+ for label, key, is_int in metric_specs:
+ cells = []
+ if isinstance(key, tuple):
+ mean_key, std_key = key
+ for method in method_order:
+ data = method_metrics[method]
+ mean_val = data.get(mean_key)
+ std_val = data.get(std_key)
+ cells.append(fmt_mean_std(mean_val, std_val))
+ else:
+ for method in method_order:
+ value = method_metrics[method].get(key)
+ if value is None:
+ cells.append("n/a")
+ else:
+ cells.append(str(int(value)) if is_int and math.isfinite(value) else fmt_float(value))
+ method_rows.append((label, cells))
+ method_header = ["Metric"] + method_order
+ print("Method depth metrics (mean ± std where applicable):")
+ print(format_table(method_header, method_rows))
+ print()
+ if args.table_csv_prefix:
+ write_table_csv(
+ make_table_csv_path(args.table_csv_prefix, "method"),
+ method_header,
+ method_rows,
+ )
+
+ if per_seq_rows:
+ seq_header = [
+ "Method/Seq",
+ "Frames",
+ "Pixels",
+ "MAE",
+ "RMSE",
+ "Log RMSE",
+ "AbsRel",
+ "SqRel",
+ "δ1",
+ "δ2",
+ "δ3",
+ ]
+ print("Per-sequence depth error under mask:")
+ print(format_table(seq_header, per_seq_rows))
+ if args.table_csv_prefix:
+ write_table_csv(
+ make_table_csv_path(args.table_csv_prefix, "per_seq"),
+ seq_header,
+ per_seq_rows,
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/third_party/drivestudio/tools/export_front_depth_nus.py b/worldbench/third_party/drivestudio/tools/export_front_depth_nus.py
new file mode 100644
index 0000000..b5f102b
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/export_front_depth_nus.py
@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+"""Render nuScenes front-view depth videos for trained runs.
+
+The script mirrors :mod:`tools.batch_eval_nus` but only exports a single
+``full_set__depths_front.mp4`` per checkpoint by rendering CAM_FRONT
+frames from ``full_image_set``.
+
+Example:
+ python tools/export_front_depth_nus.py \\
+ --root work_dirs \\
+ --pattern "drivestudio-nus-*" \\
+ --fps 12 \\
+ --skip-existing
+"""
+
+from __future__ import annotations
+
+import argparse
+import glob
+import os
+import sys
+from typing import List, Optional
+
+import torch
+from omegaconf import OmegaConf
+
+from datasets.driving_dataset import DrivingDataset
+from models.video_utils import render_images, save_videos
+from utils.misc import import_str
+
+
+def find_checkpoint(run_dir: str) -> Optional[str]:
+ final = os.path.join(run_dir, "checkpoint_final.pth")
+ if os.path.isfile(final):
+ return final
+ cands = sorted(glob.glob(os.path.join(run_dir, "checkpoint_*.pth")))
+ return cands[-1] if cands else None
+
+
+def load_cfg(log_dir: str) -> OmegaConf:
+ cfg_path = os.path.join(log_dir, "config.yaml")
+ if not os.path.isfile(cfg_path):
+ raise FileNotFoundError(f"config.yaml not found in {log_dir}")
+ return OmegaConf.load(cfg_path)
+
+
+def build_dataset_trainer(cfg: OmegaConf, device: torch.device):
+ dataset = DrivingDataset(data_cfg=cfg.data)
+ trainer = import_str(cfg.trainer.type)(
+ **cfg.trainer,
+ num_timesteps=dataset.num_img_timesteps,
+ model_config=cfg.model,
+ num_train_images=len(dataset.train_image_set),
+ num_full_images=len(dataset.full_image_set),
+ test_set_indices=dataset.test_timesteps,
+ scene_aabb=dataset.get_aabb().reshape(2, 3),
+ device=device,
+ )
+ return dataset, trainer
+
+
+def ensure_dir(path: str) -> None:
+ os.makedirs(path, exist_ok=True)
+
+
+def select_cam_results(render_results: dict, cam_name: str) -> dict:
+ """Filter render outputs so only frames from ``cam_name`` remain."""
+
+ cam_names: List[str] = render_results.get("cam_names", [])
+ if not cam_names:
+ raise ValueError("render_results missing 'cam_names'")
+ keep = [idx for idx, name in enumerate(cam_names) if name == cam_name]
+ if not keep:
+ raise ValueError(f"No frames found for camera {cam_name}")
+
+ total = len(cam_names)
+ filtered = {}
+ for key, values in render_results.items():
+ if isinstance(values, list) and len(values) == total:
+ filtered[key] = [values[i] for i in keep]
+ else:
+ filtered[key] = values
+ filtered["cam_names"] = [cam_names[i] for i in keep]
+ return filtered
+
+
+def export_front_depth_video(
+ dataset: DrivingDataset,
+ trainer,
+ save_path: str,
+ fps: int,
+ cam_name: str,
+ skip_existing: bool,
+ verbose: bool,
+):
+ ensure_dir(os.path.dirname(save_path))
+ if skip_existing and os.path.isfile(save_path):
+ if verbose:
+ print(f"[Skip] front depth video exists: {save_path}")
+ return
+
+ results = render_images(
+ trainer=trainer,
+ dataset=dataset.full_image_set,
+ compute_metrics=False,
+ compute_error_map=False,
+ vis_indices=None,
+ save_image_pairs=False,
+ )
+ front_results = select_cam_results(results, cam_name)
+ num_frames = len(front_results["cam_names"])
+
+ save_videos(
+ front_results,
+ save_path,
+ layout=dataset.layout,
+ num_timestamps=num_frames,
+ keys=["depths"],
+ num_cams=1,
+ save_seperate_video=False,
+ fps=fps,
+ verbose=verbose,
+ )
+
+
+def main():
+ parser = argparse.ArgumentParser("Export nuScenes front depth videos")
+ parser.add_argument("--root", type=str, default="work_dirs", help="Root of work dirs")
+ parser.add_argument("--pattern", type=str, default="drivestudio-nus-*", help="Glob for method dirs")
+ parser.add_argument("--fps", type=int, default=12, help="FPS for the depth video")
+ parser.add_argument("--cam-name", type=str, default="CAM_FRONT", help="Camera name to render")
+ parser.add_argument("--skip-existing", action="store_true", help="Skip runs with existing outputs")
+ parser.add_argument("--verbose", action="store_true", help="Log progress")
+ args = parser.parse_args()
+
+ method_dirs = sorted(glob.glob(os.path.join(args.root, args.pattern)))
+ if not method_dirs:
+ print(f"No method dirs matched under {args.root} with pattern {args.pattern}")
+ sys.exit(0)
+
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+ for mdir in method_dirs:
+ for run in sorted(os.listdir(mdir)):
+ run_dir = os.path.join(mdir, run)
+ if not os.path.isdir(run_dir):
+ continue
+
+ ckpt = find_checkpoint(run_dir)
+ if ckpt is None:
+ continue
+ log_dir = os.path.dirname(ckpt)
+
+ step_hint = ckpt.split("_")[-1].split(".")[0]
+ if step_hint == "final":
+ step_hint = "30000"
+ tentative_video = os.path.join(log_dir, "videos", f"full_set_{step_hint}_depths_front.mp4")
+ if args.skip_existing and os.path.isfile(tentative_video):
+ if args.verbose:
+ print(f"[Skip] {tentative_video} already exists (checkpoint hint)")
+ continue
+
+ if args.verbose:
+ print(f"[Eval] {ckpt}")
+
+ cfg = load_cfg(log_dir)
+ dataset, trainer = build_dataset_trainer(cfg, device)
+ trainer.resume_from_checkpoint(ckpt, load_only_model=True)
+ step = trainer.step
+ video_path = os.path.join(log_dir, "videos", f"full_set_{step}_depths_front.mp4")
+
+ export_front_depth_video(
+ dataset=dataset,
+ trainer=trainer,
+ save_path=video_path,
+ fps=args.fps,
+ cam_name=args.cam_name,
+ skip_existing=args.skip_existing,
+ verbose=args.verbose,
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/third_party/drivestudio/tools/export_paper_frames.py b/worldbench/third_party/drivestudio/tools/export_paper_frames.py
new file mode 100644
index 0000000..46f398a
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/export_paper_frames.py
@@ -0,0 +1,385 @@
+#!/usr/bin/env python3
+"""
+Export per-camera frames and metrics for paper figures.
+
+For a given clip + model + render types (train/depth/novel), this script:
+
+- Renders original per-camera RGB/depth (no mp4 decode) using the same
+ dataset/trainer pipeline as tools/eval.py and tools/batch_eval_nus.py.
+- Saves images under:
+
+ data/metric_compare/paper/{render_type}/{clip_id}/{model_name}/{camera_name}/{frame_id}.png
+
+ where:
+ - render_type ∈ {train, depth, novel}
+ - camera_name:
+ - train/depth: real camera name, e.g., CAM_FRONT, CAM_FRONT_LEFT, ...
+ - novel: one of {front_center_interp, lateral_offset, lateral_offset_left}
+ - frame_id: frame_idx (0-based) for train/depth; novel trajectory frame index for novel.
+
+- Saves metrics under:
+
+ data/metric_compare/paper/train/{clip_id}/{model_name}/train_metric.json
+ data/metric_compare/paper/depth/{clip_id}/{model_name}/depth_metric.json
+
+ train_metric.json:
+ - latest metrics/*.json for this run and clip (same source as tools/metrics_analysis.py),
+ e.g., image_metrics/full/lpips, psnr, ssim, etc.
+
+ depth_metric.json:
+ - row from depth_tables_3/per_seq.csv (fallback depth_tables_2/per_seq.csv)
+ for Method/Seq='{method_label}/{clip_id}', including AbsRel, MAE, RMSE, etc.
+
+Novel currently has no metrics and is skipped for metric export.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import csv
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from PIL import Image
+
+from utils.visualization import depth_visualizer
+from datasets.driving_dataset import DrivingDataset
+from utils.misc import import_str
+from utils.camera import lateral_offset_trajectory
+
+
+def load_latest_metric_file(metrics_dir: Path) -> Optional[Path]:
+ json_files = sorted(
+ (p for p in metrics_dir.glob("*.json") if p.is_file()),
+ key=lambda p: p.stat().st_mtime,
+ )
+ return json_files[-1] if json_files else None
+
+
+def load_cfg(run_dir: Path) -> OmegaConf:
+ """Load config.yaml from a run directory (per-clip run dir)."""
+ cfg_path = run_dir / "config.yaml"
+ if not cfg_path.is_file():
+ raise SystemExit(f"config.yaml not found in {run_dir}")
+ return OmegaConf.load(cfg_path)
+
+
+def find_checkpoint(run_dir: Path) -> Optional[Path]:
+ """Find checkpoint within a run directory (per-clip run dir)."""
+ final = run_dir / "checkpoint_final.pth"
+ if final.is_file():
+ return final
+ cands = sorted(run_dir.glob("checkpoint_*.pth"))
+ return cands[-1] if cands else None
+
+
+def build_dataset_trainer(cfg: OmegaConf, device: torch.device):
+ dataset = DrivingDataset(data_cfg=cfg.data)
+ trainer = import_str(cfg.trainer.type)(
+ **cfg.trainer,
+ num_timesteps=dataset.num_img_timesteps,
+ model_config=cfg.model,
+ num_train_images=len(dataset.train_image_set),
+ num_full_images=len(dataset.full_image_set),
+ test_set_indices=dataset.test_timesteps,
+ scene_aabb=dataset.get_aabb().reshape(2, 3),
+ device=device,
+ )
+ return dataset, trainer
+
+
+def save_png_uint8(rgb: np.ndarray, path: Path) -> None:
+ """Save float RGB in [0,1] as 8-bit PNG."""
+ arr = (np.clip(rgb, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8)
+ path.parent.mkdir(parents=True, exist_ok=True)
+ Image.fromarray(arr).save(path)
+
+
+def save_depth_color(depth_m: np.ndarray, opacity: Optional[np.ndarray], path: Path) -> None:
+ """Colorize depth using depth_visualizer (same as eval videos)."""
+ d = depth_m
+ if d.ndim == 3 and d.shape[-1] == 1:
+ d = d[..., 0]
+ w = opacity
+ if w is not None and isinstance(w, np.ndarray) and w.ndim == 3 and w.shape[-1] == 1:
+ w = w[..., 0]
+ color = depth_visualizer(d, w)
+ arr = (np.clip(color, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8)
+ path.parent.mkdir(parents=True, exist_ok=True)
+ Image.fromarray(arr).save(path)
+
+
+def load_train_metrics(run_dir: Path, clip_id: str) -> Dict[str, float]:
+ """Load latest metrics JSON for a given model/clip (train metrics)."""
+ metrics_dir = run_dir / clip_id / "metrics"
+ latest = load_latest_metric_file(metrics_dir)
+ if latest is None:
+ return {}
+ try:
+ payload = json.loads(latest.read_text())
+ except Exception:
+ return {}
+ return payload
+
+
+def method_label(name: str) -> str:
+ return name.split("drivestudio-nus-", 1)[-1].replace("_", " ").replace("-", " ")
+
+
+def load_depth_metrics(
+ depth_csv_paths: List[Path],
+ model_name: str,
+ clip_id: str,
+ absrel_col: str = "AbsRel",
+) -> Dict[str, float]:
+ """Load depth metrics for a given model+clip from per_seq CSVs.
+
+ Uses Method/Seq='{method}/{clip_id}', where method is method_label(model)
+ with spaces removed (same logic as compose_depth_grid.py).
+ """
+ meth = method_label(model_name).replace(" ", "")
+ key1 = (meth, clip_id)
+ key2 = (meth.lower(), clip_id)
+ result: Dict[str, float] = {}
+ for p in depth_csv_paths:
+ if not p.is_file():
+ continue
+ try:
+ with p.open("r", newline="") as f:
+ reader = csv.DictReader(f)
+ if "Method/Seq" not in reader.fieldnames:
+ continue
+ for row in reader:
+ tag = row.get("Method/Seq", "").strip()
+ if "/" not in tag:
+ continue
+ method, seq = tag.split("/", 1)
+ key = (method.strip(), seq.strip())
+ if key not in (key1, key2):
+ continue
+ # dump all numeric columns for this entry
+ for col, val in row.items():
+ if col == "Method/Seq":
+ continue
+ try:
+ result[col] = float(val)
+ except (TypeError, ValueError):
+ continue
+ return result
+ except Exception:
+ continue
+ return result
+
+
+@torch.no_grad()
+def render_train_and_depth_frames(
+ dataset: DrivingDataset,
+ trainer,
+ clip_id: str,
+ model_name: str,
+ out_root: Path,
+ types: List[str],
+ device: torch.device,
+) -> None:
+ """Render all train/depth frames for all cameras.
+
+ Saves:
+ train/{cam_name}/{frame_idx}.png (RGB)
+ depth/{cam_name}/{frame_idx}.png (colorized depth)
+ """
+ need_train = "train" in types
+ need_depth = "depth" in types
+ if not (need_train or need_depth):
+ return
+
+ ds = dataset.full_image_set
+ camera_downscale = trainer._get_downscale_factor()
+ for idx in range(len(ds)):
+ image_infos, cam_infos = ds.get_image(idx, camera_downscale)
+ # move tensors to device
+ for k, v in list(image_infos.items()):
+ if isinstance(v, torch.Tensor):
+ image_infos[k] = v.to(device, non_blocking=True)
+ for k, v in list(cam_infos.items()):
+ if isinstance(v, torch.Tensor):
+ cam_infos[k] = v.to(device, non_blocking=True)
+
+ out = trainer(image_infos=image_infos, camera_infos=cam_infos)
+ rgb = out["rgb"].clamp(0.0, 1.0).detach().cpu().numpy()
+ depth = out.get("depth", None)
+ opacity = out.get("opacity", None)
+ depth_np = depth.detach().cpu().numpy() if isinstance(depth, torch.Tensor) else None
+ opacity_np = opacity.detach().cpu().numpy() if isinstance(opacity, torch.Tensor) else None
+
+ cam_name = cam_infos["cam_name"]
+ frame_idx = int(image_infos["frame_idx"].flatten()[0].cpu().item())
+
+ if need_train:
+ base_dir_train = out_root / "train" / clip_id / model_name
+ out_path = base_dir_train / cam_name / f"{frame_idx:04d}.png"
+ save_png_uint8(rgb, out_path)
+ if need_depth and depth_np is not None:
+ base_dir_depth = out_root / "depth" / clip_id / model_name
+ out_path = base_dir_depth / cam_name / f"{frame_idx:04d}.png"
+ save_depth_color(depth_np, opacity_np, out_path)
+
+
+@torch.no_grad()
+def render_novel_frames(
+ dataset: DrivingDataset,
+ trainer,
+ clip_id: str,
+ model_name: str,
+ out_root: Path,
+ device: torch.device,
+ traj_types: List[str] = None,
+ cam_id: int = 0,
+) -> None:
+ """Render novel views for standard trajectories.
+
+ Saves:
+ novel/{traj_type}/{frame_id}.png
+ """
+ if traj_types is None:
+ traj_types = ["front_center_interp", "lateral_offset", "lateral_offset_left"]
+
+ base_dir = out_root / "novel" / clip_id / model_name
+ for traj_type in traj_types:
+ # build trajectory
+ if traj_type == "lateral_offset_left":
+ per_cam_poses: Dict[int, torch.Tensor] = {
+ c: dataset.pixel_source.camera_data[c].cam_to_worlds
+ for c in dataset.pixel_source.camera_list
+ }
+ traj = lateral_offset_trajectory(
+ dataset_type=dataset.type,
+ per_cam_poses=per_cam_poses,
+ original_frames=dataset.frame_num,
+ target_frames=dataset.frame_num,
+ offset_distance=-1.0,
+ )
+ else:
+ trajs = dataset.get_novel_render_traj(
+ traj_types=[traj_type],
+ target_frames=dataset.frame_num,
+ )
+ if traj_type not in trajs:
+ print(f"[WARN] clip {clip_id} traj '{traj_type}' unavailable; skip")
+ continue
+ traj = trajs[traj_type]
+
+ render_list = dataset.prepare_novel_view_render_data(traj, cam_id=cam_id)
+ for frame_idx, frame_data in enumerate(render_list):
+ # move to device
+ for k, v in list(frame_data["cam_infos"].items()):
+ frame_data["cam_infos"][k] = v.to(device, non_blocking=True)
+ for k, v in list(frame_data["image_infos"].items()):
+ frame_data["image_infos"][k] = v.to(device, non_blocking=True)
+ out = trainer(
+ image_infos=frame_data["image_infos"],
+ camera_infos=frame_data["cam_infos"],
+ novel_view=True,
+ )
+ rgb = out["rgb"].detach().cpu().numpy()
+ out_path = base_dir / traj_type / f"{frame_idx:04d}.png"
+ save_png_uint8(rgb, out_path)
+
+
+def parse_args() -> argparse.Namespace:
+ p = argparse.ArgumentParser("Export per-camera frames + metrics for paper demos")
+ p.add_argument("clip_id", type=str, help="Clip ID (e.g., 813)")
+ p.add_argument("model", type=str, help="Model directory name under work_dirs (e.g., drivestudio-nus-dist4d)")
+ p.add_argument(
+ "--types",
+ type=str,
+ nargs="*",
+ default=["train", "depth", "novel"],
+ choices=["train", "depth", "novel"],
+ help="Render types to export (default: train depth novel)",
+ )
+ p.add_argument("--root", type=Path, default=Path("work_dirs"), help="Root for runs (default: work_dirs)")
+ p.add_argument("--step", type=int, default=30000, help="Training step (for depth metrics naming; default: 30000)")
+ p.add_argument("--output-root", type=Path, default=Path("data/metric_compare/paper"), help="Output root directory")
+ p.add_argument("--device", type=str, default=None, help="Device to use (e.g., cuda:0 or cpu; default: auto)")
+ p.add_argument(
+ "--depth-csv",
+ type=Path,
+ nargs="*",
+ default=None,
+ help="Depth per-seq CSVs (default: depth_tables_3/per_seq.csv, depth_tables_2/per_seq.csv)",
+ )
+ p.add_argument("--absrel-column", type=str, default="AbsRel", help="Column name for AbsRel in depth CSV")
+ return p.parse_args()
+
+
+def main() -> None:
+ args = parse_args()
+ device = torch.device(args.device) if args.device else torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+ run_root = args.root / args.model
+ clip_dir = run_root / args.clip_id
+ if not clip_dir.is_dir():
+ raise SystemExit(f"Run directory for clip not found: {clip_dir}")
+
+ cfg = load_cfg(clip_dir)
+ dataset, trainer = build_dataset_trainer(cfg, device)
+
+ ckpt = find_checkpoint(clip_dir)
+ if ckpt is None:
+ raise SystemExit(f"No checkpoint found in {clip_dir}")
+ trainer.resume_from_checkpoint(str(ckpt), load_only_model=True)
+ trainer.set_eval()
+
+ out_root = args.output_root
+
+ # --- render train/depth ---
+ render_train_and_depth_frames(
+ dataset=dataset,
+ trainer=trainer,
+ clip_id=args.clip_id,
+ model_name=args.model,
+ out_root=out_root,
+ types=args.types,
+ device=device,
+ )
+
+ # --- render novel ---
+ if "novel" in args.types:
+ render_novel_frames(
+ dataset=dataset,
+ trainer=trainer,
+ clip_id=args.clip_id,
+ model_name=args.model,
+ out_root=out_root,
+ device=device,
+ traj_types=["front_center_interp", "lateral_offset", "lateral_offset_left"],
+ cam_id=0,
+ )
+
+ # --- metrics ---
+ # train metrics JSON
+ train_metrics = load_train_metrics(run_root, args.clip_id) if "train" in args.types else {}
+ if train_metrics:
+ metric_path = out_root / "train" / args.clip_id / args.model / "train_metric.json"
+ metric_path.parent.mkdir(parents=True, exist_ok=True)
+ metric_path.write_text(json.dumps(train_metrics, indent=2))
+
+ # depth metrics JSON
+ depth_csvs = list(args.depth_csv) if args.depth_csv else [
+ Path("depth_tables_3/per_seq.csv"),
+ Path("depth_tables_2/per_seq.csv"),
+ ]
+ depth_metrics = load_depth_metrics(depth_csvs, args.model, args.clip_id, args.absrel_column) if "depth" in args.types else {}
+ if depth_metrics:
+ metric_path = out_root / "depth" / args.clip_id / args.model / "depth_metric.json"
+ metric_path.parent.mkdir(parents=True, exist_ok=True)
+ metric_path.write_text(json.dumps(depth_metrics, indent=2))
+
+ print(f"[DONE] Exported frames and metrics under {out_root}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/third_party/drivestudio/tools/extract_front_frames.py b/worldbench/third_party/drivestudio/tools/extract_front_frames.py
new file mode 100644
index 0000000..dad9832
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/extract_front_frames.py
@@ -0,0 +1,255 @@
+#!/usr/bin/env python3
+"""
+Extract per-frame RGB and (if available) depth images for a given clip and frame
+from: (1) training views, and (2) specific novel trajectories (front_center_interp,
+lateral_offset, lateral_offset_left) across multiple methods (including GT).
+
+Outputs a simple folder layout under --output-dir:
+ ///
+ train_rgb_f.png
+ train_depth_f.png (if available)
+ front_center_interp_rgb_f.png
+ front_center_interp_depth_f.png (if available)
+ lateral_offset_rgb_f.png
+ lateral_offset_depth_f.png (if available)
+ lateral_offset_left_rgb_f.png
+ lateral_offset_left_depth_f.png (if available)
+
+Assumptions:
+ - Training videos live at work_dirs///videos/full_set__rgbs.mp4 (and *_depths.mp4).
+ - Novel videos live at work_dirs///videos//.mp4, where
+ defaults to the lexicographically last 'novel_*' under videos/.
+ - Some runs may not carry novel depths; we skip them with a warning.
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import shutil
+import subprocess
+from pathlib import Path
+from typing import Dict, List, Optional, Sequence, Tuple
+
+
+DEFAULT_MODELS = [
+ "drivestudio-nus-gt",
+ "drivestudio-nus-dist4d",
+ "drivestudio-nus-drivedreamer2",
+ "drivestudio-nus-opendwm",
+ "drivestudio-nus-dreamforge",
+ "drivestudio-nus-xscene",
+ "drivestudio-nus-magicdrive",
+]
+
+DEFAULT_VIEWS = [
+ "front_center_interp",
+ "lateral_offset",
+ "lateral_offset_left",
+]
+
+
+def parse_args() -> argparse.Namespace:
+ p = argparse.ArgumentParser("Extract RGB/Depth images for selected frame across methods and views")
+ p.add_argument("clip_id", type=str, help="Clip identifier (e.g., 813)")
+ p.add_argument("frame", type=int, help="Zero-based frame index to extract")
+ p.add_argument("--root", type=Path, default=Path("work_dirs"), help="Root containing per-model runs (default: work_dirs)")
+ p.add_argument("--models", type=str, nargs="*", default=DEFAULT_MODELS, help="Model directory names under root (default: 7 incl. GT)")
+ p.add_argument("--step", type=int, default=30000, help="Training step in filenames (default: 30000)")
+ p.add_argument("--novel-dir", type=str, default=None, help="Explicit novel dir name (e.g., novel_30000_benchmark). Default: latest novel_* under videos.")
+ p.add_argument("--views", type=str, nargs="*", default=DEFAULT_VIEWS, help="Novel view names to extract (default: front_center_interp lateral_offset lateral_offset_left)")
+ p.add_argument("--output-dir", type=Path, default=Path("data/demo_frames"), help="Directory to write PNGs (default: data/demo_frames)")
+ p.add_argument("--ffmpeg-bin", type=str, default="ffmpeg", help="FFmpeg executable (default: ffmpeg)")
+ p.add_argument("--force", action="store_true", help="Overwrite existing files")
+ p.add_argument("--debug", action="store_true", help="Print debug info (paths, nb_frames)")
+ return p.parse_args()
+
+
+def resolve_bin(preferred: str, fallback_env: str, default: str) -> str:
+ for cand in [preferred, os.environ.get(fallback_env), default]:
+ if cand and shutil.which(cand):
+ return cand
+ raise SystemExit(f"Executable not found: {preferred} / ${fallback_env} / {default}")
+
+
+def ffprobe_bin(ffmpeg_bin: str) -> Optional[str]:
+ p = Path(ffmpeg_bin)
+ cand = p.parent / "ffprobe" if p.parent else None
+ if cand and cand.exists():
+ return str(cand)
+ auto = shutil.which("ffprobe")
+ return auto
+
+
+def probe_nb_frames(ffprobe: Optional[str], video: Path) -> Optional[int]:
+ if not ffprobe:
+ return None
+ try:
+ out = subprocess.check_output([
+ ffprobe,
+ "-v", "error",
+ "-select_streams", "v:0",
+ "-count_frames",
+ "-show_entries", "stream=nb_read_frames",
+ "-of", "csv=p=0",
+ str(video),
+ ], text=True).strip()
+ return int(out) if out.isdigit() else None
+ except Exception:
+ return None
+
+
+def select_novel_dir(videos_dir: Path, explicit: Optional[str]) -> Path:
+ if explicit:
+ d = videos_dir / explicit
+ if not d.is_dir():
+ raise SystemExit(f"Novel dir '{explicit}' not found under {videos_dir}")
+ return d
+ cands = sorted([p for p in videos_dir.iterdir() if p.is_dir() and p.name.startswith("novel_")])
+ if not cands:
+ raise SystemExit(f"No novel_* directory found inside {videos_dir}")
+ return cands[-1]
+
+
+def extract_frame(ffmpeg: str, src: Path, frame_idx: int, dst: Path, force: bool = False) -> None:
+ if dst.exists() and not force:
+ return
+ dst.parent.mkdir(parents=True, exist_ok=True)
+ # select=eq(n,frame_idx) picks the exact decoded frame; setsar=1 for square pixels.
+ cmd = [
+ ffmpeg,
+ "-y",
+ "-i", str(src),
+ "-vf", f"select='eq(n\,{frame_idx})',setsar=1",
+ "-vframes", "1",
+ str(dst),
+ ]
+ subprocess.run(cmd, check=True)
+
+
+def main() -> None:
+ args = parse_args()
+ ffmpeg = resolve_bin(args.ffmpeg_bin, "FFMPEG_BIN", "ffmpeg")
+ ffprobe = ffprobe_bin(ffmpeg)
+
+ out_root = args.output_dir / args.clip_id
+ methods = list(args.models)
+ if "drivestudio-nus-gt" not in methods:
+ methods.insert(0, "drivestudio-nus-gt")
+
+ # Train view sources
+ train_srcs: Dict[str, Tuple[Path, Optional[Path]]] = {}
+ for m in methods:
+ base = args.root / m / args.clip_id / "videos"
+ rgb = base / f"full_set_{args.step}_rgbs.mp4"
+ dpt = base / f"full_set_{args.step}_depths.mp4"
+ if not rgb.is_file():
+ print(f"[MISS] Train RGB not found: {rgb}")
+ continue
+ train_srcs[m] = (rgb, dpt if dpt.is_file() else None)
+
+ # Novel view sources
+ view_map: Dict[str, Dict[str, Path]] = {v: {} for v in args.views}
+ view_depth_map: Dict[str, Dict[str, Path]] = {v: {} for v in args.views}
+ for m in methods:
+ videos_dir = args.root / m / args.clip_id / "videos"
+ if not videos_dir.is_dir():
+ print(f"[MISS] videos dir for {m}: {videos_dir}")
+ continue
+ try:
+ novel_dir = select_novel_dir(videos_dir, args.novel_dir)
+ except SystemExit as e:
+ print(f"[MISS] {e}")
+ continue
+ files = {p.stem: p for p in novel_dir.glob("*.mp4")}
+ # try both exact and relaxed matches
+ for v in args.views:
+ # exact
+ if v in files:
+ view_map[v][m] = files[v]
+ else:
+ # relaxed: case-insensitive and underscore-insensitive contains
+ key = None
+ vn = v.lower().replace("_", "")
+ for stem, path in files.items():
+ sn = stem.lower().replace("_", "")
+ if vn == sn:
+ key = stem; break
+ if key is None:
+ for stem, path in files.items():
+ sn = stem.lower().replace("_", "")
+ if vn in sn:
+ key = stem; break
+ if key:
+ view_map[v][m] = files[key]
+ # depth candidate names
+ # common patterns: _depths.mp4 or _depth.mp4
+ for suf in ("_depths", "_depth"):
+ cand = novel_dir / f"{v}{suf}.mp4"
+ if cand.is_file():
+ view_depth_map[v][m] = cand
+ break
+
+ # Extract
+ for m in methods:
+ # Train RGB
+ if m in train_srcs:
+ rgb, dpt = train_srcs[m]
+ dst_rgb = out_root / m / f"train_rgb_f{args.frame:04d}.png"
+ try:
+ nb = probe_nb_frames(ffprobe, rgb)
+ if args.debug:
+ print(f"[DBG] {m} train rgb nb_frames={nb}")
+ if nb is not None and args.frame >= nb:
+ print(f"[SKIP] {m} train rgb frame {args.frame} out of range (nb={nb})")
+ else:
+ extract_frame(ffmpeg, rgb, args.frame, dst_rgb, args.force)
+ except subprocess.CalledProcessError:
+ print(f"[WARN] ffmpeg failed extracting train RGB for {m}")
+ if dpt and dpt.is_file():
+ dst_d = out_root / m / f"train_depth_f{args.frame:04d}.png"
+ try:
+ nb = probe_nb_frames(ffprobe, dpt)
+ if args.debug:
+ print(f"[DBG] {m} train depth nb_frames={nb}")
+ if nb is not None and args.frame >= nb:
+ print(f"[SKIP] {m} train depth frame {args.frame} out of range (nb={nb})")
+ else:
+ extract_frame(ffmpeg, dpt, args.frame, dst_d, args.force)
+ except subprocess.CalledProcessError:
+ print(f"[WARN] ffmpeg failed extracting train DEPTH for {m}")
+ # Novel views
+ for v in args.views:
+ src = view_map.get(v, {}).get(m)
+ if src is None:
+ print(f"[MISS] {m} novel view '{v}' not found")
+ continue
+ dst_rgb = out_root / m / f"{v}_rgb_f{args.frame:04d}.png"
+ try:
+ nb = probe_nb_frames(ffprobe, src)
+ if args.debug:
+ print(f"[DBG] {m} {v} rgb nb_frames={nb}")
+ if nb is not None and args.frame >= nb:
+ print(f"[SKIP] {m} {v} rgb frame {args.frame} out of range (nb={nb})")
+ else:
+ extract_frame(ffmpeg, src, args.frame, dst_rgb, args.force)
+ except subprocess.CalledProcessError:
+ print(f"[WARN] ffmpeg failed extracting novel RGB for {m}:{v}")
+ dsrc = view_depth_map.get(v, {}).get(m)
+ if dsrc and dsrc.is_file():
+ dst_d = out_root / m / f"{v}_depth_f{args.frame:04d}.png"
+ try:
+ nb = probe_nb_frames(ffprobe, dsrc)
+ if args.debug:
+ print(f"[DBG] {m} {v} depth nb_frames={nb}")
+ if nb is not None and args.frame >= nb:
+ print(f"[SKIP] {m} {v} depth frame {args.frame} out of range (nb={nb})")
+ else:
+ extract_frame(ffmpeg, dsrc, args.frame, dst_d, args.force)
+ except subprocess.CalledProcessError:
+ print(f"[WARN] ffmpeg failed extracting novel DEPTH for {m}:{v}")
+
+ print(f"[DONE] Wrote images under {out_root}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/third_party/drivestudio/tools/gen_token_clip_mapping.py b/worldbench/third_party/drivestudio/tools/gen_token_clip_mapping.py
new file mode 100644
index 0000000..38d8a23
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/gen_token_clip_mapping.py
@@ -0,0 +1,149 @@
+#!/usr/bin/env python3
+import argparse
+import json
+import pickle
+from pathlib import Path
+from typing import Dict, List, Optional
+
+
+def parse_args():
+ parser = argparse.ArgumentParser(
+ description="Generate mapping from first_sample_token to clip folder id."
+ )
+ parser.add_argument(
+ "--scene-json",
+ default="data/nuscenes_trainval/raw/v1.0-trainval/scene.json",
+ help="Path to NuScenes scene.json.",
+ )
+ parser.add_argument(
+ "--split",
+ default="val",
+ help="Split name defined in nuscenes.utils.splits (e.g., val, mini_val).",
+ )
+ parser.add_argument(
+ "--split-file",
+ default=None,
+ help="Path to a text file listing scene names to include (one per line).",
+ )
+ parser.add_argument(
+ "--num-scenes",
+ type=int,
+ default=150,
+ help="Max number of scenes to keep (mirrors preprocess default).",
+ )
+ parser.add_argument(
+ "--json-output",
+ default=None,
+ help="Optional path to save the mapping JSON. Prints to stdout when omitted.",
+ )
+ parser.add_argument(
+ "--pickle-output",
+ default=None,
+ help=(
+ "Optional path to save a pickle containing both token→clip and clip→token mappings."
+ ),
+ )
+ parser.add_argument(
+ "--pad-width",
+ type=int,
+ default=3,
+ help="Zero padding width for clip ids (folders).",
+ )
+ return parser.parse_args()
+
+
+def load_split(split_name: str) -> List[str]:
+ try:
+ from nuscenes.utils import splits as nusc_splits
+ except ModuleNotFoundError as exc:
+ raise RuntimeError(
+ "nuscenes-devkit not available; install it or provide --split-file."
+ ) from exc
+
+ if not hasattr(nusc_splits, split_name):
+ available = [attr for attr in dir(nusc_splits) if not attr.startswith("_")]
+ raise ValueError(f"Unknown split '{split_name}'. Available: {available}")
+ return getattr(nusc_splits, split_name)
+
+
+def load_split_from_file(split_file: Path) -> List[str]:
+ with Path(split_file).open("r") as fp:
+ return [line.strip() for line in fp if line.strip()]
+
+
+def build_mapping(
+ scene_json: Path,
+ split_names: Optional[List[str]],
+ num_scenes: Optional[int],
+ pad_width: int,
+) -> Dict[str, str]:
+ with Path(scene_json).open("r") as fp:
+ scenes = json.load(fp)
+
+ if split_names is None:
+ selected_indices = list(range(len(scenes)))
+ else:
+ selected_indices = [
+ idx for idx, scene in enumerate(scenes) if scene["name"] in split_names
+ ]
+ if num_scenes is not None:
+ selected_indices = selected_indices[:num_scenes]
+
+ mapping = {}
+ for scene_idx in selected_indices:
+ scene = scenes[scene_idx]
+ token = scene["first_sample_token"]
+ clip_id = str(scene_idx).zfill(pad_width)
+ mapping[token] = clip_id
+ return mapping
+
+
+def invert_mapping(mapping: Dict[str, str]) -> Dict[str, str]:
+ """Return clip_id -> token dict ensuring uniqueness."""
+ inverted: Dict[str, str] = {}
+ for token, clip_id in mapping.items():
+ if clip_id in inverted:
+ raise ValueError(
+ f"Duplicate clip id '{clip_id}' for tokens '{inverted[clip_id]}' and '{token}'."
+ )
+ inverted[clip_id] = token
+ return inverted
+
+
+def main():
+ args = parse_args()
+ if args.split_file:
+ split_names = load_split_from_file(args.split_file)
+ elif args.split:
+ split_names = load_split(args.split)
+ else:
+ split_names = None
+
+ mapping = build_mapping(
+ args.scene_json,
+ split_names,
+ args.num_scenes,
+ args.pad_width,
+ )
+ clip_to_token = invert_mapping(mapping)
+ output_str = json.dumps(mapping, indent=2, ensure_ascii=False)
+
+ if args.json_output:
+ Path(args.json_output).parent.mkdir(parents=True, exist_ok=True)
+ with Path(args.json_output).open("w") as fp:
+ fp.write(output_str)
+ else:
+ print(output_str)
+
+ if args.pickle_output:
+ Path(args.pickle_output).parent.mkdir(parents=True, exist_ok=True)
+ payload = {
+ "token_to_clip": mapping,
+ "clip_to_token": clip_to_token,
+ }
+ with Path(args.pickle_output).open("wb") as fp:
+ pickle.dump(payload, fp)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/third_party/drivestudio/tools/make_all_demo_videos.py b/worldbench/third_party/drivestudio/tools/make_all_demo_videos.py
new file mode 100644
index 0000000..87974c8
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/make_all_demo_videos.py
@@ -0,0 +1,214 @@
+#!/usr/bin/env python3
+"""
+Batch compose RGB/depth grid videos for many sequences, saved under data/.
+
+It discovers sequences that exist for all specified models (and have both
+RGB and depth videos), then calls:
+ - tools/compose_train_view_grid.py (RGB)
+ - tools/compose_depth_grid.py (Depth)
+
+Default models: drivestudio-nus-(dist4d, drivedreamer2, opendwm, dreamforge, xscene, magicdrive)
+Default step: 30000
+Default output dir: data/demo_videos/all_6m/
+
+Examples:
+ python3 tools/make_all_demo_videos.py # discover ~150 seqs, make both RGB+Depth
+ python3 tools/make_all_demo_videos.py --seqs 813 694 614 # only selected seqs
+ python3 tools/make_all_demo_videos.py --step 40000 # use different training step suffix
+
+You need ffmpeg installed (or set env FFMPEG_BIN).
+"""
+from __future__ import annotations
+
+import argparse
+import shutil
+import subprocess
+from pathlib import Path
+from typing import Iterable, List, Sequence, Set
+
+
+DEFAULT_MODELS = [
+ "drivestudio-nus-dist4d",
+ "drivestudio-nus-drivedreamer2",
+ "drivestudio-nus-opendwm",
+ "drivestudio-nus-dreamforge",
+ "drivestudio-nus-xscene",
+ "drivestudio-nus-magicdrive",
+]
+
+
+def parse_args() -> argparse.Namespace:
+ p = argparse.ArgumentParser(description="Batch compose RGB+Depth grid videos across many sequences.")
+ p.add_argument("--root", type=Path, default=Path("work_dirs"), help="Root directory containing model runs.")
+ p.add_argument("--models", type=str, nargs="*", default=DEFAULT_MODELS, help="Model directory names under root.")
+ p.add_argument("--with-gt", action="store_true", help="Include drivestudio-nus-gt as the first model if present.")
+ p.add_argument("--discover", type=str, choices=["union", "common"], default="union", help="Discover mode for sequences across models (default: union).")
+ p.add_argument("--step", type=int, default=30000, help="Training step used in file naming (default: 30000).")
+ p.add_argument("--seqs", type=str, nargs="*", default=None, help="Optional explicit list of clip IDs (e.g., 002 013 ...). If omitted, auto-discover by --discover mode.")
+ p.add_argument("--output-dir", type=Path, default=Path("data/demo_videos/all_6m"), help="Directory to write outputs.")
+ p.add_argument("--ffmpeg-bin", type=str, default=None, help="ffmpeg binary (overrides compose_* scripts' auto-detect).")
+ p.add_argument("--font", type=Path, default=None, help="Optional TTF/OTF font path for labels (propagated).")
+ p.add_argument("--fontsize", type=int, default=26, help="Label font size.")
+ p.add_argument("--gap", type=int, default=8, help="Gap in pixels between tiles.")
+ p.add_argument("--tile-width", type=int, default=0, help="Per-tile width; 0 disables resizing (default: 0).")
+ p.add_argument("--tile-height", type=int, default=0, help="Per-tile height; 0 disables resizing (default: 0).")
+ p.add_argument("--fit", type=str, choices=["contain", "cover"], default="contain", help="Resize mode when tile size is set: contain (scale+pad) or cover (scale+crop).")
+ p.add_argument("--fps", type=int, default=0, help="Optional FPS override (0 = keep source).")
+ p.add_argument("--frames", type=int, default=0, help="Optional trim to N frames (0 = keep full length).")
+ p.add_argument("--debug", action="store_true", help="Pass --debug to compose scripts for alignment prints.")
+ p.add_argument("--novel", action="store_true", help="Also compose novel-view grid (no metrics). Requires novel_* dirs.")
+ p.add_argument("--dry-run", action="store_true", help="Only print planned commands, do not run.")
+ return p.parse_args()
+
+
+def has_both_videos(root: Path, model: str, seq: str, step: int) -> bool:
+ base = root / model / seq / "videos"
+ rgb = base / f"full_set_{step}_rgbs.mp4"
+ dpt = base / f"full_set_{step}_depths.mp4"
+ return rgb.is_file() and dpt.is_file()
+
+
+def discover_seqs(root: Path, models: Sequence[str], step: int, mode: str) -> List[str]:
+ found_sets: List[Set[str]] = []
+ for m in models:
+ seqs: Set[str] = set()
+ model_dir = root / m
+ if not model_dir.is_dir():
+ continue
+ for d in sorted(model_dir.iterdir()):
+ if not d.is_dir():
+ continue
+ seq = d.name
+ if has_both_videos(root, m, seq, step):
+ seqs.add(seq)
+ found_sets.append(seqs)
+ if not found_sets:
+ return []
+ if mode == "common":
+ acc = set.intersection(*found_sets)
+ else:
+ acc = set.union(*found_sets)
+ return sorted(acc)
+
+
+def run(cmd: List[str], dry_run: bool = False) -> int:
+ if dry_run:
+ print("DRY-RUN:", " ".join(cmd))
+ return 0
+ try:
+ subprocess.run(cmd, check=True)
+ return 0
+ except subprocess.CalledProcessError as e:
+ print(f"[WARN] command failed (exit {e.returncode}):", " ".join(cmd))
+ return e.returncode
+
+
+def main() -> None:
+ args = parse_args()
+ args.output_dir.mkdir(parents=True, exist_ok=True)
+
+ # ffmpeg presence hint (compose_* will also check internally)
+ if args.ffmpeg_bin is None and shutil.which("ffmpeg") is None and shutil.which("avconv") is None:
+ print("[WARN] ffmpeg not found in PATH; set --ffmpeg-bin or environment FFMPEG_BIN if compose_* fails.")
+
+ models = list(args.models)
+ if args.with_gt and "drivestudio-nus-gt" not in models:
+ models.insert(0, "drivestudio-nus-gt")
+ if args.seqs is None or len(args.seqs) == 0:
+ seqs = discover_seqs(args.root, models, args.step, args.discover)
+ if not seqs:
+ raise SystemExit(f"No sequences discovered across models (mode={args.discover}).")
+ print(f"[INFO] Discovered {len(seqs)} sequences (mode={args.discover}).")
+ else:
+ seqs = list(args.seqs)
+ print(f"[INFO] Using user-provided {len(seqs)} sequences.")
+
+ ffmpeg_args: List[str] = []
+ if args.ffmpeg_bin:
+ ffmpeg_args = ["--ffmpeg-bin", args.ffmpeg_bin]
+
+ font_args: List[str] = []
+ if args.font:
+ font_args = ["--font", str(args.font), "--fontsize", str(args.fontsize)]
+ else:
+ font_args = ["--fontsize", str(args.fontsize)]
+ # Always pass gap; pass tile args when provided (support height-only to mirror compose_* behavior)
+ tile_args: List[str] = ["--gap", str(args.gap)]
+ if args.tile_height and args.tile_height > 0:
+ # Respect height-only mode; do not force a width to avoid introducing internal letterbox.
+ tile_args += ["--tile-height", str(args.tile_height)]
+ if args.tile_width and args.tile_width > 0:
+ tile_args += ["--tile-width", str(args.tile_width)]
+ if (args.tile_height and args.tile_height > 0) or (args.tile_width and args.tile_width > 0):
+ tile_args += ["--fit", args.fit]
+ fps_args: List[str] = []
+ if args.fps and args.fps > 0:
+ fps_args += ["--fps", str(args.fps)]
+ if args.frames and args.frames > 0:
+ fps_args += ["--frames", str(args.frames)]
+ debug_args: List[str] = ["--debug"] if args.debug else []
+
+ ok = 0
+ failed = 0
+ for seq in seqs:
+ # Per-sequence available models (in case some method is missing videos for this seq)
+ models_this_seq = [m for m in models if has_both_videos(args.root, m, seq, args.step)]
+ if len(models_this_seq) == 0:
+ print(f"[SKIP] {seq}: no available models with both RGB+Depth videos.")
+ continue
+ if len(models_this_seq) < len(models):
+ print(f"[WARN] {seq}: missing some methods; composing with {len(models_this_seq)} models: {', '.join(models_this_seq)}")
+ print(f"[RGB ] composing clip {seq}")
+ # dynamic suffix reflecting model count and GT presence
+ has_gt = args.with_gt and ("drivestudio-nus-gt" in models_this_seq)
+ suffix = f"{len(models_this_seq)}m" + ("_gt" if has_gt else "")
+ rc1 = run([
+ "python3", "tools/compose_train_view_grid.py", seq,
+ "--root", str(args.root),
+ "--models", *models_this_seq,
+ "--step", str(args.step),
+ "--output", str(args.output_dir / f"clip_{seq}_train_grid_{suffix}.mp4"),
+ *tile_args,
+ *ffmpeg_args,
+ *font_args,
+ *fps_args,
+ *debug_args,
+ ], dry_run=args.dry_run)
+
+ print(f"[DEPTH] composing clip {seq}")
+ rc2 = run([
+ "python3", "tools/compose_depth_grid.py", seq,
+ "--root", str(args.root),
+ "--models", *models_this_seq,
+ "--step", str(args.step),
+ "--output", str(args.output_dir / f"clip_{seq}_depth_grid_{suffix}.mp4"),
+ *tile_args,
+ *ffmpeg_args,
+ *font_args,
+ *fps_args,
+ *debug_args,
+ ], dry_run=args.dry_run)
+
+ if rc1 == 0: ok += 1
+ else: failed += 1
+ if rc2 == 0: ok += 1
+ else: failed += 1
+
+ # Optional: compose novel-view grid (with GT if requested & available)
+ if args.novel and args.with_gt and "drivestudio-nus-gt" in models_this_seq:
+ print(f"[NOVEL+GT] composing clip {seq}")
+ rc3 = run([
+ "python3", "tools/compose_novel_view_grid.py", seq,
+ "--root", str(args.root),
+ "--models", *models_this_seq,
+ "--output", str(args.output_dir / f"clip_{seq}_novel_grid_{suffix}.mp4"),
+ "--fontsize", str(args.fontsize),
+ ], dry_run=args.dry_run)
+ if rc3 == 0: ok += 1
+ else: failed += 1
+
+ print(f"[DONE] Wrote videos to {args.output_dir} (ok={ok}, failed={failed})")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/third_party/drivestudio/tools/make_metric_compare_videos.py b/worldbench/third_party/drivestudio/tools/make_metric_compare_videos.py
new file mode 100644
index 0000000..d1ea64f
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/make_metric_compare_videos.py
@@ -0,0 +1,277 @@
+#!/usr/bin/env python3
+"""Make per-clip metric-compare demo videos for LPIPS and AbsRel.
+
+For each clip and each metric type (LPIPS for RGB, AbsRel for depth),
+pick the best-2 and worst-2 methods (excluding GT) and compose a 2x2
+grid video:
+
+ rows: 1st row = best two methods (left: best, right: 2nd best)
+ 2nd row = worst two methods (left: 2nd worst, right: worst)
+
+Each tile is the original 6-view training/depth video for that method
+and clip (full_set__rgbs.mp4 / full_set__depths.mp4).
+Metric values are overlaid by compose_*_grid.py via drawtext.
+
+Output layout:
+ data/metric_compare/LPIPS/.mp4
+ data/metric_compare/AbsRel/.mp4
+
+This script reuses:
+ - work_dirs_analysis_metrics.json (LPIPS per-seq)
+ - depth_tables_3/per_seq.csv (fallback depth_tables_2/per_seq.csv) (AbsRel per-seq)
+ - tools/compose_train_view_grid.py (RGB grid with LPIPS labels)
+ - tools/compose_depth_grid.py (Depth grid with AbsRel labels)
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import subprocess
+from pathlib import Path
+from typing import Dict, List, Optional, Sequence, Tuple
+
+
+DEFAULT_MODELS = [
+ "drivestudio-nus-dist4d",
+ "drivestudio-nus-drivedreamer2",
+ "drivestudio-nus-opendwm",
+ "drivestudio-nus-dreamforge",
+ "drivestudio-nus-xscene",
+ "drivestudio-nus-magicdrive",
+]
+
+
+def parse_args() -> argparse.Namespace:
+ p = argparse.ArgumentParser("Compose LPIPS/AbsRel metric-compare videos per clip")
+ # If no clip IDs are given, we will auto-discover all clips present in the
+ # depth per-seq CSVs (i.e., all clips with AbsRel metrics).
+ p.add_argument("clip_id", nargs="*", help="Clip IDs (e.g., 813 694). If omitted, auto-discover from depth tables.")
+ p.add_argument("--root", type=Path, default=Path("work_dirs"), help="Root directory containing model runs")
+ p.add_argument("--models", type=str, nargs="*", default=None,
+ help="Model directory names under root (excluding GT). Default: common drivestudio-nus-* models.")
+ p.add_argument("--step", type=int, default=30000, help="Training step used in file names (default: 30000)")
+ p.add_argument("--metrics-json", type=Path, default=Path("work_dirs_analysis_metrics.json"),
+ help="LPIPS metrics JSON (default: work_dirs_analysis_metrics.json)")
+ p.add_argument("--metric-key", type=str, default="image_metrics/full/lpips",
+ help="Metric key in JSON for LPIPS (default: image_metrics/full/lpips)")
+ p.add_argument("--depth-csv", type=Path, nargs="*", default=None,
+ help="Depth per-seq CSVs (default: depth_tables_3/per_seq.csv, depth_tables_2/per_seq.csv)")
+ p.add_argument("--absrel-column", type=str, default="AbsRel", help="Column name for AbsRel in depth CSV")
+ p.add_argument("--output-root", type=Path, default=Path("data/metric_compare"),
+ help="Root directory for metric-compare videos")
+ p.add_argument("--ffmpeg-bin", type=str, default=None, help="Optional ffmpeg binary (propagated to compose_* scripts)")
+ p.add_argument("--font", type=Path, default=None, help="Optional TTF/OTF font for labels (propagated)")
+ p.add_argument("--fontsize", type=int, default=26, help="Font size for labels")
+ p.add_argument("--gap", type=int, default=8, help="Gap in pixels between tiles (passed to compose scripts)")
+ p.add_argument("--tile-width", type=int, default=1280, help="Tile width (default: 1280; 0=auto)")
+ p.add_argument("--tile-height", type=int, default=720, help="Tile height (default: 720; 0=auto)")
+ p.add_argument("--fit", type=str, choices=["contain", "cover"], default="contain",
+ help="Resize mode for tiles (default: contain)")
+ p.add_argument("--fps", type=int, default=0, help="Optional FPS override for composed videos (0=keep source)")
+ p.add_argument("--frames", type=int, default=0, help="Optional trim to N frames (0=full length)")
+ p.add_argument("--debug", action="store_true", help="Print detailed commands and selection info")
+ p.add_argument("--dry-run", action="store_true", help="Only print planned commands, do not run")
+ return p.parse_args()
+
+
+def method_label(name: str) -> str:
+ return name.split("drivestudio-nus-", 1)[-1].replace("_", " ").replace("-", " ")
+
+
+def load_lpips(path: Path) -> Dict[str, Dict[str, Dict[str, float]]]:
+ if not path.is_file():
+ raise SystemExit(f"LPIPS metrics JSON not found: {path}")
+ try:
+ return json.loads(path.read_text())
+ except Exception as e:
+ raise SystemExit(f"Failed to parse {path}: {e}")
+
+
+def load_absrel(paths: Sequence[Path], col: str) -> Dict[Tuple[str, str], float]:
+ lookup: Dict[Tuple[str, str], float] = {}
+ for p in paths:
+ if not p.is_file():
+ continue
+ try:
+ with p.open("r", newline="") as f:
+ reader = csv.DictReader(f)
+ if "Method/Seq" not in reader.fieldnames or col not in reader.fieldnames:
+ continue
+ for row in reader:
+ tag = row.get("Method/Seq", "").strip()
+ val = row.get(col)
+ if not tag or val is None:
+ continue
+ if "/" not in tag:
+ continue
+ method, seq = tag.split("/", 1)
+ try:
+ lookup[(method.strip(), seq.strip())] = float(val)
+ except ValueError:
+ continue
+ except Exception as e:
+ print(f"[WARN] failed reading {p}: {e}")
+ return lookup
+
+
+def has_video(root: Path, model: str, clip: str, step: int, kind: str) -> bool:
+ base = root / model / clip / "videos"
+ if kind == "rgb":
+ return (base / f"full_set_{step}_rgbs.mp4").is_file()
+ if kind == "depth":
+ return (base / f"full_set_{step}_depths.mp4").is_file()
+ return False
+
+
+def run(cmd: List[str], dry_run: bool = False, debug: bool = False) -> int:
+ if debug or dry_run:
+ print("RUN:", " ".join(str(c) for c in cmd))
+ if dry_run:
+ return 0
+ try:
+ subprocess.run(cmd, check=True)
+ return 0
+ except subprocess.CalledProcessError as e:
+ print(f"[WARN] command failed (exit {e.returncode}):", " ".join(str(c) for c in cmd))
+ return e.returncode
+
+
+def select_best_worst(values: Dict[str, float]) -> Optional[List[str]]:
+ """Given model -> metric (lower is better), return [best1, best2, worst2, worst1]."""
+ if len(values) < 4:
+ return None
+ # sort ascending by metric (lower is better)
+ ordered = sorted(values.items(), key=lambda kv: kv[1])
+ best1, best2 = ordered[0][0], ordered[1][0]
+ worst1, worst2 = ordered[-1][0], ordered[-2][0]
+ return [best1, best2, worst2, worst1]
+
+
+def main() -> None:
+ args = parse_args()
+
+ lpips_json = load_lpips(args.metrics_json)
+ depth_paths: List[Path] = list(args.depth_csv) if args.depth_csv else [
+ Path("depth_tables_3/per_seq.csv"),
+ Path("depth_tables_2/per_seq.csv"),
+ ]
+ absrel_lookup = load_absrel(depth_paths, args.absrel_column)
+
+ models_base: List[str] = list(args.models) if args.models else list(DEFAULT_MODELS)
+ # ensure no GT in base list
+ models_base = [m for m in models_base if not m.endswith("-gt")]
+
+ # shared compose args
+ tile_args: List[str] = ["--gap", str(args.gap)]
+ if args.tile_height and args.tile_height > 0:
+ tile_args += ["--tile-height", str(args.tile_height)]
+ if args.tile_width and args.tile_width > 0:
+ tile_args += ["--tile-width", str(args.tile_width), "--fit", args.fit]
+ fps_args: List[str] = []
+ if args.fps and args.fps > 0:
+ fps_args += ["--fps", str(args.fps)]
+ if args.frames and args.frames > 0:
+ fps_args += ["--frames", str(args.frames)]
+
+ font_args: List[str] = []
+ if args.font:
+ font_args = ["--font", str(args.font), "--fontsize", str(args.fontsize)]
+ else:
+ font_args = ["--fontsize", str(args.fontsize)]
+
+ ffmpeg_args: List[str] = []
+ if args.ffmpeg_bin:
+ ffmpeg_args = ["--ffmpeg-bin", args.ffmpeg_bin]
+
+ out_lpips_root = args.output_root / "LPIPS"
+ out_absrel_root = args.output_root / "AbsRel"
+ out_lpips_root.mkdir(parents=True, exist_ok=True)
+ out_absrel_root.mkdir(parents=True, exist_ok=True)
+
+ # Determine clips: use provided list, or auto-discover from depth tables.
+ if args.clip_id:
+ clips = list(args.clip_id)
+ else:
+ clips = sorted({seq for (_, seq) in absrel_lookup.keys()})
+ if not clips:
+ raise SystemExit("No clips found in depth tables; please provide clip_id explicitly.")
+
+ for clip in clips:
+ # Collect per-model LPIPS and AbsRel for this clip.
+ # LPIPS only requires RGB video + LPIPS metric; AbsRel only requires depth video + AbsRel.
+ lpips_vals: Dict[str, float] = {}
+ absrel_vals: Dict[str, float] = {}
+ for model in models_base:
+ # LPIPS (train RGB)
+ if has_video(args.root, model, clip, args.step, "rgb"):
+ seq_map = lpips_json.get(model, {})
+ if isinstance(seq_map, dict):
+ clip_metrics = seq_map.get(clip, {})
+ if isinstance(clip_metrics, dict):
+ v = clip_metrics.get(args.metric_key)
+ if isinstance(v, (int, float)):
+ lpips_vals[model] = float(v)
+ # AbsRel (depth)
+ if has_video(args.root, model, clip, args.step, "depth"):
+ meth = method_label(model).replace(" ", "")
+ absrel_val: Optional[float] = None
+ key1 = (meth, clip)
+ key2 = (meth.lower(), clip)
+ if key1 in absrel_lookup:
+ absrel_val = absrel_lookup[key1]
+ elif key2 in absrel_lookup:
+ absrel_val = absrel_lookup[key2]
+ if absrel_val is not None:
+ absrel_vals[model] = absrel_val
+
+ if args.debug:
+ print(f"\n[DEBUG] clip {clip} candidates:")
+ for m in models_base:
+ print(f" {m}: LPIPS={lpips_vals.get(m)} AbsRel={absrel_vals.get(m)}")
+
+ # LPIPS selection
+ order_lpips = select_best_worst(lpips_vals)
+ if order_lpips is None:
+ print(f"[WARN] clip {clip}: not enough models with LPIPS to select 4 (have {len(lpips_vals)})")
+ else:
+ out_path = out_lpips_root / f"{clip}.mp4"
+ cmd = [
+ "python3", "tools/compose_train_view_grid.py", clip,
+ "--root", str(args.root),
+ "--models", *order_lpips,
+ "--step", str(args.step),
+ "--output", str(out_path),
+ *tile_args,
+ *ffmpeg_args,
+ *font_args,
+ *fps_args,
+ ]
+ rc = run(cmd, dry_run=args.dry_run, debug=args.debug)
+ if rc == 0:
+ print(f"[OK] LPIPS compare for clip {clip} -> {out_path}")
+
+ # AbsRel selection
+ order_absrel = select_best_worst(absrel_vals)
+ if order_absrel is None:
+ print(f"[WARN] clip {clip}: not enough models with AbsRel to select 4 (have {len(absrel_vals)})")
+ else:
+ out_path = out_absrel_root / f"{clip}.mp4"
+ cmd = [
+ "python3", "tools/compose_depth_grid.py", clip,
+ "--root", str(args.root),
+ "--models", *order_absrel,
+ "--step", str(args.step),
+ "--output", str(out_path),
+ *tile_args,
+ *ffmpeg_args,
+ *font_args,
+ *fps_args,
+ ]
+ rc = run(cmd, dry_run=args.dry_run, debug=args.debug)
+ if rc == 0:
+ print(f"[OK] AbsRel compare for clip {clip} -> {out_path}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/third_party/drivestudio/tools/metrics_analysis.py b/worldbench/third_party/drivestudio/tools/metrics_analysis.py
new file mode 100644
index 0000000..6e1b08e
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/metrics_analysis.py
@@ -0,0 +1,247 @@
+#!/usr/bin/env python3
+"""
+Aggregate DriveStudio evaluation metrics and report per-model statistics.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import math
+from collections import OrderedDict, defaultdict
+from pathlib import Path
+from statistics import mean, pstdev
+from typing import Dict, Iterable, List, Mapping, MutableMapping, Optional, Sequence, Tuple
+
+
+DEFAULT_METRICS = OrderedDict(
+ [
+ ("image_metrics/full/psnr", "PSNR"),
+ ("image_metrics/full/ssim", "SSIM"),
+ ("image_metrics/full/lpips", "LPIPS"),
+ ("image_metrics/full/masked_psnr", "Masked PSNR"),
+ ("image_metrics/full/masked_ssim", "Masked SSIM"),
+ ]
+)
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(
+ description="Summarize DriveStudio reconstruction metrics across models."
+ )
+ parser.add_argument(
+ "--work-dirs",
+ type=Path,
+ default=Path("work_dirs"),
+ help="Root directory that contains per-model subdirectories (default: work_dirs).",
+ )
+ parser.add_argument(
+ "--reference",
+ type=str,
+ default="drivestudio-nus-gt",
+ help="Reference model used for delta statistics (default: drivestudio-nus-gt).",
+ )
+ parser.add_argument(
+ "--metrics-json",
+ type=Path,
+ default=Path("work_dirs_analysis_metrics.json"),
+ help="Optional cache file to store raw metrics JSON (default: work_dirs_analysis_metrics.json).",
+ )
+ parser.add_argument(
+ "--metrics",
+ type=str,
+ nargs="*",
+ default=list(DEFAULT_METRICS.keys()),
+ help="Metric keys to include (defaults to a standard set).",
+ )
+ parser.add_argument(
+ "--models",
+ type=str,
+ nargs="*",
+ default=None,
+ help="Subset of model directory names to analyze (default: all under work_dirs).",
+ )
+ return parser.parse_args()
+
+
+def load_latest_metric_file(metrics_dir: Path) -> Optional[Path]:
+ json_files = sorted(
+ (p for p in metrics_dir.glob("*.json") if p.is_file()),
+ key=lambda p: p.stat().st_mtime,
+ )
+ return json_files[-1] if json_files else None
+
+
+def collect_metrics(root: Path, metric_keys: Iterable[str]) -> Dict[str, Dict[str, Dict[str, float]]]:
+ data: Dict[str, Dict[str, Dict[str, float]]] = {}
+ for model_dir in sorted(root.iterdir()):
+ if not model_dir.is_dir():
+ continue
+ model_name = model_dir.name
+ model_metrics: Dict[str, Dict[str, float]] = {}
+ for clip_dir in sorted(model_dir.iterdir()):
+ if not clip_dir.is_dir():
+ continue
+ latest = load_latest_metric_file(clip_dir / "metrics")
+ if latest is None:
+ continue
+ try:
+ payload = json.loads(latest.read_text())
+ except json.JSONDecodeError:
+ continue
+ filtered = {k: payload[k] for k in metric_keys if k in payload}
+ if filtered:
+ model_metrics[clip_dir.name] = filtered
+ if model_metrics:
+ data[model_name] = model_metrics
+ return data
+
+
+def compute_stats(values: Sequence[float]) -> Tuple[float, float]:
+ if not values:
+ return (math.nan, math.nan)
+ if len(values) == 1:
+ return (values[0], 0.0)
+ return (mean(values), pstdev(values))
+
+
+def summarize_common_clips(
+ data: Mapping[str, Mapping[str, Mapping[str, float]]],
+ models: Sequence[str],
+ chosen_metrics: Mapping[str, str],
+) -> Tuple[List[str], Dict[str, Dict[str, Tuple[float, float, int]]]]:
+ clip_sets = [set(data[m].keys()) for m in models if m in data]
+ if not clip_sets:
+ return [], {}
+ common = sorted(set.intersection(*clip_sets))
+ summary: Dict[str, Dict[str, Tuple[float, float, int]]] = {}
+ for model in models:
+ model_summary: Dict[str, Tuple[float, float, int]] = {}
+ clips = data.get(model, {})
+ for metric_key, metric_label in chosen_metrics.items():
+ values = [clips[c][metric_key] for c in common if metric_key in clips.get(c, {})]
+ if values:
+ m, s = compute_stats(values)
+ model_summary[metric_label] = (m, s, len(values))
+ if model_summary:
+ summary[model] = model_summary
+ return common, summary
+
+
+def summarize_deltas(
+ data: Mapping[str, Mapping[str, Mapping[str, float]]],
+ models: Sequence[str],
+ reference: str,
+ clip_subset: Sequence[str],
+ chosen_metrics: Mapping[str, str],
+) -> Dict[str, Dict[str, Tuple[float, float, int]]]:
+ deltas: Dict[str, Dict[str, Tuple[float, float, int]]] = {}
+ ref_clips = data.get(reference, {})
+ for model in models:
+ if model == reference or model not in data:
+ continue
+ model_clips = data[model]
+ metric_summary: Dict[str, Tuple[float, float, int]] = {}
+ for metric_key, metric_label in chosen_metrics.items():
+ diff_vals: List[float] = []
+ for clip in clip_subset:
+ if metric_key in model_clips.get(clip, {}) and metric_key in ref_clips.get(clip, {}):
+ diff_vals.append(model_clips[clip][metric_key] - ref_clips[clip][metric_key])
+ if diff_vals:
+ m, s = compute_stats(diff_vals)
+ metric_summary[metric_label] = (m, s, len(diff_vals))
+ if metric_summary:
+ deltas[model] = metric_summary
+ return deltas
+
+
+def write_raw_cache(path: Path, payload: Mapping[str, Mapping[str, Mapping[str, float]]]) -> None:
+ with path.open("w") as f:
+ json.dump(payload, f, indent=2)
+
+
+def format_table(
+ header: Sequence[str],
+ rows: Sequence[Tuple[str, Sequence[str]]],
+) -> str:
+ widths = [len(col) for col in header]
+ for _, cells in rows:
+ for idx, cell in enumerate(cells):
+ widths[idx + 1] = max(widths[idx + 1], len(cell))
+ line_parts = ["{:<{w}}".format(header[0], w=widths[0])]
+ for idx, col in enumerate(header[1:], start=1):
+ line_parts.append("{:>{w}}".format(col, w=widths[idx]))
+ output = [" ".join(line_parts)]
+ output.append(" ".join("-" * w for w in widths))
+ for label, cells in rows:
+ row_parts = ["{:<{w}}".format(label, w=widths[0])]
+ for idx, cell in enumerate(cells, start=1):
+ row_parts.append("{:>{w}}".format(cell, w=widths[idx]))
+ output.append(" ".join(row_parts))
+ return "\n".join(output)
+
+
+def main() -> None:
+ args = parse_args()
+
+ metric_map = OrderedDict((k, DEFAULT_METRICS.get(k, k)) for k in args.metrics)
+ data = collect_metrics(args.work_dirs, metric_map.keys())
+ if not data:
+ raise SystemExit("No metrics found under {}".format(args.work_dirs))
+
+ models = args.models if args.models else sorted(data.keys())
+ if args.reference not in data:
+ raise SystemExit(f"Reference model '{args.reference}' not found in metrics.")
+
+ write_raw_cache(args.metrics_json, data)
+
+ common_clips, summary = summarize_common_clips(data, models, metric_map)
+ if not common_clips:
+ raise SystemExit("No common clips across selected models; aborting.")
+
+ print(f"Common clips across models ({len(common_clips)}): {', '.join(sorted(common_clips))}")
+ print()
+
+ header = ["Metric"] + [m for m in models if m in summary]
+ rows: List[Tuple[str, Sequence[str]]] = []
+ for metric_label in metric_map.values():
+ cells = []
+ for model in models:
+ stats = summary.get(model, {}).get(metric_label)
+ if stats:
+ mean_val, std_val, count = stats
+ cells.append(f"{mean_val:.3f}±{std_val:.3f} (n={count})")
+ else:
+ cells.append("n/a")
+ rows.append((metric_label, cells))
+ print("Per-model statistics (mean±std over common clips):")
+ print(format_table(header, rows))
+ print()
+
+ delta_summary = summarize_deltas(
+ data=data,
+ models=models,
+ reference=args.reference,
+ clip_subset=common_clips,
+ chosen_metrics=metric_map,
+ )
+ if delta_summary:
+ header = ["Metric"] + [m for m in models if m != args.reference and m in delta_summary]
+ rows = []
+ for metric_label in metric_map.values():
+ cells = []
+ for model in models:
+ if model == args.reference:
+ continue
+ stats = delta_summary.get(model, {}).get(metric_label)
+ if stats:
+ mean_val, std_val, count = stats
+ cells.append(f"{mean_val:+.3f}±{std_val:.3f} (n={count})")
+ else:
+ cells.append("n/a")
+ rows.append((metric_label, cells))
+ print(f"Deltas relative to reference '{args.reference}':")
+ print(format_table(header, rows))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/third_party/drivestudio/tools/render_all_for_model.py b/worldbench/third_party/drivestudio/tools/render_all_for_model.py
new file mode 100644
index 0000000..03d59a7
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/render_all_for_model.py
@@ -0,0 +1,401 @@
+#!/usr/bin/env python3
+"""Render and export train/depth/novel results for a given model across clips.
+
+For each clip and task (train, depth, novel) this script:
+
+- Re-renders images via DrivingDataset + trainer (no mp4 decode).
+- Saves per-camera (or per-view) JPEG frames and a video under:
+
+ data/render_all/{task_name}/{clip_id}/{model_name}/{camera_name}/
+ images/ # all frames, jpg
+ {task_name}.mp4 # video composed from images
+
+ Metrics (per clip+model, shared across cameras):
+
+ data/render_all/train/{clip_id}/{model_name}/train_metric.json
+ data/render_all/depth/{clip_id}/{model_name}/depth_metric.json
+
+Where:
+ - task_name ∈ {train, depth, novel}
+ - camera_name:
+ * train/depth: real camera name (e.g., CAM_FRONT)
+ * novel: one of the standard novel trajectories (treated as "views"):
+ - s_curve
+ - front_center_interp
+ - lateral_offset
+ - lateral_offset_left
+
+Metrics:
+ - train_metric.json:
+ taken from work_dirs_analysis_metrics.json for this model+clip
+ (same LPIPS/PSNR/SSIM entries as tools/metrics_analysis.py uses).
+ - depth_metric.json:
+ taken from depth_tables_3/per_seq.csv (fallback depth_tables_2/per_seq.csv),
+ selecting Method/Seq='{method}/{clip_id}' where method is method_label(model).
+ Contains AbsRel 等所有数值列.
+
+Novel currently has no metrics and thus no *_metric.json.
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import os
+import subprocess
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import imageio
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from PIL import Image
+
+from datasets.driving_dataset import DrivingDataset
+from utils.misc import import_str
+from utils.camera import lateral_offset_trajectory
+from utils.visualization import depth_visualizer
+
+
+def load_cfg(clip_dir: Path) -> OmegaConf:
+ cfg_path = clip_dir / "config.yaml"
+ if not cfg_path.is_file():
+ raise SystemExit(f"config.yaml not found in {clip_dir}")
+ return OmegaConf.load(cfg_path)
+
+
+def find_checkpoint(clip_dir: Path) -> Optional[Path]:
+ final = clip_dir / "checkpoint_final.pth"
+ if final.is_file():
+ return final
+ cands = sorted(clip_dir.glob("checkpoint_*.pth"))
+ return cands[-1] if cands else None
+
+
+def build_dataset_trainer(cfg: OmegaConf, device: torch.device):
+ dataset = DrivingDataset(data_cfg=cfg.data)
+ trainer = import_str(cfg.trainer.type)(
+ **cfg.trainer,
+ num_timesteps=dataset.num_img_timesteps,
+ model_config=cfg.model,
+ num_train_images=len(dataset.train_image_set),
+ num_full_images=len(dataset.full_image_set),
+ test_set_indices=dataset.test_timesteps,
+ scene_aabb=dataset.get_aabb().reshape(2, 3),
+ device=device,
+ )
+ return dataset, trainer
+
+
+def method_label(name: str) -> str:
+ return name.split("drivestudio-nus-", 1)[-1].replace("_", " ").replace("-", " ")
+
+
+def save_jpg(rgb: np.ndarray, path: Path) -> None:
+ arr = (np.clip(rgb, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8)
+ path.parent.mkdir(parents=True, exist_ok=True)
+ Image.fromarray(arr).save(path, format="JPEG", quality=95)
+
+
+def save_depth_color(depth_m: np.ndarray, opacity: Optional[np.ndarray], path: Path) -> None:
+ d = depth_m
+ if d.ndim == 3 and d.shape[-1] == 1:
+ d = d[..., 0]
+ w = opacity
+ if w is not None and isinstance(w, np.ndarray) and w.ndim == 3 and w.shape[-1] == 1:
+ w = w[..., 0]
+ color = depth_visualizer(d, w)
+ save_jpg(color, path)
+
+
+def write_video_from_images(img_dir: Path, out_path: Path, fps: int) -> None:
+ files = sorted(p for p in img_dir.glob("*.jpg") if p.is_file())
+ if not files:
+ return
+ out_path.parent.mkdir(parents=True, exist_ok=True)
+ writer = imageio.get_writer(out_path.as_posix(), mode="I", fps=fps)
+ try:
+ for p in files:
+ frame = imageio.imread(p.as_posix())
+ writer.append_data(frame)
+ finally:
+ writer.close()
+
+
+def load_lpips_json(path: Path) -> Dict[str, Dict[str, Dict[str, float]]]:
+ if not path.is_file():
+ return {}
+ try:
+ return json.loads(path.read_text())
+ except Exception:
+ return {}
+
+
+def load_absrel_tables(paths: List[Path], col: str = "AbsRel") -> Dict[Tuple[str, str], Dict[str, float]]:
+ """Load per-seq depth metrics indexed by (method, seq) -> row-dict."""
+ lookup: Dict[Tuple[str, str], Dict[str, float]] = {}
+ for p in paths:
+ if not p.is_file():
+ continue
+ try:
+ with p.open("r", newline="") as f:
+ reader = csv.DictReader(f)
+ if "Method/Seq" not in reader.fieldnames:
+ continue
+ for row in reader:
+ tag = row.get("Method/Seq", "").strip()
+ if not tag or "/" not in tag:
+ continue
+ method, seq = tag.split("/", 1)
+ key = (method.strip(), seq.strip())
+ # store all numeric columns
+ vals: Dict[str, float] = {}
+ for k, v in row.items():
+ if k == "Method/Seq":
+ continue
+ try:
+ vals[k] = float(v)
+ except (TypeError, ValueError):
+ continue
+ lookup[key] = vals
+ except Exception:
+ continue
+ return lookup
+
+
+@torch.no_grad()
+def render_train_depth_for_clip(
+ dataset: DrivingDataset,
+ trainer,
+ clip_id: str,
+ model_name: str,
+ out_root: Path,
+ tasks: List[str],
+ device: torch.device,
+) -> None:
+ need_train = "train" in tasks
+ need_depth = "depth" in tasks
+ if not (need_train or need_depth):
+ return
+
+ ds = dataset.full_image_set
+ camera_downscale = trainer._get_downscale_factor()
+ for idx in range(len(ds)):
+ image_infos, cam_infos = ds.get_image(idx, camera_downscale)
+ for k, v in list(image_infos.items()):
+ if isinstance(v, torch.Tensor):
+ image_infos[k] = v.to(device, non_blocking=True)
+ for k, v in list(cam_infos.items()):
+ if isinstance(v, torch.Tensor):
+ cam_infos[k] = v.to(device, non_blocking=True)
+ out = trainer(image_infos=image_infos, camera_infos=cam_infos)
+ rgb = out["rgb"].clamp(0.0, 1.0).detach().cpu().numpy()
+ depth = out.get("depth", None)
+ opacity = out.get("opacity", None)
+ depth_np = depth.detach().cpu().numpy() if isinstance(depth, torch.Tensor) else None
+ opacity_np = opacity.detach().cpu().numpy() if isinstance(opacity, torch.Tensor) else None
+
+ cam_name = cam_infos["cam_name"]
+ frame_idx = int(image_infos["frame_idx"].flatten()[0].cpu().item())
+
+ if need_train:
+ img_dir = out_root / "train" / clip_id / model_name / cam_name / "images"
+ out_path = img_dir / f"{frame_idx:04d}.jpg"
+ save_jpg(rgb, out_path)
+ gt_pixels = image_infos.get("pixels")
+ if isinstance(gt_pixels, torch.Tensor):
+ gt_dir = out_root / "train" / clip_id / model_name / cam_name / "gt_images"
+ gt_path = gt_dir / f"{frame_idx:04d}.jpg"
+ save_jpg(gt_pixels.detach().cpu().numpy(), gt_path)
+ if need_depth and depth_np is not None:
+ img_dir = out_root / "depth" / clip_id / model_name / cam_name / "images"
+ out_path = img_dir / f"{frame_idx:04d}.jpg"
+ save_depth_color(depth_np, opacity_np, out_path)
+
+
+@torch.no_grad()
+def render_novel_for_clip(
+ dataset: DrivingDataset,
+ trainer,
+ clip_id: str,
+ model_name: str,
+ out_root: Path,
+ device: torch.device,
+ cam_id: int = 0,
+) -> None:
+ """Render multiple standard novel trajectories for the front camera.
+
+ We treat each trajectory type as a "camera_name" under the 'novel' task.
+ Trajectories: s_curve, front_center_interp, lateral_offset, lateral_offset_left.
+ """
+ ds = dataset
+ traj_types = ["s_curve", "front_center_interp", "lateral_offset"]
+
+ # First three trajectories from dataset helper
+ trajs = ds.get_novel_render_traj(traj_types=traj_types, target_frames=ds.frame_num)
+
+ # Lateral offset left (-1m) trajectory constructed explicitly, mirroring batch_eval_nus
+ per_cam_poses: Dict[int, torch.Tensor] = {
+ c: ds.pixel_source.camera_data[c].cam_to_worlds for c in ds.pixel_source.camera_list
+ }
+ traj_left = lateral_offset_trajectory(
+ dataset_type=ds.type,
+ per_cam_poses=per_cam_poses,
+ original_frames=ds.frame_num,
+ target_frames=ds.frame_num,
+ offset_distance=-1.0,
+ )
+ trajs["lateral_offset_left"] = traj_left
+
+ for traj_type, traj in trajs.items():
+ base_dir = out_root / "novel" / clip_id / model_name / traj_type / "images"
+ render_list = ds.prepare_novel_view_render_data(traj, cam_id=cam_id)
+ for frame_idx, frame_data in enumerate(render_list):
+ for k, v in list(frame_data["cam_infos"].items()):
+ frame_data["cam_infos"][k] = v.to(device, non_blocking=True)
+ for k, v in list(frame_data["image_infos"].items()):
+ frame_data["image_infos"][k] = v.to(device, non_blocking=True)
+ out = trainer(
+ image_infos=frame_data["image_infos"],
+ camera_infos=frame_data["cam_infos"],
+ novel_view=True,
+ )
+ rgb = out["rgb"].detach().cpu().numpy()
+ out_path = base_dir / f"{frame_idx:04d}.jpg"
+ save_jpg(rgb, out_path)
+
+
+def parse_args() -> argparse.Namespace:
+ p = argparse.ArgumentParser("Render all clips for a model into per-camera images/videos + metrics.")
+ p.add_argument("model", type=str, help="Model directory name under work_dirs (e.g., drivestudio-nus-dist4d)")
+ p.add_argument("clip_id", nargs="*", help="Optional list of clip IDs. If empty, auto-discover clips under the model.")
+ p.add_argument(
+ "--tasks",
+ type=str,
+ nargs="*",
+ default=["train", "depth", "novel"],
+ choices=["train", "depth", "novel"],
+ help="Tasks to render (default: train depth novel)",
+ )
+ p.add_argument("--root", type=Path, default=Path("work_dirs"), help="Root for runs (default: work_dirs)")
+ p.add_argument("--output-root", type=Path, default=Path("data/render_all"), help="Output root directory")
+ p.add_argument("--device", type=str, default=None, help="Device (cuda:0 or cpu; default: auto)")
+ p.add_argument("--fps", type=int, default=10, help="FPS for output videos (default: 10)")
+ p.add_argument("--metrics-json", type=Path, default=Path("work_dirs_analysis_metrics.json"), help="LPIPS metrics JSON")
+ p.add_argument(
+ "--depth-csv",
+ type=Path,
+ nargs="*",
+ default=None,
+ help="Depth per-seq CSVs (default: depth_tables_3/per_seq.csv, depth_tables_2/per_seq.csv)",
+ )
+ p.add_argument("--absrel-column", type=str, default="AbsRel", help="Column name for AbsRel in depth CSV")
+ return p.parse_args()
+
+
+def main() -> None:
+ args = parse_args()
+ device = torch.device(args.device) if args.device else torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+ model_root = args.root / args.model
+ if not model_root.is_dir():
+ raise SystemExit(f"Model directory not found: {model_root}")
+
+ # Discover clips if not provided
+ if args.clip_id:
+ clips = list(args.clip_id)
+ else:
+ clips = sorted(d.name for d in model_root.iterdir() if d.is_dir() and (d / "config.yaml").is_file())
+ if not clips:
+ raise SystemExit(f"No clips with config.yaml found under {model_root}")
+
+ # Load metrics sources once
+ lpips_json = load_lpips_json(args.metrics_json)
+ depth_csvs = list(args.depth_csv) if args.depth_csv else [
+ Path("depth_tables_3/per_seq.csv"),
+ Path("depth_tables_2/per_seq.csv"),
+ ]
+ absrel_table = load_absrel_tables(depth_csvs, args.absrel_column)
+
+ out_root = args.output_root
+
+ for clip in clips:
+ clip_dir = model_root / clip
+ if not clip_dir.is_dir():
+ print(f"[SKIP] {clip_dir} is not a directory; skip")
+ continue
+ cfg = load_cfg(clip_dir)
+ dataset, trainer = build_dataset_trainer(cfg, device)
+
+ ckpt = find_checkpoint(clip_dir)
+ if ckpt is None:
+ print(f"[SKIP] clip {clip}: no checkpoint found in {clip_dir}")
+ continue
+ trainer.resume_from_checkpoint(str(ckpt), load_only_model=True)
+ trainer.set_eval()
+
+ # Render train/depth
+ render_train_depth_for_clip(
+ dataset=dataset,
+ trainer=trainer,
+ clip_id=clip,
+ model_name=args.model,
+ out_root=out_root,
+ tasks=args.tasks,
+ device=device,
+ )
+
+ # Render novel (front camera only)
+ if "novel" in args.tasks:
+ render_novel_for_clip(
+ dataset=dataset,
+ trainer=trainer,
+ clip_id=clip,
+ model_name=args.model,
+ out_root=out_root,
+ device=device,
+ cam_id=0,
+ )
+
+ # Build videos per camera for each task
+ for task in args.tasks:
+ task_root = out_root / task / clip / args.model
+ if not task_root.is_dir():
+ continue
+ for cam_dir in task_root.iterdir():
+ if not cam_dir.is_dir():
+ continue
+ img_dir = cam_dir / "images"
+ if not img_dir.is_dir():
+ continue
+ video_path = cam_dir / f"{task}.mp4"
+ write_video_from_images(img_dir, video_path, fps=args.fps)
+
+ # Metrics per task (one per clip+model, shared across cameras)
+ # Train metrics from JSON cache
+ if "train" in args.tasks:
+ train_metrics = lpips_json.get(args.model, {}).get(clip, {})
+ if isinstance(train_metrics, dict) and train_metrics:
+ task_root = out_root / "train" / clip / args.model
+ task_root.mkdir(parents=True, exist_ok=True)
+ metric_path = task_root / "train_metric.json"
+ metric_path.write_text(json.dumps(train_metrics, indent=2))
+
+ # Depth metrics from per_seq tables
+ if "depth" in args.tasks:
+ meth = method_label(args.model).replace(" ", "")
+ key1 = (meth, clip)
+ key2 = (meth.lower(), clip)
+ depth_metrics = absrel_table.get(key1) or absrel_table.get(key2)
+ if isinstance(depth_metrics, dict) and depth_metrics:
+ task_root = out_root / "depth" / clip / args.model
+ task_root.mkdir(parents=True, exist_ok=True)
+ metric_path = task_root / "depth_metric.json"
+ metric_path.write_text(json.dumps(depth_metrics, indent=2))
+
+ print(f"[DONE] clip {clip} for model {args.model} -> {out_root}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/third_party/drivestudio/tools/render_front_frames.py b/worldbench/third_party/drivestudio/tools/render_front_frames.py
new file mode 100644
index 0000000..757e55c
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/render_front_frames.py
@@ -0,0 +1,348 @@
+#!/usr/bin/env python3
+"""
+Render ORIGINAL per-frame RGB and Depth (no video decode) for a given clip and frame
+from:
+ - Training/original FRONT view
+ - Novel views: front_center_interp, lateral_offset, lateral_offset_left
+
+One output set per method (including GT) is saved as lossless PNG (+ raw depth .npy):
+ data/render_frames///
+ train_rgb_fXXXX.png
+ train_depth_fXXXX.png (COLOR visualization, same as batch_eval)
+ train_depth_fXXXX_raw.npy (only if --with-raw-npy)
+ train_depth_fXXXX_raw16cm.png (only if --save-raw-image; or *_raw16mm.png / _raw.exr)
+ train_depth_fXXXX_vis16.png (only if --save-raw-image and --with-depth-vis)
+ front_center_interp_rgb_fXXXX.png
+ front_center_interp_depth_fXXXX.png
+ front_center_interp_depth_fXXXX_raw.npy (only if --with-raw-npy and --with-front-center)
+ front_center_interp_depth_fXXXX_raw16cm.png (only if --save-raw-image)
+ front_center_interp_depth_fXXXX_vis16.png (only if --save-raw-image and --with-depth-vis)
+ lateral_offset_rgb_fXXXX.png
+ lateral_offset_depth_fXXXX.png
+ lateral_offset_depth_fXXXX_raw.npy (only if --with-raw-npy)
+ lateral_offset_depth_fXXXX_raw16cm.png (only if --save-raw-image)
+ lateral_offset_depth_fXXXX_vis16.png (only if --save-raw-image and --with-depth-vis)
+ lateral_offset_left_rgb_fXXXX.png
+ lateral_offset_left_depth_fXXXX.png
+ lateral_offset_left_depth_fXXXX_raw.npy (only if --with-raw-npy)
+ lateral_offset_left_depth_fXXXX_raw16cm.png (only if --save-raw-image)
+ lateral_offset_left_depth_fXXXX_vis16.png (only if --save-raw-image and --with-depth-vis)
+
+This script builds the dataset and trainer from each run's config.yaml under
+work_dirs/// and renders the requested single frame with the model,
+without passing through any video encoding. It matches the evaluation utilities
+in tools/batch_eval_nus.py and tools/eval.py.
+"""
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from PIL import Image
+from utils.visualization import depth_visualizer
+
+from datasets.driving_dataset import DrivingDataset
+from utils.misc import import_str
+from utils.camera import lateral_offset_trajectory
+
+
+DEFAULT_MODELS = [
+ "drivestudio-nus-gt",
+ "drivestudio-nus-dist4d",
+ "drivestudio-nus-drivedreamer2",
+ "drivestudio-nus-opendwm",
+ "drivestudio-nus-dreamforge",
+ "drivestudio-nus-xscene",
+ "drivestudio-nus-magicdrive",
+]
+
+
+def parse_args() -> argparse.Namespace:
+ p = argparse.ArgumentParser("Render original per-frame RGB/Depth for FRONT+novel views")
+ p.add_argument("clip_id", type=str, help="Sequence/clip id, e.g., 813")
+ p.add_argument("frame", type=int, help="Zero-based frame index (0..num_timesteps-1)")
+ p.add_argument("--root", type=Path, default=Path("work_dirs"), help="Root for runs (default: work_dirs)")
+ p.add_argument("--models", type=str, nargs="*", default=DEFAULT_MODELS, help="Model directory names (default: all 6m+GT)")
+ p.add_argument("--device", type=str, default=None, help="cuda device like 'cuda:0' or 'cpu'. Default: auto")
+ p.add_argument("--output-dir", type=Path, default=Path("data/render_frames"), help="Output root directory")
+ p.add_argument("--save-raw-image", action="store_true",
+ help="Also save raw depth as an image (disabled by default).")
+ p.add_argument("--depth-format", type=str, choices=["png16cm", "png16mm", "exr"], default="png16cm",
+ help="Raw depth image format when --save-raw-image is set: 16-bit PNG in centimeters (png16cm) or millimeters (png16mm), or EXR float32 (exr). Default: png16cm.")
+ p.add_argument("--with-depth-vis", action="store_true",
+ help="When --save-raw-image is set, also save a 16-bit grayscale visualization (_vis16.png). Disabled by default.")
+ p.add_argument("--with-raw-npy", action="store_true", help="Also save raw depth as .npy (_raw.npy). Disabled by default.")
+ p.add_argument("--with-front-center", action="store_true", help="Also render the 'front_center_interp' novel view. Disabled by default.")
+ p.add_argument("--novel-cam-id", type=int, default=0, help="Camera id for novel-view rendering reference (front=0 for nuScenes)")
+ p.add_argument("--cam-id", type=int, default=None, help="Alias for front camera id used for both train/original and novel views. If set, overrides --novel-cam-id.")
+ p.add_argument("--overwrite", action="store_true", help="Overwrite existing files")
+ p.add_argument("--quiet", action="store_true", help="Fewer prints")
+ return p.parse_args()
+
+
+def find_checkpoint(run_dir: Path) -> Optional[Path]:
+ final = run_dir / "checkpoint_final.pth"
+ if final.is_file():
+ return final
+ cands = sorted(run_dir.glob("checkpoint_*.pth"))
+ return cands[-1] if cands else None
+
+
+def load_cfg(run_dir: Path) -> OmegaConf:
+ cfg_path = run_dir / "config.yaml"
+ if not cfg_path.is_file():
+ raise SystemExit(f"config.yaml not found: {cfg_path}")
+ return OmegaConf.load(cfg_path)
+
+
+def build_dataset_trainer(cfg: OmegaConf, device: torch.device):
+ dataset = DrivingDataset(data_cfg=cfg.data)
+ trainer = import_str(cfg.trainer.type)(
+ **cfg.trainer,
+ num_timesteps=dataset.num_img_timesteps,
+ model_config=cfg.model,
+ num_train_images=len(dataset.train_image_set),
+ num_full_images=len(dataset.full_image_set),
+ test_set_indices=dataset.test_timesteps,
+ scene_aabb=dataset.get_aabb().reshape(2, 3),
+ device=device,
+ )
+ return dataset, trainer
+
+
+def save_png_uint8(rgb: np.ndarray, path: Path):
+ # rgb in [0,1], HxWx3
+ arr = (np.clip(rgb, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8)
+ Image.fromarray(arr).save(path)
+
+
+def save_depth_raw(depth_m: np.ndarray, path_stem: Path, fmt: str) -> Path:
+ """Save raw depth as image.
+ - png16cm: uint16 centimeters (depth[m]*100) clipped to [0,65535]
+ - png16mm: uint16 millimeters (depth[m]*1000) clipped to [0,65535]
+ - exr: float32 EXR (if supported by imageio)
+ Returns the written file path.
+ """
+ d = depth_m.astype(np.float32)
+ d[~np.isfinite(d)] = 0.0
+ if fmt == "png16cm":
+ arr = np.clip(d * 100.0, 0.0, 65535.0).astype(np.uint16)
+ out = path_stem.with_suffix("")
+ out = out.parent / (out.name + "_raw16cm.png")
+ Image.fromarray(arr).save(out)
+ return out
+ if fmt == "png16mm":
+ arr = np.clip(d * 1000.0, 0.0, 65535.0).astype(np.uint16)
+ out = path_stem.with_suffix("")
+ out = out.parent / (out.name + "_raw16mm.png")
+ Image.fromarray(arr).save(out)
+ return out
+ if fmt == "exr":
+ try:
+ import imageio
+ out = path_stem.with_suffix("")
+ out = out.parent / (out.name + "_raw.exr")
+ imageio.imwrite(out, d.astype(np.float32))
+ return out
+ except Exception:
+ # Fallback to png16cm if EXR writer unavailable
+ arr = np.clip(d * 100.0, 0.0, 65535.0).astype(np.uint16)
+ out = path_stem.with_suffix("")
+ out = out.parent / (out.name + "_raw16cm.png")
+ Image.fromarray(arr).save(out)
+ return out
+ raise ValueError(f"Unknown depth format: {fmt}")
+
+
+def save_depth_vis(depth_m: np.ndarray, path_stem: Path, clip_max: Optional[float] = None) -> Path:
+ d = depth_m.astype(np.float32)
+ d[~np.isfinite(d)] = 0.0
+ if clip_max is None:
+ mask = np.isfinite(d)
+ clip_max = float(np.percentile(d[mask], 99.0)) if mask.any() else 50.0
+ d = np.clip(d, 0.0, clip_max) / (clip_max + 1e-8)
+ png16 = (d * 65535.0 + 0.5).astype(np.uint16)
+ out = path_stem.with_suffix("")
+ out = out.parent / (out.name + "_vis16.png")
+ Image.fromarray(png16).save(out)
+ return out
+
+
+def save_depth_color(depth_m: np.ndarray, opacity: Optional[np.ndarray], path: Path) -> None:
+ """Save colorized depth like tools/batch_eval_nus.py (turbo colormap with -log curve).
+ Inputs are meters and an opacity mask in [0,1] (may be HxW or HxWx1)."""
+ d = depth_m
+ if d.ndim == 3 and d.shape[-1] == 1:
+ d = d[..., 0]
+ w = opacity
+ if w is not None and isinstance(w, np.ndarray) and w.ndim == 3 and w.shape[-1] == 1:
+ w = w[..., 0]
+ color = depth_visualizer(d, w)
+ arr = (np.clip(color, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8)
+ Image.fromarray(arr).save(path)
+
+
+def save_depth_npy(depth_m: np.ndarray, path_stem: Path) -> Path:
+ out = path_stem.with_suffix("")
+ out = out.parent / (out.name + "_raw.npy")
+ d = depth_m.astype(np.float32)
+ if d.ndim == 3 and d.shape[-1] == 1:
+ d = d[..., 0]
+ d[~np.isfinite(d)] = 0.0
+ np.save(out, d)
+ return out
+
+
+@torch.no_grad()
+def render_train_frame(dataset: DrivingDataset, trainer, frame_idx: int, cam_id: int, device: torch.device) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
+ # Front camera assumed cam_id = 0 for nuScenes; image index = t * num_cams + cam_id
+ num_cams = dataset.num_cams
+ if not (0 <= cam_id < num_cams):
+ raise SystemExit(f"Invalid cam_id {cam_id}; dataset has {num_cams} cameras")
+ img_idx = int(frame_idx) * num_cams + int(cam_id)
+ # render a single image via render_images with vis_indices
+ from models.video_utils import render_images
+ rr = render_images(trainer=trainer, dataset=dataset.full_image_set, compute_metrics=False, vis_indices=[img_idx])
+ rgb = rr["rgbs"][0]
+ depth = rr["depths"][0]
+ opacity = rr.get("opacities", [None])[0]
+ # rr stores numpy arrays already for depths/opacities; ensure numpy
+ if hasattr(rgb, "cpu"):
+ rgb = rgb.cpu().numpy()
+ return rgb, depth, opacity
+
+
+@torch.no_grad()
+def render_novel_frame(dataset: DrivingDataset, trainer, traj_type: str, frame_idx: int, cam_id: int, device: torch.device) -> Tuple[np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]:
+ # Prepare rendering data for the trajectory.
+ if traj_type == "lateral_offset_left":
+ # build left-shifted (-1 m) lateral trajectory explicitly
+ per_cam_poses: Dict[int, torch.Tensor] = {
+ c: dataset.pixel_source.camera_data[c].cam_to_worlds for c in dataset.pixel_source.camera_list
+ }
+ traj = lateral_offset_trajectory(
+ dataset_type=dataset.type,
+ per_cam_poses=per_cam_poses,
+ original_frames=dataset.frame_num,
+ target_frames=dataset.frame_num,
+ offset_distance=-1.0,
+ )
+ else:
+ trajs = dataset.get_novel_render_traj(traj_types=[traj_type], target_frames=dataset.frame_num)
+ if traj_type not in trajs:
+ raise SystemExit(f"Novel trajectory '{traj_type}' unavailable for dataset {dataset.type}")
+ traj = trajs[traj_type]
+
+ render_list = dataset.prepare_novel_view_render_data(traj, cam_id=cam_id)
+ if not (0 <= frame_idx < len(render_list)):
+ raise SystemExit(f"Frame {frame_idx} out of range for novel traj '{traj_type}' (0..{len(render_list)-1})")
+ fd = render_list[frame_idx]
+ # move to device
+ for k, v in list(fd["cam_infos"].items()):
+ fd["cam_infos"][k] = v.to(device, non_blocking=True)
+ for k, v in list(fd["image_infos"].items()):
+ fd["image_infos"][k] = v.to(device, non_blocking=True)
+ # render
+ out = trainer(image_infos=fd["image_infos"], camera_infos=fd["cam_infos"], novel_view=True)
+ rgb = out["rgb"].detach().cpu().numpy()
+ depth = out.get("depth", None)
+ opacity = out.get("opacity", None)
+ depth_np = depth.detach().cpu().numpy() if isinstance(depth, torch.Tensor) else None
+ opacity_np = opacity.detach().cpu().numpy() if isinstance(opacity, torch.Tensor) else None
+ return rgb, depth_np, opacity_np
+
+
+def main() -> None:
+ args = parse_args()
+ device = torch.device(args.device) if args.device else torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+ out_root = args.output_dir / args.clip_id
+ out_root.mkdir(parents=True, exist_ok=True)
+
+ methods = list(args.models)
+ if "drivestudio-nus-gt" not in methods:
+ methods.insert(0, "drivestudio-nus-gt")
+
+ # resolve camera id to use
+ cam_id = args.cam_id if args.cam_id is not None else args.novel_cam_id
+ # per method loop
+ for m in methods:
+ run_dir = args.root / m / args.clip_id
+ if not run_dir.is_dir():
+ print(f"[MISS] {run_dir} not found; skip {m}")
+ continue
+ cfg = load_cfg(run_dir)
+ dataset, trainer = build_dataset_trainer(cfg, device)
+ ckpt = find_checkpoint(run_dir)
+ if ckpt is None:
+ print(f"[MISS] checkpoint not found in {run_dir}; skip {m}")
+ continue
+ trainer.resume_from_checkpoint(str(ckpt), load_only_model=True)
+ trainer.set_eval()
+
+ # bounds check on frame
+ if not (0 <= args.frame < int(dataset.num_img_timesteps)):
+ print(f"[SKIP] {m}: frame {args.frame} out of [0,{dataset.num_img_timesteps-1}]")
+ continue
+
+ out_dir = out_root / m
+ out_dir.mkdir(parents=True, exist_ok=True)
+ ftag = f"f{args.frame:04d}"
+
+ # 1) train/original FRONT
+ try:
+ rgb, depth, opac = render_train_frame(dataset, trainer, args.frame, cam_id, device)
+ if not args.quiet:
+ print(f"[OK] {m} train frame {args.frame}")
+ save_png_uint8(rgb, out_dir / f"train_rgb_{ftag}.png")
+ if depth is not None:
+ # colorized depth (primary visualization)
+ save_depth_color(depth, opac, out_dir / f"train_depth_{ftag}.png")
+ # raw depth npy (default enabled)
+ stem = out_dir / f"train_depth_{ftag}.png"
+ if args.with_raw_npy:
+ save_depth_npy(depth, stem)
+ # optional raw depth image and optional grayscale vis
+ if args.save_raw_image:
+ save_depth_raw(depth.squeeze(-1), stem, args.depth_format)
+ if args.with_depth_vis:
+ save_depth_vis(depth.squeeze(-1), stem)
+ except Exception as e:
+ print(f"[WARN] {m} train render failed: {e}")
+
+ # 2) novel views
+ novel_views: List[str] = []
+ if args.with_front_center:
+ novel_views.append("front_center_interp")
+ novel_views += ["lateral_offset", "lateral_offset_left"]
+ for v in novel_views:
+ try:
+ rgb, depth, opac = render_novel_frame(dataset, trainer, v, args.frame, cam_id, device)
+ save_png_uint8(rgb, out_dir / f"{v}_rgb_{ftag}.png")
+ if depth is not None:
+ # colorized depth (primary visualization)
+ save_depth_color(depth, opac, out_dir / f"{v}_depth_{ftag}.png")
+ # raw npy (default)
+ stem = out_dir / f"{v}_depth_{ftag}.png"
+ if args.with_raw_npy:
+ save_depth_npy(depth, stem)
+ # optional raw image + grayscale vis
+ if args.save_raw_image:
+ save_depth_raw(depth.squeeze(-1), stem, args.depth_format)
+ if args.with_depth_vis:
+ save_depth_vis(depth.squeeze(-1), stem)
+ if not args.quiet:
+ print(f"[OK] {m} {v} frame {args.frame}")
+ except Exception as e:
+ print(f"[WARN] {m} {v} render failed: {e}")
+
+ print(f"[DONE] saved frames to {out_root}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/third_party/drivestudio/tools/select_demo_sequences.py b/worldbench/third_party/drivestudio/tools/select_demo_sequences.py
new file mode 100644
index 0000000..651a573
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/select_demo_sequences.py
@@ -0,0 +1,301 @@
+#!/usr/bin/env python3
+"""
+Select demo-friendly sequences where methods differ most on LPIPS (RGB render,
+original view) and AbsRel (depth), computed per-sequence.
+
+Inputs the aggregated image metrics JSON (from tools/metrics_analysis.py
+--metrics-json) and per-sequence depth CSV tables (from tools/eval_depth_diff.py
+--table-csv-prefix /per_seq.csv). Produces ranked lists and CSVs.
+
+Usage example:
+ python3 tools/select_demo_sequences.py \
+ --metrics-json work_dirs_analysis_metrics.json \
+ --depth-per-seq depth_tables_3/per_seq.csv depth_tables_2/per_seq.csv \
+ --top-k 15 \
+ --output-dir demo_tables
+
+This prints Top-K tables for LPIPS and AbsRel, and writes CSVs under output-dir.
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+from dataclasses import dataclass
+import math
+from pathlib import Path
+from typing import Dict, Iterable, List, Mapping, Optional, Sequence, Tuple
+
+
+PREFIX = "drivestudio-nus-"
+
+
+@dataclass
+class Spread:
+ seq: str
+ best_method: str
+ best_value: float
+ worst_method: str
+ worst_value: float
+ delta: float
+ n_methods: int
+
+
+def load_lpips_per_seq(metrics_json: Path, include_methods: Optional[Sequence[str]] = None) -> Dict[str, Dict[str, float]]:
+ """Return mapping: method -> {seq -> LPIPS}.
+
+ Expects the top-level keys like 'drivestudio-nus-'.
+ """
+ payload = json.loads(metrics_json.read_text())
+ out: Dict[str, Dict[str, float]] = {}
+ for top_key, clip_map in payload.items():
+ if not top_key.startswith(PREFIX):
+ continue
+ method = top_key[len(PREFIX) :]
+ if include_methods is not None and method not in include_methods:
+ continue
+ seq_to_lpips: Dict[str, float] = {}
+ for seq, metrics in clip_map.items():
+ v = metrics.get("image_metrics/full/lpips")
+ if v is not None:
+ seq_to_lpips[seq] = float(v)
+ if seq_to_lpips:
+ out[method] = seq_to_lpips
+ return out
+
+
+def _read_csv(path: Path) -> List[List[str]]:
+ with path.open("r", newline="") as f:
+ return list(csv.reader(f))
+
+
+def load_absrel_per_seq(depth_per_seq_files: Sequence[Path], include_methods: Optional[Sequence[str]] = None) -> Dict[str, Dict[str, float]]:
+ """Return mapping: method -> {seq -> AbsRel} from one or more per_seq.csv files."""
+ out: Dict[str, Dict[str, float]] = {}
+ for csv_path in depth_per_seq_files:
+ rows = _read_csv(csv_path)
+ if not rows:
+ continue
+ header = rows[0]
+ col_map = {name: idx for idx, name in enumerate(header)}
+ try:
+ idx_method_seq = col_map["Method/Seq"]
+ idx_absrel = col_map["AbsRel"]
+ except KeyError:
+ # Not a per_seq.csv table; skip
+ continue
+ for row in rows[1:]:
+ if not row or len(row) <= max(idx_method_seq, idx_absrel):
+ continue
+ label = row[idx_method_seq].strip()
+ if "/" not in label:
+ continue
+ method, seq = label.split("/", 1)
+ if include_methods is not None and method not in include_methods:
+ continue
+ try:
+ v = float(row[idx_absrel])
+ except ValueError:
+ continue
+ out.setdefault(method, {})[seq] = v
+ return out
+
+
+def _quantiles(sorted_vals: List[float], q: float) -> float:
+ """Linear interpolation quantile for small sample size without numpy.
+
+ q in [0,1]. Assumes sorted_vals non-empty and sorted.
+ """
+ n = len(sorted_vals)
+ if n == 1:
+ return sorted_vals[0]
+ pos = (n - 1) * q
+ lo = int(math.floor(pos))
+ hi = int(math.ceil(pos))
+ if lo == hi:
+ return sorted_vals[lo]
+ frac = pos - lo
+ return sorted_vals[lo] * (1 - frac) + sorted_vals[hi] * frac
+
+
+def compute_spread(
+ per_seq: Mapping[str, Mapping[str, float]],
+ min_methods: int = 2,
+ score: str = "range",
+ trim: float = 0.1,
+) -> List[Spread]:
+ """Compute a sequence-level dispersion score across methods.
+
+ score ∈ {"range", "iqr", "std", "mad", "pairwise", "trimmed"}:
+ - range: max - min
+ - iqr: Q3 - Q1 (robust to single outliers)
+ - std: population std
+ - mad: median absolute deviation (unscaled)
+ - pairwise: mean absolute difference across all method pairs
+ - trimmed: P(1-trim) - P(trim), e.g., 90th - 10th percentiles
+ """
+ # sequences -> list of (method, value)
+ seq_to_vals: Dict[str, List[Tuple[str, float]]] = {}
+ for method, seq_map in per_seq.items():
+ for seq, v in seq_map.items():
+ seq_to_vals.setdefault(seq, []).append((method, v))
+ spreads: List[Spread] = []
+ for seq, items in seq_to_vals.items():
+ if len(items) < min_methods:
+ continue
+ items_sorted = sorted(items, key=lambda x: x[1])
+ best_method, best_val = items_sorted[0]
+ worst_method, worst_val = items_sorted[-1]
+ vals = [v for _, v in items_sorted]
+ m = len(vals)
+ # default: range
+ if score == "range":
+ delta = worst_val - best_val
+ elif score == "iqr":
+ q1 = _quantiles(vals, 0.25)
+ q3 = _quantiles(vals, 0.75)
+ delta = q3 - q1
+ elif score == "std":
+ mu = sum(vals) / m
+ delta = math.sqrt(sum((v - mu) ** 2 for v in vals) / m)
+ elif score == "mad":
+ med = _quantiles(vals, 0.5)
+ devs = sorted(abs(v - med) for v in vals)
+ delta = _quantiles(devs, 0.5)
+ elif score == "pairwise":
+ if m == 1:
+ delta = 0.0
+ else:
+ s = 0.0
+ cnt = 0
+ for i in range(m):
+ for j in range(i + 1, m):
+ s += abs(vals[i] - vals[j])
+ cnt += 1
+ delta = s / cnt
+ elif score == "trimmed":
+ lo = _quantiles(vals, trim)
+ hi = _quantiles(vals, 1.0 - trim)
+ delta = hi - lo
+ else:
+ raise SystemExit(f"Unknown score '{score}'.")
+ spreads.append(
+ Spread(
+ seq=seq,
+ best_method=best_method,
+ best_value=best_val,
+ worst_method=worst_method,
+ worst_value=worst_val,
+ delta=delta,
+ n_methods=m,
+ )
+ )
+ spreads.sort(key=lambda s: (s.delta, s.seq), reverse=True)
+ return spreads
+
+
+def intersect_methods(a: Mapping[str, Mapping[str, float]], b: Mapping[str, Mapping[str, float]]) -> List[str]:
+ return sorted(set(a.keys()) & set(b.keys()))
+
+
+def filter_methods(per_seq: Dict[str, Dict[str, float]], allowed: Sequence[str]) -> Dict[str, Dict[str, float]]:
+ return {m: d for m, d in per_seq.items() if m in allowed}
+
+
+def print_table(title: str, spreads: Sequence[Spread], top_k: int) -> None:
+ print(title)
+ header = ["Seq", "Δ", "Best(method,val)", "Worst(method,val)", "#M"]
+ print(" ".join(header))
+ for s in spreads[:top_k]:
+ best = f"{s.best_method},{s.best_value:.4f}"
+ worst = f"{s.worst_method},{s.worst_value:.4f}"
+ print(f"{s.seq} {s.delta:.4f} {best} {worst} {s.n_methods}")
+ print()
+
+
+def write_csv(path: Path, spreads: Sequence[Spread]) -> None:
+ path.parent.mkdir(parents=True, exist_ok=True)
+ with path.open("w", newline="") as f:
+ w = csv.writer(f)
+ w.writerow(["seq", "delta", "best_method", "best_value", "worst_method", "worst_value", "n_methods"])
+ for s in spreads:
+ w.writerow([s.seq, f"{s.delta:.6f}", s.best_method, f"{s.best_value:.6f}", s.worst_method, f"{s.worst_value:.6f}", s.n_methods])
+
+
+def parse_args() -> argparse.Namespace:
+ p = argparse.ArgumentParser(description="Select sequences with largest per-seq metric gaps across methods.")
+ p.add_argument("--metrics-json", type=Path, default=Path("work_dirs_analysis_metrics.json"))
+ p.add_argument("--depth-per-seq", type=Path, nargs="*", default=[Path("depth_tables_3/per_seq.csv")])
+ p.add_argument("--methods", type=str, nargs="*", default=None, help="Restrict to these method names.")
+ p.add_argument("--top-k", type=int, default=15)
+ p.add_argument("--score-lpips", type=str, default="range", choices=["range","iqr","std","mad","pairwise","trimmed"], help="Dispersion score for LPIPS.")
+ p.add_argument("--score-absrel", type=str, default="range", choices=["range","iqr","std","mad","pairwise","trimmed"], help="Dispersion score for AbsRel.")
+ p.add_argument("--trim", type=float, default=0.1, help="Trim fraction for 'trimmed' score (e.g., 0.1 → P90-P10).")
+ p.add_argument("--min-methods", type=int, default=2, help="Min methods per sequence to consider.")
+ p.add_argument("--output-dir", type=Path, default=Path("demo_tables"))
+ return p.parse_args()
+
+
+def main() -> None:
+ args = parse_args()
+
+ # Load raw per-seq metrics
+ lpips_map = load_lpips_per_seq(args.metrics_json, include_methods=args.methods)
+ absrel_map = load_absrel_per_seq(args.depth_per_seq, include_methods=args.methods)
+
+ if not lpips_map:
+ raise SystemExit("No LPIPS data found (check --metrics-json).")
+ if not absrel_map:
+ raise SystemExit("No AbsRel per-seq data found (check --depth-per-seq).")
+
+ # Use common methods if not specified
+ if args.methods is None:
+ common_methods = intersect_methods(lpips_map, absrel_map)
+ if not common_methods:
+ raise SystemExit("No common methods between LPIPS JSON and AbsRel CSVs.")
+ lpips_map = filter_methods(lpips_map, common_methods)
+ absrel_map = filter_methods(absrel_map, common_methods)
+
+ # Compute spreads with requested scores
+ lpips_spreads = compute_spread(lpips_map, min_methods=args.min_methods, score=args.score_lpips, trim=args.trim)
+ absrel_spreads = compute_spread(absrel_map, min_methods=args.min_methods, score=args.score_absrel, trim=args.trim)
+
+ print_table(f"Top sequences by LPIPS gap [{args.score_lpips}] (lower better):", lpips_spreads, args.top_k)
+ print_table(f"Top sequences by AbsRel gap [{args.score_absrel}] (lower better):", absrel_spreads, args.top_k)
+
+ # Write CSVs
+ outdir = args.output_dir
+ write_csv(outdir / "lpips_spread.csv", lpips_spreads)
+ write_csv(outdir / "absrel_spread.csv", absrel_spreads)
+
+ # Intersection candidates: rank by min rank across both lists
+ rank_lp = {s.seq: i for i, s in enumerate(lpips_spreads)}
+ rank_ab = {s.seq: i for i, s in enumerate(absrel_spreads)}
+ both = []
+ for seq in set(rank_lp) & set(rank_ab):
+ r = max(rank_lp[seq], rank_ab[seq]) # conservative: good in both lists
+ both.append((r, seq))
+ both.sort()
+ top_both = [seq for _, seq in both[: args.top_k]]
+ if top_both:
+ print("Top sequences that are high gap in both LPIPS and AbsRel:")
+ print(", ".join(top_both))
+ print()
+ # Write a small CSV joining both spreads for convenience
+ join_rows: List[List[str]] = [["seq", "lpips_delta", "absrel_delta"]]
+ lp_map = {s.seq: s for s in lpips_spreads}
+ ab_map = {s.seq: s for s in absrel_spreads}
+ for seq in top_both:
+ join_rows.append([
+ seq,
+ f"{lp_map[seq].delta:.6f}",
+ f"{ab_map[seq].delta:.6f}",
+ ])
+ out = outdir / "both_spread_top.csv"
+ out.parent.mkdir(parents=True, exist_ok=True)
+ with out.open("w", newline="") as f:
+ csv.writer(f).writerows(join_rows)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/third_party/drivestudio/tools/show_clip_metrics.py b/worldbench/third_party/drivestudio/tools/show_clip_metrics.py
new file mode 100644
index 0000000..88bda51
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/show_clip_metrics.py
@@ -0,0 +1,122 @@
+#!/usr/bin/env python3
+"""Show per-clip train/depth metrics (LPIPS, AbsRel) for selected models."""
+from __future__ import annotations
+
+import argparse
+import json
+import csv
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+
+DEFAULT_MODELS = [
+ "drivestudio-nus-gt",
+ "drivestudio-nus-dist4d",
+ "drivestudio-nus-drivedreamer2",
+ "drivestudio-nus-opendwm",
+ "drivestudio-nus-dreamforge",
+ "drivestudio-nus-xscene",
+ "drivestudio-nus-magicdrive",
+]
+
+
+def parse_args() -> argparse.Namespace:
+ p = argparse.ArgumentParser("Print LPIPS/AbsRel metrics for specified clip IDs")
+ p.add_argument("clip_id", nargs="+", help="Clip IDs (e.g., 813 694)")
+ p.add_argument("--models", type=str, nargs="*", default=None, help="Model directory names. Default: all known models containing the clip in metrics/depth tables.")
+ p.add_argument("--metrics-json", type=Path, default=Path("work_dirs_analysis_metrics.json"), help="Path to LPIPS metrics JSON (default: work_dirs_analysis_metrics.json)")
+ p.add_argument("--metric-key", type=str, default="image_metrics/full/lpips", help="Metric key inside JSON (default: image_metrics/full/lpips)")
+ p.add_argument("--depth-csv", type=Path, nargs="*", default=None, help="Candidate depth per-seq CSVs. Default: depth_tables_3/per_seq.csv, depth_tables_2/per_seq.csv")
+ p.add_argument("--absrel-column", type=str, default="AbsRel", help="Column name for AbsRel in depth CSV (default: AbsRel)")
+ return p.parse_args()
+
+
+def method_label(name: str) -> str:
+ return name.split("drivestudio-nus-", 1)[-1].replace("_", " ").replace("-", " ")
+
+
+def load_lpips(path: Path) -> Dict[str, Dict[str, Dict[str, float]]]:
+ if not path.is_file():
+ raise SystemExit(f"LPIPS metrics JSON not found: {path}")
+ try:
+ return json.loads(path.read_text())
+ except Exception as e:
+ raise SystemExit(f"Failed to parse {path}: {e}")
+
+
+def load_absrel(paths: List[Path], col: str) -> Dict[Tuple[str, str], float]:
+ lookup: Dict[Tuple[str, str], float] = {}
+ for p in paths:
+ if not p.is_file():
+ continue
+ try:
+ with p.open("r", newline="") as f:
+ reader = csv.DictReader(f)
+ if "Method/Seq" not in reader.fieldnames or col not in reader.fieldnames:
+ continue
+ for row in reader:
+ tag = row.get("Method/Seq", "").strip()
+ val = row.get(col)
+ if not tag or val is None:
+ continue
+ if "/" not in tag:
+ continue
+ method, seq = tag.split("/", 1)
+ try:
+ lookup[(method.strip(), seq.strip())] = float(val)
+ except ValueError:
+ continue
+ except Exception as e:
+ print(f"[WARN] failed reading {p}: {e}")
+ return lookup
+
+
+def main() -> None:
+ args = parse_args()
+ metrics = load_lpips(args.metrics_json)
+ depth_paths = args.depth_csv
+ if not depth_paths:
+ depth_paths = [Path("depth_tables_3/per_seq.csv"), Path("depth_tables_2/per_seq.csv")]
+ absrel_lookup = load_absrel(depth_paths, args.absrel_column)
+
+ for clip in args.clip_id:
+ # determine models
+ if args.models:
+ models = list(args.models)
+ else:
+ models = set()
+ for model, seqs in metrics.items():
+ if isinstance(seqs, dict) and clip in seqs:
+ models.add(model)
+ for (method, seq), _ in absrel_lookup.items():
+ if seq == clip:
+ models.add(f"drivestudio-nus-{method.lower()}")
+ if not models:
+ models = set(DEFAULT_MODELS)
+ models = sorted(models)
+
+ print(f"\nClip {clip}")
+ print(f"{'Model':<25} {'LPIPS':>10} {'AbsRel':>10}")
+ print("-" * 50)
+ for model in models:
+ lpips = None
+ seq_map = metrics.get(model, {}) if isinstance(metrics.get(model), dict) else {}
+ if isinstance(seq_map, dict):
+ val = seq_map.get(clip, {}).get(args.metric_key)
+ if isinstance(val, (int, float)):
+ lpips = val
+ meth_key = method_label(model).replace(" ", "")
+ absrel = None
+ key1 = (meth_key, clip)
+ key2 = (meth_key.lower(), clip)
+ if key1 in absrel_lookup:
+ absrel = absrel_lookup[key1]
+ elif key2 in absrel_lookup:
+ absrel = absrel_lookup[key2]
+ lpips_str = f"{lpips:.3f}" if lpips is not None else "-"
+ absrel_str = f"{absrel:.3f}" if absrel is not None else "-"
+ print(f"{method_label(model):<25} {lpips_str:>10} {absrel_str:>10}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/third_party/drivestudio/tools/train.py b/worldbench/third_party/drivestudio/tools/train.py
new file mode 100644
index 0000000..4268286
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/train.py
@@ -0,0 +1,543 @@
+from omegaconf import OmegaConf
+import numpy as np
+import os
+import time
+import wandb
+import random
+import imageio
+import logging
+import argparse
+import socket
+import atexit
+
+import torch
+from tools.eval import do_evaluation
+from utils.misc import import_str
+from utils.backup import backup_project
+from utils.logging import MetricLogger, setup_logging
+from models.video_utils import render_images, save_videos
+from datasets.driving_dataset import DrivingDataset
+
+logger = logging.getLogger()
+current_time = time.strftime("%Y-%m-%d_%H-%M-%S", time.localtime())
+
+def set_seeds(seed=31):
+ """
+ Fix random seeds.
+ """
+ torch.manual_seed(seed)
+ torch.cuda.manual_seed_all(seed)
+ np.random.seed(seed)
+ random.seed(seed)
+
+def setup(args):
+ # get config
+ cfg = OmegaConf.load(args.config_file)
+
+ # parse datasets
+ args_from_cli = OmegaConf.from_cli(args.opts)
+ if "dataset" in args_from_cli:
+ cfg.dataset = args_from_cli.pop("dataset")
+
+ assert "dataset" in cfg or "data" in cfg, \
+ "Please specify dataset in config or data in config"
+
+ if "dataset" in cfg:
+ dataset_type = cfg.pop("dataset")
+ dataset_cfg = OmegaConf.load(
+ os.path.join("configs", "datasets", f"{dataset_type}.yaml")
+ )
+ # merge data
+ cfg = OmegaConf.merge(cfg, dataset_cfg)
+
+ # merge cli
+ cfg = OmegaConf.merge(cfg, args_from_cli)
+ log_dir = os.path.join(args.output_root, args.project, args.run_name)
+
+ # update config and create log dir
+ cfg.log_dir = log_dir
+ os.makedirs(log_dir, exist_ok=True)
+ for folder in ["images", "videos", "metrics", "configs_bk", "buffer_maps", "backup"]:
+ os.makedirs(os.path.join(log_dir, folder), exist_ok=True)
+
+ # setup wandb
+ if args.enable_wandb:
+ # sometimes wandb fails to init in cloud machines, so we give it several (many) tries
+ while (
+ wandb.init(
+ project=args.project,
+ entity=args.entity,
+ sync_tensorboard=True,
+ settings=wandb.Settings(start_method="fork"),
+ )
+ is not wandb.run
+ ):
+ continue
+ wandb.run.name = args.run_name
+ wandb.run.save()
+ wandb.config.update(OmegaConf.to_container(cfg, resolve=True))
+ wandb.config.update(args)
+
+ # setup random seeds
+ set_seeds(cfg.seed)
+
+ global logger
+ setup_logging(output=log_dir, level=logging.INFO, time_string=current_time)
+ logger.info("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items())))
+
+ # save config
+ logger.info(f"Config:\n{OmegaConf.to_yaml(cfg)}")
+ saved_cfg_path = os.path.join(log_dir, "config.yaml")
+ with open(saved_cfg_path, "w") as f:
+ OmegaConf.save(config=cfg, f=f)
+
+ # also save a backup copy
+ saved_cfg_path_bk = os.path.join(log_dir, "configs_bk", f"config_{current_time}.yaml")
+ with open(saved_cfg_path_bk, "w") as f:
+ OmegaConf.save(config=cfg, f=f)
+ logger.info(f"Full config saved to {saved_cfg_path}, and {saved_cfg_path_bk}")
+
+ # Backup codes
+ backup_project(
+ os.path.join(log_dir, 'backup'), "./",
+ ["configs", "datasets", "models", "utils", "tools"],
+ [".py", ".h", ".cpp", ".cuh", ".cu", ".sh", ".yaml"]
+ )
+ return cfg
+
+def train(cfg, device):
+
+ # build dataset
+ dataset = DrivingDataset(data_cfg=cfg.data)
+
+ # setup trainer
+ trainer = import_str(cfg.trainer.type)(
+ **cfg.trainer,
+ num_timesteps=dataset.num_img_timesteps,
+ model_config=cfg.model,
+ num_train_images=len(dataset.train_image_set),
+ num_full_images=len(dataset.full_image_set),
+ test_set_indices=dataset.test_timesteps,
+ scene_aabb=dataset.get_aabb().reshape(2, 3),
+ device=device
+ )
+
+ # NOTE: If resume, gaussians will be loaded from checkpoint
+ # If not, gaussians will be initialized from dataset
+ if args.resume_from is not None:
+ trainer.resume_from_checkpoint(
+ ckpt_path=args.resume_from,
+ load_only_model=True
+ )
+ logger.info(
+ f"Resuming training from {args.resume_from}, starting at step {trainer.step}"
+ )
+ else:
+ trainer.init_gaussians_from_dataset(dataset=dataset)
+ logger.info(
+ f"Training from scratch, initializing gaussians from dataset, starting at step {trainer.step}"
+ )
+
+ if args.enable_viewer:
+ # a simple viewer for background visualization
+ trainer.init_viewer(port=args.viewer_port)
+
+ # define render keys
+ render_keys = [
+ "gt_rgbs",
+ "rgbs",
+ "Background_rgbs",
+ "Dynamic_rgbs",
+ "RigidNodes_rgbs",
+ "DeformableNodes_rgbs",
+ "SMPLNodes_rgbs",
+ # "depths",
+ # "Background_depths",
+ # "Dynamic_depths",
+ # "RigidNodes_depths",
+ # "DeformableNodes_depths",
+ # "SMPLNodes_depths",
+ # "mask"
+ ]
+ if cfg.render.vis_lidar:
+ render_keys.insert(0, "lidar_on_images")
+ if cfg.render.vis_sky:
+ render_keys += ["rgb_sky_blend", "rgb_sky"]
+ if cfg.render.vis_error:
+ render_keys.insert(render_keys.index("rgbs") + 1, "rgb_error_maps")
+
+ # setup optimizer
+ trainer.initialize_optimizer()
+
+ # setup metric logger
+ metrics_file = os.path.join(cfg.log_dir, "metrics.json")
+ metric_logger = MetricLogger(delimiter=" ", output_file=metrics_file)
+ all_iters = np.arange(trainer.step, trainer.num_iters + 1)
+
+ # DEBUG USE
+ # do_evaluation(
+ # step=0,
+ # cfg=cfg,
+ # trainer=trainer,
+ # dataset=dataset,
+ # render_keys=render_keys,
+ # args=args,
+ # )
+
+ for step in metric_logger.log_every(all_iters, cfg.logging.print_freq):
+ #----------------------------------------------------------------------------
+ #---------------------------- Validate ------------------------------
+ if step % cfg.logging.vis_freq == 0 and cfg.logging.vis_freq > 0:
+ logger.info("Visualizing...")
+ vis_timestep = np.linspace(
+ 0,
+ dataset.num_img_timesteps,
+ trainer.num_iters // cfg.logging.vis_freq + 1,
+ endpoint=False,
+ dtype=int,
+ )[step // cfg.logging.vis_freq]
+ with torch.no_grad():
+ render_results = render_images(
+ trainer=trainer,
+ dataset=dataset.full_image_set,
+ compute_metrics=True,
+ compute_error_map=cfg.render.vis_error,
+ vis_indices=[
+ vis_timestep * dataset.pixel_source.num_cams + i
+ for i in range(dataset.pixel_source.num_cams)
+ ],
+ )
+ if args.enable_wandb:
+ wandb.log(
+ {
+ "image_metrics/psnr": render_results["psnr"],
+ "image_metrics/ssim": render_results["ssim"],
+ "image_metrics/occupied_psnr": render_results["occupied_psnr"],
+ "image_metrics/occupied_ssim": render_results["occupied_ssim"],
+ }
+ )
+ vis_frame_dict = save_videos(
+ render_results,
+ save_pth=os.path.join(
+ cfg.log_dir, "images", f"step_{step}.png"
+ ), # don't save the video
+ layout=dataset.layout,
+ num_timestamps=1,
+ keys=render_keys,
+ save_seperate_video=cfg.logging.save_seperate_video,
+ num_cams=dataset.pixel_source.num_cams,
+ fps=cfg.render.fps,
+ verbose=False,
+ )
+ if args.enable_wandb:
+ for k, v in vis_frame_dict.items():
+ wandb.log({"image_rendering/" + k: wandb.Image(v)})
+ del render_results
+ torch.cuda.empty_cache()
+
+
+ #----------------------------------------------------------------------------
+ #---------------------------- training step -------------------------------
+ # prepare for training
+ trainer.set_train()
+ trainer.preprocess_per_train_step(step=step)
+ trainer.optimizer_zero_grad() # zero grad
+
+ # get data
+ train_step_camera_downscale = trainer._get_downscale_factor()
+ image_infos, cam_infos = dataset.train_image_set.next(train_step_camera_downscale)
+ for k, v in image_infos.items():
+ if isinstance(v, torch.Tensor):
+ image_infos[k] = v.cuda(non_blocking=True)
+ for k, v in cam_infos.items():
+ if isinstance(v, torch.Tensor):
+ cam_infos[k] = v.cuda(non_blocking=True)
+
+ # forward & backward
+ outputs = trainer(image_infos, cam_infos)
+ trainer.update_visibility_filter()
+
+ loss_dict = trainer.compute_losses(
+ outputs=outputs,
+ image_infos=image_infos,
+ cam_infos=cam_infos,
+ )
+ # check nan or inf
+ # for k, v in loss_dict.items():
+ # if torch.isnan(v).any():
+ # raise ValueError(f"NaN detected in loss {k} at step {step}")
+ # if torch.isinf(v).any():
+ # raise ValueError(f"Inf detected in loss {k} at step {step}")
+
+ # TODO: 用生成数据训练的时候Background_sharp_shape_reg会出现nan, 需要check
+ loss_dict_keys = list(loss_dict.keys())
+ for k in loss_dict_keys:
+ if torch.isnan(loss_dict[k]).any() or torch.isinf(loss_dict[k]).any():
+ loss_dict.pop(k)
+ trainer.backward(loss_dict)
+
+ # after training step
+ trainer.postprocess_per_train_step(step=step)
+
+ #----------------------------------------------------------------------------
+ #------------------------------- logging ----------------------------------
+ with torch.no_grad():
+ # cal stats
+ metric_dict = trainer.compute_metrics(
+ outputs=outputs,
+ image_infos=image_infos,
+ )
+ metric_logger.update(**{"train_metrics/"+k: v.item() for k, v in metric_dict.items()})
+ metric_logger.update(**{"train_stats/gaussian_num_" + k: v for k, v in trainer.get_gaussian_count().items()})
+ metric_logger.update(**{"losses/"+k: v.item() for k, v in loss_dict.items()})
+ metric_logger.update(**{"train_stats/lr_" + group['name']: group['lr'] for group in trainer.optimizer.param_groups})
+ if args.enable_wandb:
+ wandb.log({k: v.avg for k, v in metric_logger.meters.items()})
+
+ #----------------------------------------------------------------------------
+ #---------------------------- Saving --------------------------------
+ do_save = step > 0 and (
+ (step % cfg.logging.saveckpt_freq == 0) or (step == trainer.num_iters)
+ ) and (args.resume_from is None)
+ if do_save:
+ trainer.save_checkpoint(
+ log_dir=cfg.log_dir,
+ save_only_model=True,
+ is_final=step == trainer.num_iters,
+ )
+
+ #----------------------------------------------------------------------------
+ #------------------------ Cache Image Error ---------------------------
+ if (
+ step > 0 and trainer.optim_general.cache_buffer_freq > 0
+ and step % trainer.optim_general.cache_buffer_freq == 0
+ ):
+ logger.info("Caching image error...")
+ trainer.set_eval()
+ with torch.no_grad():
+ dataset.pixel_source.update_downscale_factor(
+ 1 / dataset.pixel_source.buffer_downscale
+ )
+ render_results = render_images(
+ trainer=trainer,
+ dataset=dataset.full_image_set,
+ )
+ dataset.pixel_source.reset_downscale_factor()
+ dataset.pixel_source.update_image_error_maps(render_results)
+
+ # save error maps
+ merged_error_video = dataset.pixel_source.get_image_error_video(
+ dataset.layout
+ )
+ imageio.mimsave(
+ os.path.join(
+ cfg.log_dir, "buffer_maps", f"buffer_maps_{step}.mp4"
+ ),
+ merged_error_video,
+ fps=cfg.render.fps,
+ )
+ logger.info("Done caching rgb error maps")
+
+
+ logger.info("Training done!")
+
+ do_evaluation(
+ step=step,
+ cfg=cfg,
+ trainer=trainer,
+ dataset=dataset,
+ render_keys=render_keys,
+ args=args,
+ )
+
+ if args.enable_viewer:
+ print("Viewer running... Ctrl+C to exit.")
+ time.sleep(1000000)
+
+ return step
+
+def is_data_ready(scene_dir):
+ return 'fine_dynamic_masks' in os.listdir(scene_dir) and \
+ 'humanpose' in os.listdir(scene_dir) and \
+ len(os.listdir(os.path.join(scene_dir, 'fine_dynamic_masks', 'all'))) == len(os.listdir(os.path.join(scene_dir, 'images'))) and \
+ len(os.listdir(os.path.join(scene_dir, 'fine_dynamic_masks', 'human'))) == len(os.listdir(os.path.join(scene_dir, 'images'))) and \
+ len(os.listdir(os.path.join(scene_dir, 'fine_dynamic_masks', 'vehicle'))) == len(os.listdir(os.path.join(scene_dir, 'images'))) and \
+ os.path.exists(os.path.join(scene_dir, 'humanpose', 'smpl.pkl'))
+
+# ---------------------------------------------------------------
+# Simple file-based lock for multi-process scene scheduling.
+# - Lock files live under: //_locks/.lock
+# - Create uses O_CREAT|O_EXCL for atomic acquisition across processes.
+# - Always release in finally; atexit guard removes any leftover locks.
+# ---------------------------------------------------------------
+_HELD_LOCKS = set()
+
+def _release_all_locks_at_exit():
+ for p in list(_HELD_LOCKS):
+ try:
+ if os.path.exists(p):
+ os.remove(p)
+ except Exception:
+ pass
+
+atexit.register(_release_all_locks_at_exit)
+
+def acquire_scene_lock(lock_root, scene_idx):
+ """Try to acquire an exclusive lock for a scene.
+
+ Returns lock_path (str) on success, or None if already locked or on failure.
+ """
+ try:
+ os.makedirs(lock_root, exist_ok=True)
+ lock_path = os.path.join(lock_root, f"{scene_idx:03d}.lock")
+ # atomic create; fails if exists
+ fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
+ # record some metadata for debugging
+ with os.fdopen(fd, "w") as f:
+ f.write(
+ f"pid={os.getpid()} host={socket.gethostname()} time={time.strftime('%Y-%m-%d %H:%M:%S')}\n"
+ )
+ _HELD_LOCKS.add(lock_path)
+ return lock_path
+ except FileExistsError:
+ return None
+ except Exception as e:
+ logger.warning(f"Failed to acquire lock for scene {scene_idx:03d}: {e}")
+ return None
+
+def release_scene_lock(lock_path):
+ """Release a previously acquired lock (idempotent)."""
+ if not lock_path:
+ return
+ try:
+ if os.path.exists(lock_path):
+ os.remove(lock_path)
+ except Exception as e:
+ logger.warning(f"Failed to release lock {lock_path}: {e}")
+ finally:
+ _HELD_LOCKS.discard(lock_path)
+
+def main(args):
+
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ opts = dict([(x.split("=", 1)[0], x.split("=", 1)[1]) for x in args.opts])
+ # all scenes
+ if opts['data.scene_idx'] == '-1':
+
+ # 轮询间隔（秒）
+ POLL_INTERVAL = 5
+
+ # 获取所有场景索引（假设目录名是数字，如 000、001、...）
+ assert 'data.data_root' in opts
+ data_root = opts['data.data_root']
+ all_scene_idxs = sorted(
+ int(d) for d in os.listdir(data_root)
+ if d.isdigit()
+ )
+ if args.end_idx != -1:
+ all_scene_idxs = all_scene_idxs[args.start_idx:args.end_idx+1]
+
+ # 待处理集合与失败记录
+ pending = set(all_scene_idxs)
+ all_fails = {}
+
+ # 锁目录：与输出同盘，避免写到只读数据盘
+ lock_root = os.path.join(args.output_root, args.project, "_locks")
+
+ while pending:
+ locked_scene = None
+ lock_path = None
+
+ # 在所有待处理场景中，优先挑选已就绪并且成功抢到锁的那一个
+ for scene_idx in sorted(pending):
+ scene_dir = os.path.join(data_root, f"{scene_idx:03d}")
+ if not is_data_ready(scene_dir):
+ continue
+
+ # 已处理完成的直接剔除（避免无谓加锁）
+ run_dir = os.path.join(args.output_root, args.project, f"{scene_idx:03d}")
+ if os.path.exists(os.path.join(run_dir, 'checkpoint_final.pth')):
+ print(f"Scene {scene_idx:03d} already processed, skipping.")
+ pending.remove(scene_idx)
+ continue
+
+ # 抢占式加锁：原子创建锁文件，失败则尝试下一个场景
+ lock_path_try = acquire_scene_lock(lock_root, scene_idx)
+ if lock_path_try is None:
+ continue
+
+ # 加锁后再二次确认是否已完成（与其他进程竞争到末尾的竞态）
+ if os.path.exists(os.path.join(run_dir, 'checkpoint_final.pth')):
+ # 已完成则释放锁并跳过
+ release_scene_lock(lock_path_try)
+ pending.remove(scene_idx)
+ continue
+
+ locked_scene = scene_idx
+ lock_path = lock_path_try
+ break
+
+ if locked_scene is None:
+ # 没有可加锁的且 ready 的场景，等待后重试
+ print(f"No ready & unlocked scenes. Waiting for {POLL_INTERVAL} seconds...")
+ time.sleep(POLL_INTERVAL)
+ continue
+
+ # 处理这个获得锁的场景
+ args.run_name = f"{locked_scene:03d}"
+ # 更新 CLI opts 中的 scene_idx（去重后添加）
+ for item in list(args.opts):
+ if item.startswith("data.scene_idx="):
+ args.opts.remove(item)
+ args.opts.append(f"data.scene_idx={locked_scene}")
+
+ # try:
+ cfg = setup(args)
+ train(cfg, device)
+ # except Exception as e:
+ # print(f"Fail to process scene {locked_scene:03d}, due to {e}")
+ # all_fails[locked_scene] = e
+ # finally:
+ # # 失败或成功都释放锁，避免死锁
+ # release_scene_lock(lock_path)
+
+ # 从待处理集合里移除
+ pending.remove(locked_scene)
+
+ print(all_fails)
+ else:
+ cfg = setup(args)
+ return train(cfg, device)
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser("Train Gaussian Splatting for a single scene")
+ parser.add_argument("--config_file", help="path to config file", type=str)
+ parser.add_argument("--output_root", default="./work_dirs/", help="path to save checkpoints and logs", type=str)
+
+ # eval
+ parser.add_argument("--resume_from", default=None, help="path to checkpoint to resume from", type=str)
+ parser.add_argument("--render_video_postfix", type=str, default=None, help="an optional postfix for video")
+
+ # wandb logging part
+ parser.add_argument("--enable_wandb", action="store_true", help="enable wandb logging")
+ parser.add_argument("--entity", default="ziyc", type=str, help="wandb entity name")
+ parser.add_argument("--project", default="drivestudio", type=str, help="wandb project name, also used to enhance log_dir")
+ parser.add_argument("--run_name", default="omnire", type=str, help="wandb run name, also used to enhance log_dir")
+
+ # viewer
+ parser.add_argument("--enable_viewer", action="store_true", help="enable viewer")
+ parser.add_argument("--viewer_port", type=int, default=8080, help="viewer port")
+
+ # start & end
+ parser.add_argument("--start_idx", type=int, default=0, help="start idx")
+ parser.add_argument("--end_idx", type=int, default=-1, help="end idx")
+
+ # novel view rendering
+ parser.add_argument("--novel_cam_ids", type=str, default=None,
+ help="Novel view cameras for evaluation: 'all' or comma-separated ids (e.g., '0,1,2'). Default: dataset default (front cam).")
+
+ # misc
+ parser.add_argument("opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER)
+
+ args = parser.parse_args()
+ final_step = main(args)
diff --git a/worldbench/third_party/drivestudio/tools/worldlens_reconstruction_backend.py b/worldbench/third_party/drivestudio/tools/worldlens_reconstruction_backend.py
new file mode 100644
index 0000000..d91df3d
--- /dev/null
+++ b/worldbench/third_party/drivestudio/tools/worldlens_reconstruction_backend.py
@@ -0,0 +1,379 @@
+#!/usr/bin/env python3
+"""DriveStudio backend entrypoint for WorldLens Reconstruction.
+
+This script is intentionally small: WorldLens calls it as an external process
+when Reconstruction uses ``backend=drivestudio``. Heavy DriveStudio imports and
+CUDA dependencies stay in the DriveStudio environment, while WorldLens reads the
+JSON contract written here.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+from typing import Dict, Iterable, List, Optional
+
+NOVEL_VIEWS = [
+ "front_center_interp",
+ "s_curve",
+ "lateral_offset",
+ "lateral_offset_left",
+]
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser("DriveStudio backend for WorldLens Reconstruction")
+ parser.add_argument("--method-name", required=True)
+ parser.add_argument("--clips", nargs="+", required=True)
+ parser.add_argument("--reconstruction-root", type=Path, required=True)
+ parser.add_argument("--work-root", type=Path, default=Path("work_dirs"))
+ parser.add_argument("--data-root", type=Path, default=Path("data/nuscenes_trainval/raw"))
+ parser.add_argument("--target-dir", type=Path, default=Path("data/nuscenes_trainval/processed"))
+ parser.add_argument("--config-file", default="configs/omnire.yaml")
+ parser.add_argument("--step", type=int, default=30000)
+ parser.add_argument("--fps", type=int, default=12)
+ parser.add_argument("--gt-method", default="gt")
+ parser.add_argument("--cameras", default="CAM_FRONT")
+ parser.add_argument("--skip-preprocess", action="store_true")
+ parser.add_argument("--skip-train", action="store_true")
+ parser.add_argument("--skip-eval", action="store_true")
+ parser.add_argument("--skip-depth-metrics", action="store_true", help=argparse.SUPPRESS)
+ parser.add_argument("--dry-run", action="store_true")
+ return parser.parse_args()
+
+
+def write_json(path: Path, payload: Dict[str, object]) -> None:
+ path.parent.mkdir(parents=True, exist_ok=True)
+ path.write_text(json.dumps(payload, indent=2))
+
+
+def method_root(reconstruction_root: Path, method_name: str) -> Path:
+ return reconstruction_root / method_name
+
+
+def write_manifest(args: argparse.Namespace) -> None:
+ write_json(
+ method_root(args.reconstruction_root, args.method_name) / "manifest.json",
+ {
+ "schema_version": 1,
+ "method_name": args.method_name,
+ "backend": "drivestudio",
+ "mode": "dry-run" if args.dry_run else "reconstruct",
+ "contract": "assets",
+ "clips": [str(clip) for clip in args.clips],
+ "train_step": args.step,
+ "config_file": args.config_file,
+ "geometric_mask_source": args.gt_method,
+ "novel_views": NOVEL_VIEWS,
+ },
+ )
+
+
+def write_dry_run_outputs(reconstruction_root: Path, method_name: str, clips: Iterable[str]) -> None:
+ root = method_root(reconstruction_root, method_name)
+ for clip in clips:
+ clip_id = str(clip)
+ for relative_dir in (
+ f"ckpt/{clip_id}",
+ f"train/{clip_id}/pred",
+ f"train/{clip_id}/gt",
+ f"depth/{clip_id}/pred",
+ f"depth/{clip_id}/gt",
+ f"depth/{clip_id}/masks",
+ f"novel/{clip_id}",
+ ):
+ (root / relative_dir).mkdir(parents=True, exist_ok=True)
+
+
+def run(command: List[str]) -> None:
+ subprocess.run(command, check=True)
+
+
+def run_preprocess(args: argparse.Namespace) -> None:
+ run(
+ [
+ sys.executable,
+ "datasets/preprocess.py",
+ "--data_root",
+ str(args.data_root),
+ "--target_dir",
+ str(args.target_dir),
+ "--dataset",
+ "nuscenes_gen",
+ "--split",
+ "advanced_12Hz_trainval",
+ "--process_keys",
+ "images",
+ "lidar",
+ "calib",
+ "dynamic_masks",
+ "objects",
+ "--data_source",
+ args.method_name,
+ "--scene_ids",
+ *[str(int(clip)) for clip in args.clips],
+ ]
+ )
+
+
+def run_train(args: argparse.Namespace) -> None:
+ # This mirrors train_nus_gen.sh while staying shell-independent.
+ run(
+ [
+ sys.executable,
+ "tools/train.py",
+ "--config_file",
+ args.config_file,
+ "--output_root",
+ str(args.work_root),
+ "--project",
+ f"drivestudio-nus-{args.method_name}",
+ "--run_name",
+ "-1",
+ "--start_idx",
+ str(min(int(clip) for clip in args.clips)),
+ "--end_idx",
+ str(max(int(clip) for clip in args.clips)),
+ "dataset=nuscenes_gen/6cams",
+ "data.scene_idx=-1",
+ "data.start_timestep=0",
+ "data.end_timestep=-1",
+ f"data.data_root={processed_data_root(args)}",
+ ]
+ )
+
+
+def processed_data_root_for_method(args: argparse.Namespace, method_name: str) -> Path:
+ processed_root = str(args.target_dir).replace("processed", f"processed_{method_name}", 1)
+ return Path(processed_root) / "advanced_12Hz_trainval"
+
+
+def processed_data_root(args: argparse.Namespace) -> Path:
+ return processed_data_root_for_method(args, args.method_name)
+
+
+def gt_processed_data_root(args: argparse.Namespace) -> Path:
+ return processed_data_root_for_method(args, args.gt_method)
+
+
+def validate_training_inputs(data_root: Path, clips: Iterable[str]) -> None:
+ missing = []
+ for clip in clips:
+ scene_dir = data_root / f"{int(clip):03d}"
+ images_dir = scene_dir / "images"
+ fine_masks_dir = scene_dir / "fine_dynamic_masks"
+ humanpose_path = scene_dir / "humanpose" / "smpl.pkl"
+ if not images_dir.is_dir():
+ missing.append(f"{scene_dir}: images")
+ continue
+ image_count = len(list(images_dir.iterdir()))
+ for mask_name in ("all", "human", "vehicle"):
+ mask_dir = fine_masks_dir / mask_name
+ if not mask_dir.is_dir():
+ missing.append(f"{scene_dir}: fine_dynamic_masks/{mask_name}")
+ elif len(list(mask_dir.iterdir())) != image_count:
+ missing.append(f"{scene_dir}: fine_dynamic_masks/{mask_name} count mismatch")
+ if not humanpose_path.is_file():
+ missing.append(f"{scene_dir}: humanpose/smpl.pkl")
+ if missing:
+ details = "\n".join(f"- {item}" for item in missing[:20])
+ raise RuntimeError(
+ "DriveStudio 4DGS training inputs are incomplete. "
+ "Run generated-data segmentation and human-pose preprocessing first.\n"
+ f"{details}"
+ )
+
+
+def run_render_and_export(args: argparse.Namespace) -> None:
+ run(
+ [
+ sys.executable,
+ "tools/render_all_for_model.py",
+ f"drivestudio-nus-{args.method_name}",
+ *[str(clip) for clip in args.clips],
+ "--tasks",
+ "train",
+ "depth",
+ "novel",
+ "--root",
+ str(args.work_root),
+ "--output-root",
+ str(args.reconstruction_root / "_drivestudio_render_all"),
+ "--fps",
+ str(args.fps),
+ ]
+ )
+
+
+def find_checkpoint(work_root: Path, method_name: str, clip_id: str) -> Optional[Path]:
+ run_dir = work_root / f"drivestudio-nus-{method_name}" / clip_id
+ final = run_dir / "checkpoint_final.pth"
+ if final.is_file():
+ return final
+ candidates = sorted(run_dir.glob("checkpoint_*.pth"))
+ return candidates[-1] if candidates else None
+
+
+def run_eval_for_existing_checkpoints(args: argparse.Namespace) -> None:
+ for clip in args.clips:
+ checkpoint = find_checkpoint(args.work_root, args.method_name, str(clip))
+ if checkpoint is None:
+ raise FileNotFoundError(
+ f"No checkpoint found for method={args.method_name}, clip={clip} "
+ f"under {args.work_root}"
+ )
+ run(
+ [
+ sys.executable,
+ "tools/eval.py",
+ "--resume_from",
+ str(checkpoint),
+ "--include_depth",
+ "--save_raw_depth",
+ "raw_depth",
+ ]
+ )
+
+
+def copy_file_if_exists(source: Path, target: Path) -> bool:
+ if not source.is_file():
+ return False
+ target.parent.mkdir(parents=True, exist_ok=True)
+ shutil.copy2(source, target)
+ return True
+
+
+def camera_id_from_name(camera_name: str) -> Optional[int]:
+ names = {
+ "CAM_FRONT": 0,
+ "CAM_FRONT_LEFT": 1,
+ "CAM_FRONT_RIGHT": 2,
+ "CAM_BACK_LEFT": 3,
+ "CAM_BACK_RIGHT": 4,
+ "CAM_BACK": 5,
+ }
+ return names.get(camera_name)
+
+
+def export_checkpoint_assets(args: argparse.Namespace, clip_id: str) -> None:
+ checkpoint = find_checkpoint(args.work_root, args.method_name, clip_id)
+ run_dir = args.work_root / f"drivestudio-nus-{args.method_name}" / clip_id
+ out_dir = method_root(args.reconstruction_root, args.method_name) / "ckpt" / clip_id
+ copy_file_if_exists(run_dir / "config.yaml", out_dir / "config.yaml")
+ if checkpoint is not None:
+ copy_file_if_exists(checkpoint, out_dir / checkpoint.name)
+ final_target = out_dir / "checkpoint_final.pth"
+ if checkpoint.name != final_target.name and not final_target.exists():
+ shutil.copy2(checkpoint, final_target)
+
+
+def export_train_assets(args: argparse.Namespace, clip_id: str) -> None:
+ render_root = (
+ args.reconstruction_root
+ / "_drivestudio_render_all"
+ / "train"
+ / clip_id
+ / f"drivestudio-nus-{args.method_name}"
+ )
+ out_root = method_root(args.reconstruction_root, args.method_name) / "train" / clip_id
+
+ camera_dirs = sorted(path for path in render_root.iterdir() if path.is_dir()) if render_root.is_dir() else []
+ for camera_dir in camera_dirs:
+ camera_name = camera_dir.name
+ for pred_path in sorted((camera_dir / "images").glob("*.jpg")):
+ frame_id = pred_path.stem
+ target_name = f"{frame_id}_{camera_name}.jpg"
+ copy_file_if_exists(pred_path, out_root / "pred" / target_name)
+ for gt_path in sorted((camera_dir / "gt_images").glob("*.jpg")):
+ frame_id = gt_path.stem
+ target_name = f"{frame_id}_{camera_name}.jpg"
+ copy_file_if_exists(gt_path, out_root / "gt" / target_name)
+
+
+def export_depth_assets(args: argparse.Namespace, clip_id: str) -> None:
+ out_root = method_root(args.reconstruction_root, args.method_name) / "depth" / clip_id
+ pred_dir = args.work_root / f"drivestudio-nus-{args.method_name}" / clip_id / "raw_depth" / f"full_set_{args.step}"
+ gt_dir = args.work_root / f"drivestudio-nus-{args.gt_method}" / clip_id / "raw_depth" / f"full_set_{args.step}"
+ mask_root = gt_processed_data_root(args) / f"{int(clip_id):03d}" / "sam_mask"
+
+ if not pred_dir.is_dir():
+ raise FileNotFoundError(f"Missing predicted raw depth directory: {pred_dir}")
+ pred_paths = sorted(pred_dir.glob("*.npy"))
+ if not pred_paths:
+ raise FileNotFoundError(f"No predicted raw depth files found under {pred_dir}")
+
+ for pred_path in pred_paths:
+ copy_file_if_exists(pred_path, out_root / "pred" / pred_path.name)
+ gt_path = gt_dir / pred_path.name
+ if not copy_file_if_exists(gt_path, out_root / "gt" / pred_path.name):
+ raise FileNotFoundError(f"Missing GT raw depth for geometric discrepancy: {gt_path}")
+ try:
+ idx_str, cam_name = pred_path.stem.split("_CAM_")
+ except ValueError:
+ raise ValueError(f"Unexpected raw depth filename, expected .npy: {pred_path}")
+ cam_name = f"CAM_{cam_name}" if not cam_name.startswith("CAM_") else cam_name
+ camera_id = camera_id_from_name(cam_name)
+ if camera_id is None:
+ raise ValueError(f"Unsupported camera name in raw depth filename: {pred_path}")
+ frame_idx = int(idx_str) // 6
+ mask_path = mask_root / f"{frame_idx:03d}_{camera_id}.png"
+ if not copy_file_if_exists(mask_path, out_root / "masks" / f"{pred_path.stem}.png"):
+ raise FileNotFoundError(
+ f"Missing Grounded-SAM2 evaluation mask for geometric discrepancy: {mask_path}. "
+ "Generate GT masks with worldbench/videogen/reconstruction/generate_eval_masks.py."
+ )
+
+
+def copy_novel_videos(args: argparse.Namespace, clip_id: str) -> None:
+ render_root = (
+ args.reconstruction_root
+ / "_drivestudio_render_all"
+ / "novel"
+ / clip_id
+ / f"drivestudio-nus-{args.method_name}"
+ )
+ out_dir = method_root(args.reconstruction_root, args.method_name) / "novel" / clip_id
+ out_dir.mkdir(parents=True, exist_ok=True)
+ for view_name in NOVEL_VIEWS:
+ source = render_root / view_name / "novel.mp4"
+ target = out_dir / f"{view_name}.mp4"
+ if source.is_file() and not target.exists():
+ shutil.copy2(source, target)
+
+
+def export_contract_assets(args: argparse.Namespace) -> None:
+ for clip in args.clips:
+ clip_id = str(clip)
+ export_checkpoint_assets(args, clip_id)
+ export_train_assets(args, clip_id)
+ export_depth_assets(args, clip_id)
+ copy_novel_videos(args, clip_id)
+
+
+def run_reconstruct(args: argparse.Namespace) -> None:
+ if not args.skip_preprocess:
+ run_preprocess(args)
+ if not args.skip_train:
+ validate_training_inputs(processed_data_root(args), args.clips)
+ run_train(args)
+ if not args.skip_eval:
+ run_eval_for_existing_checkpoints(args)
+ run_render_and_export(args)
+ export_contract_assets(args)
+
+
+def main() -> None:
+ args = parse_args()
+ write_manifest(args)
+ if args.dry_run:
+ write_dry_run_outputs(args.reconstruction_root, args.method_name, args.clips)
+ return
+ run_reconstruct(args)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/third_party/drivestudio/utils/__init__.py b/worldbench/third_party/drivestudio/utils/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/worldbench/third_party/drivestudio/utils/backup.py b/worldbench/third_party/drivestudio/utils/backup.py
new file mode 100644
index 0000000..05097e8
--- /dev/null
+++ b/worldbench/third_party/drivestudio/utils/backup.py
@@ -0,0 +1,55 @@
+# Modified from: https://github.com/pjlab-ADG/nr3d_lib/
+import os
+import shutil
+import logging
+
+logger = logging.getLogger()
+
+def backup_folder(
+ backup_dir: str,
+ source_dir: str="./",
+ filetypes_to_copy=[".py", ".h", ".cpp", ".cuh", ".cu", ".sh"]
+):
+ filetypes_to_copy = tuple(filetypes_to_copy)
+ os.makedirs(backup_dir, exist_ok=True)
+ for file in os.listdir(source_dir):
+ if not file.endswith(filetypes_to_copy):
+ continue
+ source_file_path = os.path.join(source_dir, file)
+ target_file_path = os.path.join(backup_dir, file)
+ shutil.copy(source_file_path, target_file_path)
+
+def backup_folder_recursive(
+ backup_dir: str,
+ source_dir: str="./",
+ filetypes_to_copy=[".py", ".h", ".cpp", ".cuh", ".cu", ".sh"]
+):
+ filetypes_to_copy = tuple(filetypes_to_copy)
+ for root, _, files in os.walk(source_dir):
+ for file in files:
+ if not file.endswith(filetypes_to_copy):
+ continue
+ source_file_path = os.path.join(root, file)
+ # Keeps original directory structure
+ target_file_path = os.path.join(backup_dir, os.path.relpath(source_file_path, source_dir))
+ target_dir_path = os.path.dirname(target_file_path)
+ os.makedirs(target_dir_path, exist_ok=True)
+ shutil.copy(source_file_path, target_file_path)
+
+def backup_project(
+ backup_dir: str,
+ source_dir: str="./",
+ subdirs_to_copy=["app", "code_multi", "code_single", "dataio", "nr3d_lib"],
+ filetypes_to_copy=[".py", ".h", ".cpp", ".cuh", ".cu", ".sh"],
+):
+ filetypes_to_copy = tuple(filetypes_to_copy)
+ # Automatic backup codes
+ logger.info(f"=> Backing up from {source_dir} to {backup_dir}...")
+ # Backup project root dir, depth = 1
+ backup_folder(backup_dir, source_dir, filetypes_to_copy)
+ # Backup cared subdirs, depth = inf
+ for subdir in subdirs_to_copy:
+ sub_source_dir = os.path.join(source_dir, subdir)
+ sub_backup_dir = os.path.join(backup_dir, subdir)
+ backup_folder_recursive(sub_backup_dir, sub_source_dir, filetypes_to_copy)
+ logger.info("done.")
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/utils/camera.py b/worldbench/third_party/drivestudio/utils/camera.py
new file mode 100644
index 0000000..2010976
--- /dev/null
+++ b/worldbench/third_party/drivestudio/utils/camera.py
@@ -0,0 +1,190 @@
+# Camera pose manipulation and trajectory generation.
+import os
+import torch
+import numpy as np
+from typing import Dict
+
+from scipy.spatial.transform import Slerp
+from scipy.spatial.transform import Rotation as R
+
+def interpolate_poses(key_poses: torch.Tensor, target_frames: int) -> torch.Tensor:
+ """
+ Interpolate between key poses to generate a smooth trajectory.
+
+ Args:
+ key_poses (torch.Tensor): Tensor of shape (N, 4, 4) containing key camera poses.
+ target_frames (int): Number of frames to interpolate.
+
+ Returns:
+ torch.Tensor: Interpolated poses of shape (target_frames, 4, 4).
+ """
+ device = key_poses.device
+ key_poses = key_poses.cpu().numpy()
+
+ # Separate translation and rotation
+ translations = key_poses[:, :3, 3]
+ rotations = key_poses[:, :3, :3]
+
+ # Create time array
+ times = np.linspace(0, 1, len(key_poses))
+ target_times = np.linspace(0, 1, target_frames)
+
+ # Interpolate translations
+ interp_translations = np.stack([
+ np.interp(target_times, times, translations[:, i])
+ for i in range(3)
+ ], axis=-1)
+
+ # Interpolate rotations using Slerp
+ key_rots = R.from_matrix(rotations)
+ slerp = Slerp(times, key_rots)
+ interp_rotations = slerp(target_times).as_matrix()
+
+ # Combine interpolated translations and rotations
+ interp_poses = np.eye(4)[None].repeat(target_frames, axis=0)
+ interp_poses[:, :3, :3] = interp_rotations
+ interp_poses[:, :3, 3] = interp_translations
+
+ return torch.tensor(interp_poses, dtype=torch.float32, device=device)
+
+def look_at_rotation(direction: torch.Tensor, up: torch.Tensor = torch.tensor([0., 0., 1.])) -> torch.Tensor:
+ """Calculate rotation matrix to look at a specific direction."""
+ front = torch.nn.functional.normalize(direction, dim=-1)
+ right = torch.nn.functional.normalize(torch.cross(front, up), dim=-1)
+ up = torch.cross(right, front)
+ rotation_matrix = torch.stack([right, up, -front], dim=-1)
+ return rotation_matrix
+
+def get_interp_novel_trajectories(
+ dataset_type: str,
+ scene_idx: str,
+ per_cam_poses: Dict[int, torch.Tensor],
+ traj_type: str = "front_center_interp",
+ target_frames: int = 100
+) -> torch.Tensor:
+ original_frames = per_cam_poses[list(per_cam_poses.keys())[0]].shape[0]
+
+ trajectory_generators = {
+ "front_center_interp": front_center_interp,
+ "s_curve": s_curve,
+ "three_key_poses": three_key_poses_trajectory,
+ "lateral_offset": lateral_offset_trajectory,
+ }
+
+ if traj_type not in trajectory_generators:
+ raise ValueError(f"Unknown trajectory type: {traj_type}")
+
+ return trajectory_generators[traj_type](dataset_type, per_cam_poses, original_frames, target_frames)
+
+def front_center_interp(
+ dataset_type: str, per_cam_poses: Dict[int, torch.Tensor], original_frames: int, target_frames: int, num_loops: int = 1
+) -> torch.Tensor:
+ """Interpolate key frames from the front center camera."""
+ assert 0 in per_cam_poses.keys(), "Front center camera (ID 0) is required for front_center_interp"
+ key_poses = per_cam_poses[0][::original_frames//4] # Select every 4th frame as key frame
+ return interpolate_poses(key_poses, target_frames)
+
+def s_curve(
+ dataset_type: str, per_cam_poses: Dict[int, torch.Tensor], original_frames: int, target_frames: int
+) -> torch.Tensor:
+ """Create an S-shaped trajectory using the front three cameras."""
+ assert all(cam in per_cam_poses.keys() for cam in [0, 1, 2]), "Front three cameras (IDs 0, 1, 2) are required for s_curve"
+ key_poses = torch.cat([
+ per_cam_poses[0][0:1],
+ per_cam_poses[1][original_frames//4:original_frames//4+1],
+ per_cam_poses[0][original_frames//2:original_frames//2+1],
+ per_cam_poses[2][3*original_frames//4:3*original_frames//4+1],
+ per_cam_poses[0][-1:]
+ ], dim=0)
+ return interpolate_poses(key_poses, target_frames)
+
+def three_key_poses_trajectory(
+ dataset_type: str,
+ per_cam_poses: Dict[int, torch.Tensor],
+ original_frames: int,
+ target_frames: int
+) -> torch.Tensor:
+ """
+ Create a trajectory using three key poses:
+ 1. First frame of front center camera
+ 2. Middle frame with interpolated rotation and position from camera 1 or 2
+ 3. Last frame of front center camera
+
+ The rotation of the middle pose is calculated using Slerp between
+ the start frame and the middle frame of camera 1 or 2.
+
+ Args:
+ dataset_type (str): Type of the dataset (e.g., "waymo", "pandaset", etc.).
+ per_cam_poses (Dict[int, torch.Tensor]): Dictionary of camera poses.
+ original_frames (int): Number of original frames.
+ target_frames (int): Number of frames in the output trajectory.
+
+ Returns:
+ torch.Tensor: Trajectory of shape (target_frames, 4, 4).
+ """
+ assert 0 in per_cam_poses.keys(), "Front center camera (ID 0) is required"
+ assert 1 in per_cam_poses.keys() or 2 in per_cam_poses.keys(), "Either camera 1 or camera 2 is required"
+
+ # First key pose: First frame of front center camera
+ start_pose = per_cam_poses[0][0]
+ key_poses = [start_pose]
+
+ # Select camera for middle frame
+ middle_frame = int(original_frames // 2)
+ chosen_cam = np.random.choice([1, 2])
+
+ middle_pose = per_cam_poses[chosen_cam][middle_frame]
+
+ # Calculate interpolated rotation for middle pose
+ start_rotation = R.from_matrix(start_pose[:3, :3].cpu().numpy())
+ middle_rotation = R.from_matrix(middle_pose[:3, :3].cpu().numpy())
+ slerp = Slerp([0, 1], R.from_quat([start_rotation.as_quat(), middle_rotation.as_quat()]))
+ interpolated_rotation = slerp(0.5).as_matrix()
+
+ # Create middle key pose with interpolated rotation and original translation
+ middle_key_pose = torch.eye(4, device=start_pose.device)
+ middle_key_pose[:3, :3] = torch.tensor(interpolated_rotation, device=start_pose.device)
+ middle_key_pose[:3, 3] = middle_pose[:3, 3] # Keep the original translation
+ key_poses.append(middle_key_pose)
+
+ # Third key pose: Last frame of front center camera
+ key_poses.append(per_cam_poses[0][-1])
+
+ # Stack the key poses and interpolate
+ key_poses = torch.stack(key_poses)
+ return interpolate_poses(key_poses, target_frames)
+
+def lateral_offset_trajectory(
+ dataset_type: str,
+ per_cam_poses: Dict[int, torch.Tensor],
+ original_frames: int,
+ target_frames: int,
+ offset_distance: float = 1.0,
+) -> torch.Tensor:
+ """
+ Translate the default reference camera trajectory sideways by a fixed amount.
+
+ Args:
+ dataset_type: Dataset name to select the default camera.
+ per_cam_poses: Original camera poses keyed by camera id.
+ original_frames: Number of frames in the original recording.
+ target_frames: Desired length of the interpolated trajectory.
+ offset_distance: Lateral offset in meters (positive = right, negative = left).
+
+ Returns:
+ torch.Tensor: Interpolated poses with lateral offset applied.
+ """
+ if len(per_cam_poses) == 0:
+ raise ValueError("per_cam_poses is empty, cannot build lateral offset trajectory")
+
+ # Select the same default camera used for front-center renderings.
+ if dataset_type == "argoverse":
+ default_cam = 1 if 1 in per_cam_poses else 0
+ else:
+ default_cam = 0 if 0 in per_cam_poses else list(per_cam_poses.keys())[0]
+
+ base_poses = per_cam_poses[default_cam].clone()
+ lateral_axis = base_poses[:, :3, 0]
+ base_poses[:, :3, 3] = base_poses[:, :3, 3] + lateral_axis * offset_distance
+
+ return interpolate_poses(base_poses, target_frames)
diff --git a/worldbench/third_party/drivestudio/utils/gather_results.py b/worldbench/third_party/drivestudio/utils/gather_results.py
new file mode 100644
index 0000000..b6bc36a
--- /dev/null
+++ b/worldbench/third_party/drivestudio/utils/gather_results.py
@@ -0,0 +1,62 @@
+import os
+import json
+
+# novel view synthesis results
+root = "logs/omnire_waymo/nvs"
+split = "test"
+# scene reconstruction results
+# root = "logs/omnire_waymo/recon"
+# split = "full"
+
+scene_ids = os.listdir(root)
+scene_ids.sort()
+metrics = {}
+for scene_id in scene_ids:
+ try:
+ metric = os.listdir(os.path.join(root, str(scene_id), "metrics"))
+ metric = [m for m in metric if m.startswith(f"images_{split}")][0]
+ with open(os.path.join(root, str(scene_id), "metrics", metric), "r") as f:
+ data = json.load(f)
+ except:
+ print("error", scene_id)
+ continue
+ try:
+ metrics[scene_id] = {
+ "psnr": data[f"image_metrics/{split}/psnr"],
+ "ssim": data[f"image_metrics/{split}/ssim"],
+ "human_psnr": data[f"image_metrics/{split}/human_psnr"],
+ "human_ssim": data[f"image_metrics/{split}/human_ssim"],
+ "vehicle_psnr": data[f"image_metrics/{split}/vehicle_psnr"],
+ "vehicle_ssim": data[f"image_metrics/{split}/vehicle_ssim"],
+ }
+ except:
+ print(data)
+ breakpoint()
+for scene_id in scene_ids:
+ if scene_id in metrics:
+ print(
+ scene_id,
+ metrics[scene_id]["psnr"],
+ metrics[scene_id]["ssim"],
+ metrics[scene_id]["human_psnr"],
+ metrics[scene_id]["human_ssim"],
+ metrics[scene_id]["vehicle_psnr"],
+ metrics[scene_id]["vehicle_ssim"],
+ )
+ else:
+ print("error", scene_id)
+ continue
+# print avg
+psnr = [metrics[scene_id]["psnr"] for scene_id in metrics.keys()]
+ssim = [metrics[scene_id]["ssim"] for scene_id in metrics.keys()]
+human_psnr = [metrics[scene_id]["human_psnr"] for scene_id in metrics.keys()]
+human_ssim = [metrics[scene_id]["human_ssim"] for scene_id in metrics.keys()]
+vehicle_psnr = [metrics[scene_id]["vehicle_psnr"] for scene_id in metrics.keys()]
+vehicle_ssim = [metrics[scene_id]["vehicle_ssim"] for scene_id in metrics.keys()]
+print("avg Full Image PSNR: ", sum(psnr) / len(psnr))
+print("avg Full Image SSIM: ", sum(ssim) / len(ssim))
+print("avg Human-Only PSNR: ", sum(human_psnr) / len(human_psnr))
+print("avg Human-Only SSIM: ", sum(human_ssim) / len(human_ssim))
+print("avg Vehicle-Only PSNR: ", sum(vehicle_psnr) / len(vehicle_psnr))
+print("avg Vehicle-Only SSIM: ", sum(vehicle_ssim) / len(vehicle_ssim))
+print("len", len(psnr))
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/utils/geometry.py b/worldbench/third_party/drivestudio/utils/geometry.py
new file mode 100644
index 0000000..e0d72b3
--- /dev/null
+++ b/worldbench/third_party/drivestudio/utils/geometry.py
@@ -0,0 +1,104 @@
+# Utility functions for geometric transformations and projections.
+import numpy as np
+import torch
+from torch import Tensor
+import torch.nn.functional as F
+
+def transform_points(points, transform_matrix):
+ """
+ Apply a 4x4 transformation matrix to 3D points.
+
+ Args:
+ points: (N, 3) tensor of 3D points
+ transform_matrix: (4, 4) transformation matrix
+
+ Returns:
+ (N, 3) tensor of transformed 3D points
+ """
+ ones = torch.ones((points.shape[0], 1), dtype=points.dtype, device=points.device)
+ homo_points = torch.cat([points, ones], dim=1) # N x 4
+ transformed_points = torch.matmul(homo_points, transform_matrix.T)
+ return transformed_points[:, :3]
+
+def get_corners(l: float, w: float, h: float):
+ """
+ Get 8 corners of a 3D bounding box centered at origin.
+
+ Args:
+ l, w, h: length, width, height of the box
+
+ Returns:
+ (3, 8) array of corner coordinates
+ """
+ return np.array([
+ [-l/2, -l/2, l/2, l/2, -l/2, -l/2, l/2, l/2],
+ [w/2, -w/2, -w/2, w/2, w/2, -w/2, -w/2, w/2],
+ [h/2, h/2, h/2, h/2, -h/2, -h/2, -h/2, -h/2],
+ ])
+
+def project_camera_points_to_image(points_cam, cam_intrinsics):
+ """
+ Project 3D points from camera space to 2D image space.
+
+ Args:
+ points_cam (np.ndarray): Shape (N, 3), points in camera space.
+ cam_intrinsics (np.ndarray): Shape (3, 3), intrinsic matrix of the camera.
+
+ Returns:
+ tuple: (projected_points, depths)
+ - projected_points (np.ndarray): Shape (N, 2), projected 2D points in image space.
+ - depths (np.ndarray): Shape (N,), depth values of the projected points.
+ """
+ points_img = cam_intrinsics @ points_cam.T
+ depths = points_img[2, :]
+ projected_points = (points_img[:2, :] / (depths + 1e-6)).T
+
+ return projected_points, depths
+
+def cube_root(x):
+ return torch.sign(x) * torch.abs(x) ** (1. / 3)
+
+def spherical_to_cartesian(r, theta, phi):
+ x = r * torch.sin(theta) * torch.cos(phi)
+ y = r * torch.sin(theta) * torch.sin(phi)
+ z = r * torch.cos(theta)
+ return torch.stack([x, y, z], dim=1)
+
+def uniform_sample_sphere(num_samples, device, inverse=False):
+ """
+ refer to https://stackoverflow.com/questions/5408276/sampling-uniformly-distributed-random-points-inside-a-spherical-volume
+ sample points uniformly inside a sphere
+ """
+ if not inverse:
+ dist = torch.rand((num_samples,)).to(device)
+ dist = cube_root(dist)
+ else:
+ dist = torch.rand((num_samples,)).to(device)
+ dist = 1 / dist.clamp_min(0.02)
+ thetas = torch.arccos(2 * torch.rand((num_samples,)) - 1).to(device)
+ phis = 2 * torch.pi * torch.rand((num_samples,)).to(device)
+ pts = spherical_to_cartesian(dist, thetas, phis)
+ return pts
+
+def rotation_6d_to_matrix(d6: Tensor) -> Tensor:
+ """
+ Converts 6D rotation representation by Zhou et al. [1] to rotation matrix
+ using Gram--Schmidt orthogonalization per Section B of [1]. Adapted from pytorch3d.
+ Args:
+ d6: 6D rotation representation, of size (*, 6)
+
+ Returns:
+ batch of rotation matrices of size (*, 3, 3)
+
+ [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
+ On the Continuity of Rotation Representations in Neural Networks.
+ IEEE Conference on Computer Vision and Pattern Recognition, 2019.
+ Retrieved from http://arxiv.org/abs/1812.07035
+ """
+
+ a1, a2 = d6[..., :3], d6[..., 3:]
+ b1 = F.normalize(a1, dim=-1)
+ b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1
+ b2 = F.normalize(b2, dim=-1)
+ b3 = torch.cross(b1, b2, dim=-1)
+ return torch.stack((b1, b2, b3), dim=-2)
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/utils/logging.py b/worldbench/third_party/drivestudio/utils/logging.py
new file mode 100644
index 0000000..9a59fea
--- /dev/null
+++ b/worldbench/third_party/drivestudio/utils/logging.py
@@ -0,0 +1,300 @@
+# Modified from: https://github.com/facebookresearch/dinov2/blob/main/dinov2/logging/helpers.py
+import os
+import sys
+import time
+import json
+import logging
+import datetime
+import functools
+
+from typing import Optional
+from collections import defaultdict, deque
+
+import torch
+
+from utils.misc import get_global_rank, is_enabled, is_main_process
+
+logger = logging.getLogger()
+
+class MetricLogger(object):
+ def __init__(self, delimiter="\t", output_file=None):
+ self.meters = defaultdict(SmoothedValue)
+ self.delimiter = delimiter
+ self.output_file = output_file
+
+ def update(self, **kwargs):
+ for k, v in kwargs.items():
+ if isinstance(v, torch.Tensor):
+ v = v.item()
+ assert isinstance(v, (float, int))
+ self.meters[k].update(v)
+
+ def __getattr__(self, attr):
+ if attr in self.meters:
+ return self.meters[attr]
+ if attr in self.__dict__:
+ return self.__dict__[attr]
+ raise AttributeError(
+ f"'{type(self).__name__}' object has no attribute '{attr}'"
+ )
+
+ def __str__(self):
+ loss_str = []
+ for name, meter in self.meters.items():
+ loss_str.append(f"{name}: {str(meter)}")
+ return self.delimiter.join(loss_str)
+
+ def synchronize_between_processes(self):
+ for meter in self.meters.values():
+ meter.synchronize_between_processes()
+
+ def add_meter(self, name, meter):
+ self.meters[name] = meter
+
+ def dump_in_output_file(self, iteration, iter_time, data_time):
+ if self.output_file is None:
+ return
+ dict_to_dump = dict(
+ iteration=iteration,
+ iter_time=iter_time,
+ data_time=data_time,
+ )
+ dict_to_dump.update({k: v.median for k, v in self.meters.items()})
+ with open(self.output_file, "a") as f:
+ f.write(json.dumps(dict_to_dump) + "\n")
+ pass
+
+ def log_every(
+ self, iterable, print_freq, header=None, n_iterations=None, start_iteration=0
+ ):
+ i = start_iteration
+ if not header:
+ header = ""
+ start_time = time.time()
+ end = time.time()
+ iter_time = SmoothedValue(fmt="{avg:.6f}")
+ data_time = SmoothedValue(fmt="{avg:.6f}")
+
+ if n_iterations is None:
+ n_iterations = len(iterable)
+
+ space_fmt = ":" + str(len(str(n_iterations))) + "d"
+
+ log_list = [
+ header,
+ "[{0" + space_fmt + "}/{1}]",
+ "eta: {eta}",
+ "elapsed: {elapsed_time_str}",
+ "{meters}",
+ "time: {time}",
+ "data: {data}",
+ ]
+ if torch.cuda.is_available():
+ log_list += ["max mem: {memory:.0f}"]
+
+ log_msg = self.delimiter.join(log_list)
+ MB = 1024.0 * 1024.0
+ for obj in iterable:
+ data_time.update(time.time() - end)
+ yield obj
+ iter_time.update(time.time() - end)
+ if i % print_freq == 0 or i == n_iterations - 1:
+ self.dump_in_output_file(
+ iteration=i, iter_time=iter_time.avg, data_time=data_time.avg
+ )
+ eta_seconds = iter_time.global_avg * (n_iterations - i)
+ eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+ elapsed_time = time.time() - start_time
+ elapsed_time_str = str(datetime.timedelta(seconds=int(elapsed_time)))
+
+ if torch.cuda.is_available():
+ logger.info(
+ log_msg.format(
+ i,
+ n_iterations,
+ eta=eta_string,
+ elapsed_time_str=elapsed_time_str,
+ meters=str(self),
+ time=str(iter_time),
+ data=str(data_time),
+ memory=torch.cuda.max_memory_allocated() / MB,
+ )
+ )
+ else:
+ logger.info(
+ log_msg.format(
+ i,
+ n_iterations,
+ eta=eta_string,
+ meters=str(self),
+ time=str(iter_time),
+ data=str(data_time),
+ )
+ )
+ i += 1
+ end = time.time()
+ if i >= n_iterations:
+ break
+ total_time = time.time() - start_time
+ total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+ logger.info(
+ f"{header} Total time: {total_time_str} ({total_time / n_iterations:.6f} s / it)"
+ )
+
+
+class SmoothedValue:
+ """Track a series of values and provide access to smoothed values over a
+ window or the global series average.
+ """
+
+ def __init__(self, window_size=20, fmt=None):
+ if fmt is None:
+ fmt = "{median:.6f} ({global_avg:.6f})"
+ self.deque = deque(maxlen=window_size)
+ self.total = 0.0
+ self.count = 0
+ self.fmt = fmt
+
+ def update(self, value, num=1):
+ self.deque.append(value)
+ self.count += num
+ self.total += value * num
+
+ def synchronize_between_processes(self):
+ """
+ Distributed synchronization of the metric
+ Warning: does not synchronize the deque!
+ """
+ if not is_enabled():
+ return
+ t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda")
+ torch.distributed.barrier()
+ torch.distributed.all_reduce(t)
+ t = t.tolist()
+ self.count = int(t[0])
+ self.total = t[1]
+
+ @property
+ def median(self):
+ d = torch.tensor(list(self.deque))
+ return d.median().item()
+
+ @property
+ def avg(self):
+ d = torch.tensor(list(self.deque), dtype=torch.float32)
+ return d.mean().item()
+
+ @property
+ def global_avg(self):
+ return self.total / self.count
+
+ @property
+ def max(self):
+ return max(self.deque)
+
+ @property
+ def value(self):
+ return self.deque[-1]
+
+ def __str__(self):
+ return self.fmt.format(
+ median=self.median,
+ avg=self.avg,
+ global_avg=self.global_avg,
+ max=self.max,
+ value=self.value,
+ )
+
+
+# So that calling _configure_logger multiple times won't add many handlers
+@functools.lru_cache()
+def _configure_logger(
+ name: Optional[str] = None,
+ *,
+ level: int = logging.DEBUG,
+ output: Optional[str] = None,
+ time_string: Optional[str] = None,
+):
+ """
+ Configure a logger.
+
+ Adapted from Detectron2.
+
+ Args:
+ name: The name of the logger to configure.
+ level: The logging level to use.
+ output: A file name or a directory to save log. If None, will not save log file.
+ If ends with ".txt" or ".log", assumed to be a file name.
+ Otherwise, logs will be saved to `output/log.txt`.
+
+ Returns:
+ The configured logger.
+ """
+
+ logger = logging.getLogger(name)
+ logger.setLevel(level)
+ logger.propagate = False
+
+ # Loosely match Google glog format:
+ # [IWEF]yyyymmdd hh:mm:ss.uuuuuu threadid file:line] msg
+ # but use a shorter timestamp and include the logger name:
+ # [IWEF]yyyymmdd hh:mm:ss logger threadid file:line] msg
+ fmt_prefix = "%(levelname).1s%(asctime)s %(name)s %(filename)s:%(lineno)s] "
+ fmt_message = "%(message)s"
+ fmt = fmt_prefix + fmt_message
+ datefmt = "%Y%m%d %H:%M:%S"
+ formatter = logging.Formatter(fmt=fmt, datefmt=datefmt)
+
+ # stdout logging for main worker only
+ if is_main_process():
+ handler = logging.StreamHandler(stream=sys.stdout)
+ handler.setLevel(logging.DEBUG)
+ handler.setFormatter(formatter)
+ logger.addHandler(handler)
+
+ # file logging for all workers
+ if output:
+ if os.path.splitext(output)[-1] in (".txt", ".log"):
+ filename = output
+ else:
+ if time_string is None:
+ filename = os.path.join(output, "logs", "log.txt")
+ else:
+ filename = os.path.join(output, "logs", f"log_{time_string}.txt")
+
+ if not is_main_process():
+ global_rank = get_global_rank()
+ filename = filename + ".rank{}".format(global_rank)
+
+ os.makedirs(os.path.dirname(filename), exist_ok=True)
+
+ handler = logging.StreamHandler(open(filename, "a"))
+ handler.setLevel(logging.DEBUG)
+ handler.setFormatter(formatter)
+ logger.addHandler(handler)
+
+ return logger
+
+
+def setup_logging(
+ output: Optional[str] = None,
+ *,
+ name: Optional[str] = None,
+ level: int = logging.DEBUG,
+ capture_warnings: bool = True,
+ time_string: Optional[str] = None,
+) -> None:
+ """
+ Setup logging.
+
+ Args:
+ output: A file name or a directory to save log files. If None, log
+ files will not be saved. If output ends with ".txt" or ".log", it
+ is assumed to be a file name.
+ Otherwise, logs will be saved to `output/log.txt`.
+ name: The name of the logger to configure, by default the root logger.
+ level: The logging level to use.
+ capture_warnings: Whether warnings should be captured as logs.
+ """
+ logging.captureWarnings(capture_warnings)
+ _configure_logger(name, level=level, output=output, time_string=time_string)
diff --git a/worldbench/third_party/drivestudio/utils/misc.py b/worldbench/third_party/drivestudio/utils/misc.py
new file mode 100644
index 0000000..d0c5372
--- /dev/null
+++ b/worldbench/third_party/drivestudio/utils/misc.py
@@ -0,0 +1,124 @@
+# Miscellaneous utility functions for exporting point clouds.
+import importlib
+import logging
+import os
+
+import numpy as np
+import open3d as o3d
+import torch
+import torch.distributed as dist
+
+logger = logging.getLogger()
+
+def import_str(string: str):
+ """ Import a python module given string paths
+
+ Args:
+ string (str): The given paths
+
+ Returns:
+ Any: Imported python module / object
+ """
+ # From https://github.com/CompVis/taming-transformers
+ module, cls = string.rsplit(".", 1)
+ return getattr(importlib.import_module(module, package=None), cls)
+
+def export_points_to_ply(
+ positions: torch.tensor,
+ colors: torch.tensor,
+ save_path: str,
+ normalize: bool = False,
+ ):
+ # normalize points
+ if normalize:
+ aabb_min = positions.min(0)[0]
+ aabb_max = positions.max(0)[0]
+ positions = (positions - aabb_min) / (aabb_max - aabb_min)
+ if isinstance(colors, torch.Tensor):
+ positions = positions.cpu().numpy()
+ if isinstance(colors, torch.Tensor):
+ colors = colors.cpu().numpy()
+
+ # clamp colors
+ colors = np.clip(colors, a_min=0., a_max=1.)
+
+ pcd = o3d.geometry.PointCloud()
+ pcd.points = o3d.utility.Vector3dVector(positions)
+ pcd.colors = o3d.utility.Vector3dVector(colors)
+ o3d.io.write_point_cloud(save_path, pcd)
+
+def export_gaussians_to_ply(model, path, name='point_cloud.ply', aabb=None):
+ model.eval()
+ filename = os.path.join(path, name)
+ map_to_tensors = {}
+
+ with torch.no_grad():
+ positions = model.means
+ if aabb is not None:
+ aabb = aabb.to(positions.device)
+ aabb_min, aabb_max = aabb[:3], aabb[3:]
+ aabb_center = (aabb_min + aabb_max) / 2
+ aabb_sacle_max = (aabb_max - aabb_min).max() / 2 * 1.1
+ vis_mask = torch.logical_and(positions >= aabb_min, positions < aabb_max).all(-1)
+ else:
+ aabb_center = positions.mean(0)
+ aabb_sacle_max = (positions - aabb_center).abs().max() * 1.1
+ vis_mask = torch.ones_like(positions[:, 0], dtype=torch.bool)
+
+ positions = ((positions[vis_mask] - aabb_center) / aabb_sacle_max).cpu().numpy()
+ map_to_tensors["positions"] = o3d.core.Tensor(positions, o3d.core.float32)
+ map_to_tensors["normals"] = o3d.core.Tensor(np.zeros_like(positions), o3d.core.float32)
+
+ colors = model.colors[vis_mask].data.cpu().numpy()
+ map_to_tensors["colors"] = (colors * 255).astype(np.uint8)
+ for i in range(colors.shape[1]):
+ map_to_tensors[f"f_dc_{i}"] = colors[:, i : i + 1]
+
+ shs = model.shs_rest[vis_mask].data.cpu().numpy()
+ if model.config.sh_degree > 0:
+ shs = shs.reshape((colors.shape[0], -1, 1))
+ for i in range(shs.shape[-1]):
+ map_to_tensors[f"f_rest_{i}"] = shs[:, i]
+
+ map_to_tensors["opacity"] = model.opacities[vis_mask].data.cpu().numpy()
+
+ scales = model.scales[vis_mask].data.cpu().unsqueeze(-1).numpy()
+ for i in range(3):
+ map_to_tensors[f"scale_{i}"] = scales[:, i]
+
+ quats = model.quats[vis_mask].data.cpu().unsqueeze(-1).numpy()
+
+ for i in range(4):
+ map_to_tensors[f"rot_{i}"] = quats[:, i]
+
+ pcd = o3d.t.geometry.PointCloud(map_to_tensors)
+ o3d.t.io.write_point_cloud(str(filename), pcd)
+
+ logger.info(f"Exported point cloud to {filename}, containing {vis_mask.sum().item()} points.")
+
+def is_enabled() -> bool:
+ """
+ Returns:
+ True if distributed training is enabled
+ """
+ return dist.is_available() and dist.is_initialized()
+
+
+def get_global_rank() -> int:
+ """
+ Returns:
+ The rank of the current process within the global process group.
+ """
+ return dist.get_rank() if is_enabled() else 0
+
+
+def get_world_size():
+ return dist.get_world_size() if is_enabled() else 1
+
+
+def is_main_process() -> bool:
+ """
+ Returns:
+ True if the current process is the main one.
+ """
+ return get_global_rank() == 0
\ No newline at end of file
diff --git a/worldbench/third_party/drivestudio/utils/visualization.py b/worldbench/third_party/drivestudio/utils/visualization.py
new file mode 100644
index 0000000..83899e0
--- /dev/null
+++ b/worldbench/third_party/drivestudio/utils/visualization.py
@@ -0,0 +1,497 @@
+"""
+@brief Utilities for visualizing multi-camera data and depth
+
+This module provides functions for:
+- Combining multiple camera images into tiled layouts for various datasets
+- Visualizing depth information on images
+- Color mapping and depth visualization
+- Utility functions for image processing and color manipulation
+"""
+
+import hashlib
+from typing import List, Optional, Tuple
+
+import cv2
+import matplotlib.cm as cm
+import numpy as np
+import torch
+
+def to8b(x):
+ if isinstance(x, torch.Tensor):
+ x = x.detach().cpu().numpy()
+ return (255 * np.clip(x, 0, 1)).astype(np.uint8)
+
+def get_layout(dataset_type: str):
+ if dataset_type == "waymo":
+ layout = layout_waymo
+ elif dataset_type == "pandaset":
+ layout = layout_pandaset
+ elif dataset_type == "argoverse":
+ layout = layout_argoverse
+ elif dataset_type == "nuscenes":
+ layout = layout_nuscenes
+ elif dataset_type == "kitti":
+ layout = layout_kitti
+ elif dataset_type == "nuplan":
+ layout = layout_nuplan
+ else:
+ raise ValueError(f"dataset_type {dataset_type} not supported")
+ return layout
+
+def layout_nuplan(
+ imgs: List[np.array], cam_names: List[str]
+) -> np.array:
+ """Combine cameras into a tiled image for NuPlan dataset.
+ Layout:
+ ##############################################
+ # CAM_L0 # CAM_F0 # CAM_R0 #
+ ##############################################
+ # CAM_L1 # # CAM_R1 #
+ ##############################################
+ # CAM_L2 # CAM_B0 # CAM_R2 #
+ ##############################################
+ """
+ channel = imgs[0].shape[-1]
+ height, width = imgs[0].shape[:2]
+
+ # Create a canvas that's 3 times the height and 3 times the width of a single image
+ tiled_height = height * 3
+ tiled_width = width * 3
+ tiled_img = np.zeros((tiled_height, tiled_width, channel), dtype=np.float32)
+ filled_mask = np.zeros((tiled_height, tiled_width), dtype=np.uint8)
+
+ for idx, cam_name in enumerate(cam_names):
+ img = imgs[idx]
+ if cam_name == "CAM_F0":
+ tiled_img[:height, width:2*width] = img
+ filled_mask[:height, width:2*width] = 1
+ elif cam_name == "CAM_L0":
+ tiled_img[:height, :width] = img
+ filled_mask[:height, :width] = 1
+ elif cam_name == "CAM_R0":
+ tiled_img[:height, 2*width:] = img
+ filled_mask[:height, 2*width:] = 1
+ elif cam_name == "CAM_L1":
+ tiled_img[height:2*height, :width] = img
+ filled_mask[height:2*height, :width] = 1
+ elif cam_name == "CAM_R1":
+ tiled_img[height:2*height, 2*width:] = img
+ filled_mask[height:2*height, 2*width:] = 1
+ elif cam_name == "CAM_L2":
+ tiled_img[2*height:, :width] = img
+ filled_mask[2*height:, :width] = 1
+ elif cam_name == "CAM_R2":
+ tiled_img[2*height:, 2*width:] = img
+ filled_mask[2*height:, 2*width:] = 1
+ elif cam_name == "CAM_B0":
+ tiled_img[2*height:, width:2*width] = img
+ filled_mask[2*height:, width:2*width] = 1
+
+ # Crop the image according to the largest filled area
+ min_y, max_y = np.where(filled_mask)[0].min(), np.where(filled_mask)[0].max()
+ min_x, max_x = np.where(filled_mask)[1].min(), np.where(filled_mask)[1].max()
+ tiled_img = tiled_img[min_y:max_y+1, min_x:max_x+1]
+
+ return tiled_img
+
+def layout_waymo(
+ imgs: List[np.array], cam_names: List[str]
+) -> np.array:
+ """Combine cameras into a tiled image.
+ Layout:
+
+ ######################################################################################
+ # left_camera # front_left_camera # front_camera # front_right_camera # right_camera #
+ ######################################################################################
+ """
+ channel = imgs[0].shape[-1]
+ front_cam_idx = cam_names.index('front_camera')
+ front_img = imgs[front_cam_idx]
+ landscape_width, landscape_height = front_img.shape[1], front_img.shape[0]
+
+ height = landscape_height
+ width = landscape_width * 5
+ tiled_img = np.zeros((height, width, channel), dtype=np.float32)
+ filled_mask = np.zeros((height, width), dtype=np.uint8)
+
+ for idx, cam_name in enumerate(cam_names):
+ img = imgs[idx]
+ if cam_name == "left_camera":
+ tiled_img[landscape_height - img.shape[0]:, :landscape_width] = img
+ filled_mask[landscape_height - img.shape[0]:, :landscape_width] = 1
+ elif cam_name == "front_left_camera":
+ tiled_img[:, landscape_width : 2 * landscape_width] = img
+ filled_mask[:, landscape_width : 2 * landscape_width] = 1
+ elif cam_name == "front_camera":
+ tiled_img[:, 2 * landscape_width : 3 * landscape_width] = img
+ filled_mask[:, 2 * landscape_width : 3 * landscape_width] = 1
+ elif cam_name == "front_right_camera":
+ tiled_img[:, 3 * landscape_width : 4 * landscape_width] = img
+ filled_mask[:, 3 * landscape_width : 4 * landscape_width] = 1
+ elif cam_name == "right_camera":
+ tiled_img[landscape_height - img.shape[0]:, 4 * landscape_width :] = img
+ filled_mask[landscape_height - img.shape[0]:, 4 * landscape_width :] = 1
+
+ # crop the image according to the lagrest filled area
+ min_y, max_y = np.where(filled_mask)[0].min(), np.where(filled_mask)[0].max()
+ min_x, max_x = np.where(filled_mask)[1].min(), np.where(filled_mask)[1].max()
+ tiled_img = tiled_img[min_y:max_y, min_x:max_x]
+ return tiled_img
+
+def layout_nuscenes(
+ imgs: List[np.array], cam_names: List[str]
+) -> np.array:
+ """Combine cameras into a tiled image.
+ Layout:
+
+ ################################################################
+ # CAM_FRONT_LEFT # CAM_FRONT # CAM_FRONT_RIGHT #
+ ################################################################
+ # CAM_BACK_LEFT # CAM_BACK # CAM_BACK_RIGHT #
+ ################################################################
+ """
+ channel = imgs[0].shape[-1]
+ for img in imgs:
+ landscape_width = max(img.shape[0], img.shape[1])
+ landscape_height = min(img.shape[0], img.shape[1])
+ break
+
+ height = landscape_height * 2
+ width = landscape_width * 3
+ tiled_img = np.zeros((height, width, channel), dtype=np.float32)
+ filled_mask = np.zeros((height, width), dtype=np.uint8)
+
+ for idx, cam_name in enumerate(cam_names):
+ img = imgs[idx]
+ if cam_name == "CAM_FRONT_LEFT":
+ tiled_img[:landscape_height, :landscape_width] = img
+ filled_mask[:landscape_height, :landscape_width] = 1
+ elif cam_name == "CAM_FRONT":
+ tiled_img[:landscape_height, landscape_width : 2 * landscape_width] = img
+ filled_mask[:landscape_height, landscape_width : 2 * landscape_width] = 1
+ elif cam_name == "CAM_FRONT_RIGHT":
+ tiled_img[:landscape_height, 2 * landscape_width :] = img
+ filled_mask[:landscape_height, 2 * landscape_width :] = 1
+ elif cam_name == "CAM_BACK_LEFT":
+ tiled_img[landscape_height :, :landscape_width] = img
+ filled_mask[landscape_height :, :landscape_width] = 1
+ elif cam_name == "CAM_BACK":
+ tiled_img[landscape_height :, landscape_width : 2 * landscape_width] = img
+ filled_mask[landscape_height :, landscape_width : 2 * landscape_width] = 1
+ elif cam_name == "CAM_BACK_RIGHT":
+ tiled_img[landscape_height :, 2 * landscape_width :] = img
+ filled_mask[landscape_height :, 2 * landscape_width :] = 1
+
+ # crop the image according to the largest filled area
+ min_y, max_y = np.where(filled_mask)[0].min(), np.where(filled_mask)[0].max()
+ min_x, max_x = np.where(filled_mask)[1].min(), np.where(filled_mask)[1].max()
+ tiled_img = tiled_img[min_y:max_y, min_x:max_x]
+ return tiled_img
+
+def layout_pandaset(
+ imgs: List[np.array], cam_names: List[str]
+) -> np.array:
+ """Combine cameras into a tiled image.
+ Layout:
+
+ ################################################################
+ # front_left_camera # front_camera # front_right_camera #
+ ################################################################
+ # left_camera # back_camera # right_camera #
+ ################################################################
+ """
+ channel = imgs[0].shape[-1]
+ for img in imgs:
+ landscape_width = max(img.shape[0], img.shape[1])
+ landscape_height = min(img.shape[0], img.shape[1])
+ break
+
+ height = landscape_height + landscape_height
+ width = landscape_width + landscape_width + landscape_width
+ tiled_img = np.zeros((height, width, channel), dtype=np.float32)
+ filled_mask = np.zeros((height, width), dtype=np.uint8)
+
+ for idx, cam_name in enumerate(cam_names):
+ img = imgs[idx]
+ if cam_name == "front_left_camera":
+ tiled_img[:landscape_height, :landscape_width] = img
+ filled_mask[:landscape_height, :landscape_width] = 1
+ elif cam_name == "front_camera":
+ tiled_img[:landscape_height, landscape_width : 2 * landscape_width] = img
+ filled_mask[:landscape_height, landscape_width : 2 * landscape_width] = 1
+ elif cam_name == "front_right_camera":
+ tiled_img[:landscape_height, 2 * landscape_width :] = img
+ filled_mask[:landscape_height, 2 * landscape_width :] = 1
+ elif cam_name == "left_camera":
+ tiled_img[landscape_height :, :landscape_width] = img
+ filled_mask[landscape_height :, :landscape_width] = 1
+ elif cam_name == "back_camera":
+ tiled_img[landscape_height :, landscape_width : 2 * landscape_width] = img
+ filled_mask[landscape_height :, landscape_width : 2 * landscape_width] = 1
+ elif cam_name == "right_camera":
+ tiled_img[landscape_height :, 2 * landscape_width :] = img
+ filled_mask[landscape_height :, 2 * landscape_width :] = 1
+
+ # crop the image according to the lagrest filled area
+ min_y, max_y = np.where(filled_mask)[0].min(), np.where(filled_mask)[0].max()
+ min_x, max_x = np.where(filled_mask)[1].min(), np.where(filled_mask)[1].max()
+ tiled_img = tiled_img[min_y:max_y, min_x:max_x]
+ return tiled_img
+
+def layout_kitti(
+ imgs: List[np.array], cam_names: List[str]
+) -> np.array:
+ """Combine cameras into a tiled image.
+ Layout:
+
+ ##############################
+ # CAM_LEFT # CAM_RIGHT #
+ ##############################
+ """
+ channel = imgs[0].shape[-1]
+ height = imgs[0].shape[0]
+ width = imgs[0].shape[1] * 2
+
+ tiled_img = np.zeros((height, width, channel), dtype=np.float32)
+ filled_mask = np.zeros((height, width), dtype=np.uint8)
+
+ for idx, cam_name in enumerate(cam_names):
+ img = imgs[idx]
+ if cam_name == "CAM_LEFT":
+ tiled_img[:, :img.shape[1]] = img
+ filled_mask[:, :img.shape[1]] = 1
+ elif cam_name == "CAM_RIGHT":
+ tiled_img[:, img.shape[1]:] = img
+ filled_mask[:, img.shape[1]:] = 1
+
+ # crop the image according to the largest filled area
+ min_y, max_y = np.where(filled_mask)[0].min(), np.where(filled_mask)[0].max()
+ min_x, max_x = np.where(filled_mask)[1].min(), np.where(filled_mask)[1].max()
+ tiled_img = tiled_img[min_y:max_y, min_x:max_x]
+
+ return tiled_img
+
+def layout_argoverse(
+ imgs: List[np.array], cam_names: List[str]
+) -> np.array:
+ """Combine cameras into a tiled image.
+ Layout:
+
+ ##########################################################
+ # ring_front_left # ring_front_center # ring_front_right #
+ ##########################################################
+ # ring_side_left # # ring_side_right #
+ ##########################################################
+ ############ ring_rear_left # ring_rear_right ############
+ ##########################################################
+ """
+ channel = imgs[0].shape[-1]
+ for img in imgs:
+ landscape_width = max(img.shape[0], img.shape[1])
+ landscape_height = min(img.shape[0], img.shape[1])
+ break
+
+ height = landscape_height + landscape_height + landscape_height
+ width = landscape_width + landscape_height + landscape_width
+ tiled_img = np.zeros((height, width, channel), dtype=np.float32)
+ filled_mask = np.zeros((height, width), dtype=np.uint8)
+
+ for idx, cam_name in enumerate(cam_names):
+ img = imgs[idx]
+ if cam_name == "ring_front_left":
+ tiled_img[:landscape_height, :landscape_width] = img
+ filled_mask[:landscape_height, :landscape_width] = 1
+ elif cam_name == "ring_front_center":
+ tiled_img[:landscape_height, landscape_width : landscape_width + landscape_height] = img[:landscape_height, :]
+ filled_mask[:landscape_height, landscape_width : landscape_width + landscape_height] = 1
+ elif cam_name == "ring_front_right":
+ tiled_img[:landscape_height, landscape_width + landscape_height :] = img
+ filled_mask[:landscape_height, landscape_width + landscape_height :] = 1
+ elif cam_name == "ring_side_left":
+ tiled_img[landscape_height : 2 * landscape_height, :landscape_width] = img
+ filled_mask[landscape_height : 2 * landscape_height, :landscape_width] = 1
+ elif cam_name == "ring_side_right":
+ tiled_img[
+ landscape_height : 2 * landscape_height,
+ landscape_width + landscape_height :
+ ] = img
+ filled_mask[
+ landscape_height : 2 * landscape_height,
+ landscape_width + landscape_height :
+ ] = 1
+ elif cam_name == "ring_rear_left":
+ tiled_img[
+ 2 * landscape_height : 3 * landscape_height, int(0.5 * landscape_height):int(landscape_width + 0.5 * landscape_height)
+ ] = img
+ filled_mask[
+ 2 * landscape_height : 3 * landscape_height, int(0.5 * landscape_height):int(landscape_width + 0.5 * landscape_height)
+ ] = 1
+ elif cam_name == "ring_rear_right":
+ tiled_img[
+ 2 * landscape_height : 3 * landscape_height, int(landscape_width + 0.5 * landscape_height):int(2 * landscape_width + 0.5 * landscape_height)
+ ] = img
+ filled_mask[
+ 2 * landscape_height : 3 * landscape_height, int(landscape_width + 0.5 * landscape_height):int(2 * landscape_width + 0.5 * landscape_height)
+ ] = 1
+
+ # crop the image according to the lagrest filled area
+ min_y, max_y = np.where(filled_mask)[0].min(), np.where(filled_mask)[0].max()
+ min_x, max_x = np.where(filled_mask)[1].min(), np.where(filled_mask)[1].max()
+ tiled_img = tiled_img[min_y:max_y, min_x:max_x]
+ return tiled_img
+
+def dump_3d_bbox_on_image(
+ coords, img,
+ color: Optional[Tuple[int, int, int]] = None,
+ thickness: float = 4,
+) -> None:
+ canvas = img.copy()
+ canvas = cv2.cvtColor(canvas, cv2.COLOR_RGB2BGR)
+ coords = coords.astype(np.int32)
+ for index in range(coords.shape[0]):
+ if isinstance(color, tuple):
+ c = color
+ elif isinstance(color, list):
+ c = color[index]
+ projected_points2d = coords[index]
+ bbox = projected_points2d.tolist()
+ cv2.line(canvas, bbox[0], bbox[1], c, thickness)
+ cv2.line(canvas, bbox[0], bbox[4], c, thickness)
+ cv2.line(canvas, bbox[0], bbox[3], c, thickness)
+ cv2.line(canvas, bbox[1], bbox[2], c, thickness)
+ cv2.line(canvas, bbox[1], bbox[5], c, thickness)
+ cv2.line(canvas, bbox[2], bbox[3], c, thickness)
+ cv2.line(canvas, bbox[2], bbox[6], c, thickness)
+ cv2.line(canvas, bbox[3], bbox[7], c, thickness)
+ cv2.line(canvas, bbox[4], bbox[7], c, thickness)
+ cv2.line(canvas, bbox[4], bbox[5], c, thickness)
+ cv2.line(canvas, bbox[5], bbox[6], c, thickness)
+ cv2.line(canvas, bbox[6], bbox[7], c, thickness)
+ canvas = canvas.astype(np.uint8)
+ canvas = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)
+
+ return canvas
+
+def color_mapper(id: str) -> tuple:
+ # use SHA256 to hash the id
+ hash_object = hashlib.sha256(id.encode())
+ hash_hex = hash_object.hexdigest()
+
+ r = int(hash_hex[0:2], 16)
+ g = int(hash_hex[2:4], 16)
+ b = int(hash_hex[4:6], 16)
+ return (r, g, b)
+
+def sinebow(h):
+ """A cyclic and uniform colormap, see http://basecase.org/env/on-rainbows."""
+ f = lambda x: np.sin(np.pi * x) ** 2
+ return np.stack([f(3 / 6 - h), f(5 / 6 - h), f(7 / 6 - h)], -1)
+
+
+def matte(vis, acc, dark=0.8, light=1.0, width=8):
+ """Set non-accumulated pixels to a Photoshop-esque checker pattern."""
+ bg_mask = np.logical_xor(
+ (np.arange(acc.shape[0]) % (2 * width) // width)[:, None],
+ (np.arange(acc.shape[1]) % (2 * width) // width)[None, :],
+ )
+ bg = np.where(bg_mask, light, dark)
+ return vis * acc[:, :, None] + (bg * (1 - acc))[:, :, None]
+
+
+def weighted_percentile(x, w, ps, assume_sorted=False):
+ """Compute the weighted percentile(s) of a single vector."""
+ x = x.reshape([-1])
+ w = w.reshape([-1])
+ if not assume_sorted:
+ sortidx = np.argsort(x)
+ x, w = x[sortidx], w[sortidx]
+ acc_w = np.cumsum(w)
+ return np.interp(np.array(ps) * (acc_w[-1] / 100), acc_w, x)
+
+
+def visualize_cmap(
+ value,
+ weight,
+ colormap,
+ lo=None,
+ hi=None,
+ percentile=99.0,
+ curve_fn=lambda x: x,
+ modulus=None,
+ matte_background=True,
+):
+ """Visualize a 1D image and a 1D weighting according to some colormap.
+ from mipnerf
+
+ Args:
+ value: A 1D image.
+ weight: A weight map, in [0, 1].
+ colormap: A colormap function.
+ lo: The lower bound to use when rendering, if None then use a percentile.
+ hi: The upper bound to use when rendering, if None then use a percentile.
+ percentile: What percentile of the value map to crop to when automatically
+ generating `lo` and `hi`. Depends on `weight` as well as `value'.
+ curve_fn: A curve function that gets applied to `value`, `lo`, and `hi`
+ before the rest of visualization. Good choices: x, 1/(x+eps), log(x+eps).
+ modulus: If not None, mod the normalized value by `modulus`. Use (0, 1]. If
+ `modulus` is not None, `lo`, `hi` and `percentile` will have no effect.
+ matte_background: If True, matte the image over a checkerboard.
+
+ Returns:
+ A colormap rendering.
+ """
+ # Identify the values that bound the middle of `value' according to `weight`.
+ if lo is None or hi is None:
+ lo_auto, hi_auto = weighted_percentile(
+ value, weight, [50 - percentile / 2, 50 + percentile / 2]
+ )
+ # If `lo` or `hi` are None, use the automatically-computed bounds above.
+ eps = np.finfo(np.float32).eps
+ lo = lo or (lo_auto - eps)
+ hi = hi or (hi_auto + eps)
+
+ # Curve all values.
+ value, lo, hi = [curve_fn(x) for x in [value, lo, hi]]
+
+ # Wrap the values around if requested.
+ if modulus:
+ value = np.mod(value, modulus) / modulus
+ else:
+ # Otherwise, just scale to [0, 1].
+ value = np.nan_to_num(
+ np.clip((value - np.minimum(lo, hi)) / np.abs(hi - lo), 0, 1)
+ )
+ if weight is not None:
+ value *= weight
+ else:
+ weight = np.ones_like(value)
+ if colormap:
+ colorized = colormap(value)[..., :3]
+ else:
+ assert len(value.shape) == 3 and value.shape[-1] == 3
+ colorized = value
+
+ return matte(colorized, weight) if matte_background else colorized
+
+
+def visualize_depth(
+ x, acc=None, lo=None, hi=None, depth_curve_fn=lambda x: -np.log(x + 1e-6)
+):
+ """Visualizes depth maps."""
+ return visualize_cmap(
+ x,
+ acc,
+ cm.get_cmap("turbo"),
+ curve_fn=depth_curve_fn,
+ lo=lo,
+ hi=hi,
+ matte_background=False,
+ )
+
+depth_visualizer = lambda frame, opacity: visualize_depth(
+ frame,
+ opacity,
+ lo=4.0,
+ hi=120,
+ depth_curve_fn=lambda x: -np.log(x + 1e-6),
+)
\ No newline at end of file
diff --git a/worldbench/videogen/reconstruction/README.md b/worldbench/videogen/reconstruction/README.md
new file mode 100644
index 0000000..5868c27
--- /dev/null
+++ b/worldbench/videogen/reconstruction/README.md
@@ -0,0 +1,230 @@
+# Reconstruction Metrics
+
+Reconstruction metrics use the same WorldLens evaluator flow as the other
+video-generation metrics: enable metrics in the YAML config, then run
+`tools/evaluate.py`.
+
+WorldLens owns metric calculation and aggregation. DriveStudio is only used to
+prepare reconstruction assets before evaluation. The DriveStudio source needed
+for this path is vendored in `worldbench/third_party/drivestudio`.
+
+## Environment
+
+Start from the normal WorldLens environment described in the project README.
+Reconstruction metric computation runs in that environment.
+
+To train/render 4DGS assets with the vendored DriveStudio backend, install only
+the Reconstruction extras on top of the WorldLens environment:
+
+```bash
+cd /path/to/WorldLens
+pip install -r worldbench/videogen/reconstruction/requirements-reconstruction-extra.txt
+```
+
+Do not install
+`worldbench/third_party/drivestudio/requirements.txt` into
+the WorldLens environment. That file pins a separate DriveStudio stack,
+including older torch/torchvision/CUDA versions, and can conflict with the
+WorldLens CUDA extensions.
+
+The base WorldLens environment already includes several dependencies used by
+this path, such as `gsplat`, `lpips`, `nuscenes-devkit`, `open3d`,
+`pyquaternion`, and `torchmetrics`. The extra file adds the missing runtime
+pieces without changing the WorldLens torch/CUDA version. If `pytorch3d` is not
+already available in your environment, install a wheel that matches the
+WorldLens PyTorch/CUDA build before running 4DGS training.
+
+## Asset Contract
+
+Reconstruction metrics expect this layout:
+
+```text
+generated_results/reconstruction/
+ magicdrive/
+ manifest.json
+ ckpt/081/config.yaml
+ ckpt/081/checkpoint_final.pth
+ train/081/pred/0000_CAM_FRONT.jpg
+ train/081/gt/0000_CAM_FRONT.jpg
+ depth/081/pred/00000_CAM_FRONT.npy
+ depth/081/gt/00000_CAM_FRONT.npy
+ depth/081/masks/00000_CAM_FRONT.png
+ novel/081/front_center_interp.mp4
+ novel/081/s_curve.mp4
+ novel/081/lateral_offset.mp4
+ novel/081/lateral_offset_left.mp4
+```
+
+Legacy `train//metrics.json` and `depth//metrics.json` files are
+still accepted as a fallback when train/depth assets are not available, but the
+recommended path is asset-based metric computation inside WorldLens.
+
+## Evaluate Assets
+
+Use this after 4DGS assets have been prepared.
+
+```yaml
+dimensions:
+ reconstruction:
+ photometric_error:
+ - name: photometric_error
+ method_name: magicdrive
+ reconstruction_root: generated_results/reconstruction
+ clips: ["081"]
+
+ geometric_discrepancy:
+ - name: geometric_discrepancy
+ method_name: magicdrive
+ reconstruction_root: generated_results/reconstruction
+ clips: ["081"]
+
+ novel_view_fidelity:
+ - name: novel_view_fidelity
+ method_name: magicdrive
+ reconstruction_root: generated_results/reconstruction
+ clips: ["081"]
+
+ novel_view_consistency:
+ - name: novel_view_consistency
+ method_name: magicdrive
+ reconstruction_root: generated_results/reconstruction
+ clips: ["081"]
+ local_save_path: pretrained_models/clip/ViT-B-32.pt
+```
+
+Run through the normal WorldLens entrypoint:
+
+```bash
+cd /path/to/WorldLens
+python tools/evaluate.py method_name=magicdrive
+```
+
+## Prepare Assets With DriveStudio
+
+Use this when new generated videos need to be reconstructed with DriveStudio 4DGS
+before WorldLens evaluates them.
+
+```bash
+cd /path/to/WorldLens
+python worldbench/videogen/reconstruction/prepare_assets.py \
+ --method-name magicdrive \
+ --clips 081 \
+ --reconstruction-root generated_results/reconstruction \
+ --work-root work_dirs \
+ --data-root data/nuscenes_trainval/raw \
+ --target-dir data/nuscenes_trainval/processed
+```
+
+The exporter preprocesses generated videos, trains 4DGS, renders train/depth and
+novel-view outputs, exports standard assets, and writes a manifest. WorldLens
+then computes photometric, geometric, and novel-view metrics from those assets.
+
+The default 4DGS training config is
+`worldbench/third_party/drivestudio/configs/omnire.yaml`.
+Override it with `--config-file`; relative config paths are resolved from the
+vendored DriveStudio root:
+
+```bash
+python worldbench/videogen/reconstruction/prepare_assets.py \
+ --method-name magicdrive \
+ --clips 081 \
+ --reconstruction-root generated_results/reconstruction \
+ --work-root work_dirs \
+ --data-root data/nuscenes_trainval/raw \
+ --target-dir data/nuscenes_trainval/processed \
+ --config-file configs/streetgs.yaml
+```
+
+Available vendored configs include `configs/omnire.yaml`,
+`configs/streetgs.yaml`, `configs/pvg.yaml`, and `configs/deformablegs.yaml`.
+The backend uses DriveStudio's `nuscenes_gen/6cams` dataset config for generated
+video reconstruction and derives the training clip range from `--clips`.
+
+For reference, the wrapper launches DriveStudio training with the same effective
+arguments as:
+
+```bash
+cd worldbench/third_party/drivestudio
+python tools/train.py \
+ --config_file configs/omnire.yaml \
+ --output_root /path/to/WorldLens/work_dirs \
+ --project drivestudio-nus-magicdrive \
+ --run_name -1 \
+ --start_idx 81 \
+ --end_idx 81 \
+ dataset=nuscenes_gen/6cams \
+ data.scene_idx=-1 \
+ data.start_timestep=0 \
+ data.end_timestep=-1 \
+ data.data_root=/path/to/WorldLens/data/nuscenes_trainval/processed_magicdrive/advanced_12Hz_trainval
+```
+
+DriveStudio renders novel-view videos; WorldLens computes novel-view fidelity
+and consistency from those videos.
+
+The DriveStudio source code is vendored under `worldbench/third_party`, so users
+do not need a separate DriveStudio checkout. By default the backend runs with
+the current WorldLens Python. Use `--python-bin` only if you deliberately want
+to run the backend in a separate Python environment.
+
+## New Data Requirements
+
+Generated videos should use the DriveStudio `nuscenes_gen` layout for the
+selected `method_name`.
+
+Before 4DGS training, each requested clip must have:
+
+- processed images
+- `fine_dynamic_masks/all`
+- `fine_dynamic_masks/human`
+- `fine_dynamic_masks/vehicle`
+- `humanpose/smpl.pkl`
+
+The backend validates these inputs before launching training. Missing assets
+fail early with a direct error message.
+
+`geometric_discrepancy` is not a pure RGB-video metric. It requires reference/GT
+depth assets and Grounded-SAM2 evaluation masks in addition to generated videos.
+The DriveStudio exporter copies GT depth from `--gt-method`, which defaults to
+`gt`, and copies masks from `processed_`.
+
+Generate the evaluation masks once from GT processed images and reuse them for
+all methods. Run this in an environment where Grounded-SAM-2, GroundingDINO,
+SAM2, and their checkpoints are available:
+
+```bash
+cd /path/to/WorldLens
+python worldbench/videogen/reconstruction/generate_eval_masks.py \
+ --grounded-sam-root /path/to/Grounded-SAM-2 \
+ --processed-root data/nuscenes_trainval/processed_gt/advanced_12Hz_trainval \
+ --clips 081 \
+ --skip-existing
+```
+
+This writes
+`data/nuscenes_trainval/processed_gt/advanced_12Hz_trainval//sam_mask/.png`.
+The masks use the paper prompts `road surface` and `vehicles`. They are only for
+geometric discrepancy; DriveStudio's `fine_dynamic_masks/{all,human,vehicle}`
+remain training inputs and are not WorldLens metric masks.
+
+If checkpoints already exist and only rendering/export is needed, skip the
+preprocess and training phases:
+
+```bash
+cd /path/to/WorldLens
+python worldbench/videogen/reconstruction/prepare_assets.py \
+ --method-name magicdrive \
+ --clips 081 \
+ --reconstruction-root generated_results/reconstruction \
+ --work-root work_dirs \
+ --data-root data/nuscenes_trainval/raw \
+ --target-dir data/nuscenes_trainval/processed \
+ --skip-preprocess \
+ --skip-train
+```
+
+If all renders and raw depths already exist and only the WorldLens asset layout
+needs to be refreshed, add `--skip-eval` as well.
+
+Use `--drivestudio-root /path/to/drivestudio` only to override the vendored
+source during backend development.
diff --git a/worldbench/videogen/reconstruction/_asset_metrics.py b/worldbench/videogen/reconstruction/_asset_metrics.py
new file mode 100644
index 0000000..8b4c0d1
--- /dev/null
+++ b/worldbench/videogen/reconstruction/_asset_metrics.py
@@ -0,0 +1,292 @@
+import json
+import math
+from pathlib import Path
+
+import numpy as np
+
+
+def _clip_ids(clips):
+ if clips is None:
+ return None
+ return [str(clip) for clip in clips]
+
+
+def _read_legacy_metrics(path):
+ if not path.is_file():
+ return None
+ payload = json.loads(path.read_text())
+ return payload["metrics"]
+
+
+def _load_rgb(path):
+ from PIL import Image
+
+ return np.asarray(Image.open(path).convert("RGB"), dtype=np.float32) / 255.0
+
+
+def _load_mask(path, shape):
+ from PIL import Image
+
+ if not path.is_file():
+ return None
+ mask = np.asarray(Image.open(path).convert("L")) > 0
+ if mask.shape != shape[:2]:
+ mask = _resize_mask_nearest(mask, shape[:2])
+ return mask
+
+
+def _resize_mask_nearest(mask, target_shape):
+ if mask.shape == target_shape:
+ return mask
+ src_h, src_w = mask.shape
+ tgt_h, tgt_w = target_shape
+ if tgt_h <= 0 or tgt_w <= 0:
+ raise ValueError(f"Invalid target mask shape: {target_shape}")
+ y_idx = np.linspace(0, src_h - 1, tgt_h).round().astype(int)
+ x_idx = np.linspace(0, src_w - 1, tgt_w).round().astype(int)
+ return mask[y_idx][:, x_idx]
+
+
+def _load_required_mask(path, shape):
+ mask = _load_mask(path, shape)
+ if mask is None:
+ raise FileNotFoundError(f"Missing geometric evaluation mask: {path}")
+ return mask
+
+
+def _psnr(pred, gt):
+ mse = float(np.mean((pred - gt) ** 2))
+ if mse <= 1e-12:
+ return float("inf")
+ return -10.0 * math.log10(mse)
+
+
+def _ssim(pred, gt):
+ from skimage.metrics import structural_similarity
+
+ return float(structural_similarity(pred, gt, data_range=1.0, channel_axis=-1))
+
+
+def _mean(values):
+ finite = [float(value) for value in values if value is not None and not np.isnan(value)]
+ if not finite:
+ return float("nan")
+ return float(np.mean(finite))
+
+
+def _discover_clips(root, asset_type):
+ asset_root = root / asset_type
+ if not asset_root.is_dir():
+ raise FileNotFoundError(f"Missing Reconstruction {asset_type} asset directory: {asset_root}")
+ clips = sorted(path.name for path in asset_root.iterdir() if path.is_dir())
+ if not clips:
+ raise FileNotFoundError(f"No Reconstruction {asset_type} clips found under {asset_root}")
+ return clips
+
+
+class AssetPhotometricMetric:
+ def __init__(
+ self,
+ method_name,
+ reconstruction_root,
+ clips,
+ compute_lpips=True,
+ **kwargs,
+ ):
+ self.method_name = method_name
+ self.reconstruction_root = Path(reconstruction_root)
+ self.clips = _clip_ids(clips) or _discover_clips(self.reconstruction_root / self.method_name, "train")
+ self.compute_lpips = compute_lpips
+
+ def _legacy_path(self, clip):
+ return self.reconstruction_root / self.method_name / "train" / clip / "metrics.json"
+
+ def _clip_dir(self, clip):
+ return self.reconstruction_root / self.method_name / "train" / clip
+
+ def _clip_assets_ready(self, clip):
+ clip_dir = self._clip_dir(clip)
+ return any((clip_dir / "pred").glob("*")) and any((clip_dir / "gt").glob("*"))
+
+ def _assets_ready(self):
+ for clip in self.clips:
+ if not self._clip_assets_ready(clip):
+ return False
+ return True
+
+ def _lpips_model(self):
+ if not self.compute_lpips:
+ return None, None, None
+ import torch
+
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+ try:
+ from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+
+ return LearnedPerceptualImagePatchSimilarity(normalize=True).to(device).eval(), device, "torchmetrics"
+ except ImportError:
+ import lpips
+
+ return lpips.LPIPS(net="alex").to(device).eval(), device, "lpips"
+
+ def _lpips_score(self, model, device, backend, pred, gt):
+ if model is None:
+ return None
+ import torch
+
+ pred_t = torch.from_numpy(pred).permute(2, 0, 1).unsqueeze(0).to(device)
+ gt_t = torch.from_numpy(gt).permute(2, 0, 1).unsqueeze(0).to(device)
+ if backend == "lpips":
+ pred_t = pred_t * 2.0 - 1.0
+ gt_t = gt_t * 2.0 - 1.0
+ with torch.no_grad():
+ return float(model(pred_t, gt_t).mean().item())
+
+ def _compute_clip(self, clip, lpips_model, lpips_device, lpips_backend):
+ clip_dir = self._clip_dir(clip)
+ pred_dir = clip_dir / "pred"
+ gt_dir = clip_dir / "gt"
+ pred_files = sorted(path for path in pred_dir.glob("*") if path.suffix.lower() in {".jpg", ".jpeg", ".png"})
+ if not pred_files:
+ raise FileNotFoundError(f"No photometric prediction images found under {pred_dir}")
+
+ psnr_values = []
+ ssim_values = []
+ lpips_values = []
+
+ for pred_path in pred_files:
+ gt_path = gt_dir / pred_path.name
+ if not gt_path.is_file():
+ raise FileNotFoundError(f"Missing GT image for {pred_path.name}: {gt_path}")
+ pred = _load_rgb(pred_path)
+ gt = _load_rgb(gt_path)
+ psnr_values.append(_psnr(pred, gt))
+ ssim_values.append(_ssim(pred, gt))
+ lpips_value = self._lpips_score(lpips_model, lpips_device, lpips_backend, pred, gt)
+ if lpips_value is not None:
+ lpips_values.append(lpips_value)
+
+ metrics = {
+ "image_metrics/full/psnr": _mean(psnr_values),
+ "image_metrics/full/ssim": _mean(ssim_values),
+ }
+ if lpips_values:
+ metrics["image_metrics/full/lpips"] = _mean(lpips_values)
+ return metrics
+
+ def __call__(self):
+ lpips_model = lpips_device = lpips_backend = None
+ result = {}
+ for clip in self.clips:
+ legacy = _read_legacy_metrics(self._legacy_path(clip))
+ if legacy is not None and not self._clip_assets_ready(clip):
+ result[clip] = legacy
+ continue
+ if lpips_model is None and self.compute_lpips:
+ lpips_model, lpips_device, lpips_backend = self._lpips_model()
+ result[clip] = self._compute_clip(clip, lpips_model, lpips_device, lpips_backend)
+ return result
+
+
+def _load_depth(path):
+ depth = np.load(path)
+ if depth.ndim == 3 and depth.shape[-1] == 1:
+ depth = depth[..., 0]
+ return depth.astype(np.float64)
+
+
+class AssetGeometricMetric:
+ def __init__(
+ self,
+ method_name,
+ reconstruction_root,
+ clips,
+ **kwargs,
+ ):
+ self.method_name = method_name
+ self.reconstruction_root = Path(reconstruction_root)
+ self.clips = _clip_ids(clips) or _discover_clips(self.reconstruction_root / self.method_name, "depth")
+
+ def _legacy_path(self, clip):
+ return self.reconstruction_root / self.method_name / "depth" / clip / "metrics.json"
+
+ def _clip_dir(self, clip):
+ return self.reconstruction_root / self.method_name / "depth" / clip
+
+ def _clip_assets_ready(self, clip):
+ clip_dir = self._clip_dir(clip)
+ return any((clip_dir / "pred").glob("*.npy")) and any((clip_dir / "gt").glob("*.npy"))
+
+ def _assets_ready(self):
+ for clip in self.clips:
+ if not self._clip_assets_ready(clip):
+ return False
+ return True
+
+ def _compute_clip(self, clip):
+ clip_dir = self._clip_dir(clip)
+ pred_files = sorted((clip_dir / "pred").glob("*.npy"))
+ if not pred_files:
+ raise FileNotFoundError(f"No depth predictions found under {clip_dir / 'pred'}")
+
+ sq_sum = 0.0
+ log_sq_sum = 0.0
+ abs_rel_sum = 0.0
+ sq_rel_sum = 0.0
+ delta1 = delta2 = delta3 = 0
+ pixels = 0
+ rel_pixels = 0
+ eps = 1e-6
+
+ for pred_path in pred_files:
+ gt_path = clip_dir / "gt" / pred_path.name
+ if not gt_path.is_file():
+ raise FileNotFoundError(f"Missing GT depth for {pred_path.name}: {gt_path}")
+ pred = _load_depth(pred_path)
+ gt = _load_depth(gt_path)
+ mask = _load_required_mask(clip_dir / "masks" / pred_path.with_suffix(".png").name, pred.shape)
+ valid = mask.copy()
+ if not valid.any():
+ continue
+ pred_values = pred[valid]
+ gt_values = gt[valid]
+ diff = pred_values - gt_values
+ sq_sum += float((diff ** 2).sum())
+ pixels += int(valid.sum())
+ rel_mask = gt_values > eps
+ if rel_mask.any():
+ pred_rel = pred_values[rel_mask]
+ gt_rel = gt_values[rel_mask]
+ rel_diff = pred_rel - gt_rel
+ log_sq_sum += float(((np.log(np.maximum(pred_rel, eps)) - np.log(np.maximum(gt_rel, eps))) ** 2).sum())
+ abs_rel_sum += float((np.abs(rel_diff) / gt_rel).sum())
+ sq_rel_sum += float(((rel_diff ** 2) / gt_rel).sum())
+ ratio = np.maximum(np.maximum(pred_rel, eps) / gt_rel, gt_rel / np.maximum(pred_rel, eps))
+ delta1 += int((ratio < 1.25).sum())
+ delta2 += int((ratio < 1.25 ** 2).sum())
+ delta3 += int((ratio < 1.25 ** 3).sum())
+ rel_pixels += int(rel_mask.sum())
+
+ if pixels == 0:
+ raise ValueError(f"No valid depth pixels for method={self.method_name}, clip={clip}")
+ if rel_pixels == 0:
+ raise ValueError(f"No valid relative depth pixels for method={self.method_name}, clip={clip}")
+ return {
+ "global_rmse": math.sqrt(sq_sum / pixels),
+ "global_log_rmse": math.sqrt(log_sq_sum / rel_pixels),
+ "global_abs_rel": abs_rel_sum / rel_pixels,
+ "global_sq_rel": sq_rel_sum / rel_pixels,
+ "global_delta1": delta1 / rel_pixels,
+ "global_delta2": delta2 / rel_pixels,
+ "global_delta3": delta3 / rel_pixels,
+ }
+
+ def __call__(self):
+ result = {}
+ for clip in self.clips:
+ legacy = _read_legacy_metrics(self._legacy_path(clip))
+ if legacy is not None and not self._clip_assets_ready(clip):
+ result[clip] = legacy
+ continue
+ result[clip] = self._compute_clip(clip)
+ return result
diff --git a/worldbench/videogen/reconstruction/_novel_view_contract.py b/worldbench/videogen/reconstruction/_novel_view_contract.py
new file mode 100644
index 0000000..9be6041
--- /dev/null
+++ b/worldbench/videogen/reconstruction/_novel_view_contract.py
@@ -0,0 +1,57 @@
+from pathlib import Path
+from typing import Dict, Iterable, List, Optional
+
+
+def _clip_ids(clips: Optional[Iterable[str]]) -> Optional[List[str]]:
+ if clips is None:
+ return None
+ return [str(clip) for clip in clips]
+
+
+def novel_view_groups(
+ method_name: str,
+ generated_data_path: str,
+ reconstruction_root: Optional[str],
+ clips: Optional[Iterable[str]],
+) -> Dict[str, List[str]]:
+ if reconstruction_root is not None:
+ root = Path(reconstruction_root) / method_name / "novel"
+ clip_list = _clip_ids(clips)
+ if clip_list is None:
+ clip_list = sorted(path.name for path in root.iterdir() if path.is_dir())
+ groups = {}
+ missing = []
+ for clip in clip_list:
+ video_dir = root / str(clip)
+ videos = sorted(str(path) for path in video_dir.glob("*.mp4"))
+ if videos:
+ groups[str(clip)] = videos
+ else:
+ missing.append(str(video_dir))
+ if missing:
+ details = "\n".join(f"- {path}" for path in missing[:20])
+ raise FileNotFoundError(f"Missing novel-view videos for requested clips:\n{details}")
+ if not groups:
+ raise FileNotFoundError(f"No novel-view videos found under {root}")
+ return groups
+
+ legacy_root = Path(generated_data_path) / method_name / "reconstruction"
+ groups = {}
+ for dim_dir in sorted(path for path in legacy_root.iterdir() if path.is_dir()):
+ videos = sorted(str(path) for path in dim_dir.rglob("*.mp4"))
+ if videos:
+ groups[dim_dir.name] = videos
+ if not groups:
+ raise FileNotFoundError(f"No novel-view videos found under {legacy_root}")
+ return groups
+
+
+def novel_metric_output_dir(
+ method_name: str,
+ generated_data_path: str,
+ reconstruction_root: Optional[str],
+ metric_name: str,
+) -> Path:
+ if reconstruction_root is not None:
+ return Path(reconstruction_root) / method_name / metric_name
+ return Path(generated_data_path) / method_name / metric_name
diff --git a/worldbench/videogen/reconstruction/generate_eval_masks.py b/worldbench/videogen/reconstruction/generate_eval_masks.py
new file mode 100644
index 0000000..973aeb5
--- /dev/null
+++ b/worldbench/videogen/reconstruction/generate_eval_masks.py
@@ -0,0 +1,207 @@
+#!/usr/bin/env python3
+"""Generate GT Grounded-SAM2 masks for Reconstruction geometric evaluation."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from collections import defaultdict
+from pathlib import Path
+from typing import Iterable, List, Optional, Sequence, Tuple
+
+import numpy as np
+from PIL import Image
+
+
+IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"}
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(
+ "Generate Grounded-SAM2 road/vehicle masks for WorldLens Reconstruction geometric evaluation"
+ )
+ parser.add_argument(
+ "--processed-root",
+ type=Path,
+ default=Path("data/nuscenes_trainval/processed_gt/advanced_12Hz_trainval"),
+ help="GT processed split root that contains /images.",
+ )
+ parser.add_argument("--clips", nargs="+", help="Clip ids to process, for example 081 082.")
+ parser.add_argument("--start", type=int, default=None, help="Start clip index in sorted processed-root clips.")
+ parser.add_argument("--end", type=int, default=None, help="Inclusive end clip index in sorted processed-root clips.")
+ parser.add_argument("--output-subdir", default="sam_mask", help="Output directory name inside each clip.")
+ parser.add_argument("--prompts", nargs="+", default=["road surface", "vehicles"])
+ parser.add_argument("--box-threshold", type=float, default=0.35)
+ parser.add_argument("--text-threshold", type=float, default=0.25)
+ parser.add_argument("--device", default="cuda")
+ parser.add_argument("--skip-existing", action="store_true")
+ parser.add_argument(
+ "--grounded-sam-root",
+ type=Path,
+ default=None,
+ help="Grounded-SAM-2 checkout root. Relative model/config paths are resolved from this root.",
+ )
+ parser.add_argument("--sam2-checkpoint", default="checkpoints/sam2.1_hiera_base_plus.pt")
+ parser.add_argument("--sam2-config", default="configs/sam2.1/sam2.1_hiera_b+.yaml")
+ parser.add_argument(
+ "--grounding-dino-config",
+ default="grounding_dino/groundingdino/config/GroundingDINO_SwinB_cfg.py",
+ )
+ parser.add_argument("--grounding-dino-checkpoint", default="gdino_checkpoints/groundingdino_swinb_cogcoor.pth")
+ return parser.parse_args()
+
+
+def resolve_path(root: Optional[Path], value: str) -> str:
+ path = Path(value)
+ if path.is_absolute() or root is None:
+ return str(path)
+ return str(root / path)
+
+
+def configure_import_path(grounded_sam_root: Optional[Path]) -> None:
+ if grounded_sam_root is None:
+ return
+ root = grounded_sam_root.resolve()
+ sys.path.insert(0, str(root))
+
+
+def load_models(args: argparse.Namespace):
+ configure_import_path(args.grounded_sam_root)
+
+ from grounding_dino.groundingdino.util.inference import load_model
+ from sam2.build_sam import build_sam2
+ from sam2.sam2_image_predictor import SAM2ImagePredictor
+
+ sam2_checkpoint = resolve_path(args.grounded_sam_root, args.sam2_checkpoint)
+ grounding_config = resolve_path(args.grounded_sam_root, args.grounding_dino_config)
+ grounding_checkpoint = resolve_path(args.grounded_sam_root, args.grounding_dino_checkpoint)
+
+ sam2_model = build_sam2(args.sam2_config, sam2_checkpoint, device=args.device)
+ sam2_predictor = SAM2ImagePredictor(sam2_model)
+
+ grounding_model = load_model(
+ model_config_path=grounding_config,
+ model_checkpoint_path=grounding_checkpoint,
+ device=args.device,
+ )
+ grounding_model.eval()
+ return grounding_model, sam2_predictor
+
+
+def discover_clips(
+ processed_root: Path,
+ clips: Optional[Sequence[str]],
+ start: Optional[int],
+ end: Optional[int],
+) -> List[Path]:
+ if clips:
+ selected = [processed_root / f"{int(clip):03d}" for clip in clips]
+ else:
+ selected = sorted(path for path in processed_root.iterdir() if path.is_dir())
+ if start is not None or end is not None:
+ selected = selected[start : None if end is None else end + 1]
+ missing = [str(path) for path in selected if not (path / "images").is_dir()]
+ if missing:
+ details = "\n".join(f"- {path}/images" for path in missing[:20])
+ raise FileNotFoundError(f"Missing clip image directories under {processed_root}:\n{details}")
+ return selected
+
+
+def iter_images(images_dir: Path) -> List[Path]:
+ return sorted(path for path in images_dir.iterdir() if path.suffix.lower() in IMAGE_SUFFIXES)
+
+
+def run_grounded_sam2(
+ image_path: Path,
+ prompts: Sequence[str],
+ grounding_model,
+ sam2_predictor,
+ box_threshold: float,
+ text_threshold: float,
+) -> np.ndarray:
+ import torch
+ from grounding_dino.groundingdino.util.inference import load_image, predict
+ from torchvision.ops import box_convert
+
+ image_source, image = load_image(str(image_path))
+ sam2_predictor.set_image(image_source)
+ h, w = image_source.shape[:2]
+ combined = np.zeros((h, w), dtype=bool)
+
+ with torch.no_grad():
+ for prompt in prompts:
+ boxes, _, _ = predict(
+ model=grounding_model,
+ image=image,
+ caption=prompt.rstrip(".") + ".",
+ box_threshold=box_threshold,
+ text_threshold=text_threshold,
+ )
+ if boxes.shape[0] == 0:
+ continue
+ boxes = boxes * torch.tensor([w, h, w, h], dtype=boxes.dtype, device=boxes.device)
+ input_boxes = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").cpu().numpy()
+ masks, _, _ = sam2_predictor.predict(box=input_boxes, multimask_output=False)
+ if masks.ndim == 4:
+ masks = np.squeeze(masks, axis=1)
+ combined |= masks.astype(bool).any(axis=0)
+
+ return combined
+
+
+def progress(items: Iterable[object], desc: str):
+ try:
+ from tqdm import tqdm
+
+ return tqdm(items, desc=desc)
+ except ImportError:
+ return items
+
+
+def generate_masks(args: argparse.Namespace) -> Tuple[int, int]:
+ processed_root = args.processed_root.resolve()
+ clips = discover_clips(processed_root, args.clips, args.start, args.end)
+
+ written = 0
+ skipped = 0
+ pending_by_clip = defaultdict(list)
+ for clip_dir in clips:
+ image_paths = iter_images(clip_dir / "images")
+ if not image_paths:
+ raise FileNotFoundError(f"No images found in {clip_dir / 'images'}")
+ output_dir = clip_dir / args.output_subdir
+ output_dir.mkdir(parents=True, exist_ok=True)
+ for image_path in image_paths:
+ output_path = output_dir / f"{image_path.stem}.png"
+ if output_path.exists() and args.skip_existing:
+ skipped += 1
+ continue
+ pending_by_clip[clip_dir.name].append((image_path, output_path))
+
+ if not pending_by_clip:
+ return written, skipped
+
+ grounding_model, sam2_predictor = load_models(args)
+ for clip_name, pending_items in pending_by_clip.items():
+ for image_path, output_path in progress(pending_items, desc=clip_name):
+ mask = run_grounded_sam2(
+ image_path=image_path,
+ prompts=args.prompts,
+ grounding_model=grounding_model,
+ sam2_predictor=sam2_predictor,
+ box_threshold=args.box_threshold,
+ text_threshold=args.text_threshold,
+ )
+ Image.fromarray(mask.astype(np.uint8) * 255).save(output_path)
+ written += 1
+ return written, skipped
+
+
+def main() -> None:
+ args = parse_args()
+ written, skipped = generate_masks(args)
+ print(f"Generated {written} masks; skipped {skipped} existing masks.")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/videogen/reconstruction/geometric_discrepancy/__init__.py b/worldbench/videogen/reconstruction/geometric_discrepancy/__init__.py
new file mode 100644
index 0000000..0a6236c
--- /dev/null
+++ b/worldbench/videogen/reconstruction/geometric_discrepancy/__init__.py
@@ -0,0 +1 @@
+from .geometric_discrepancy import GEOMETRIC_DISCREPANCY
diff --git a/worldbench/videogen/reconstruction/geometric_discrepancy/geometric_discrepancy.py b/worldbench/videogen/reconstruction/geometric_discrepancy/geometric_discrepancy.py
new file mode 100644
index 0000000..af6a478
--- /dev/null
+++ b/worldbench/videogen/reconstruction/geometric_discrepancy/geometric_discrepancy.py
@@ -0,0 +1,5 @@
+from worldbench.videogen.reconstruction._asset_metrics import AssetGeometricMetric
+
+
+class GEOMETRIC_DISCREPANCY(AssetGeometricMetric):
+ pass
diff --git a/worldbench/videogen/reconstruction/novel_view_consistency/novel_view_consistency.py b/worldbench/videogen/reconstruction/novel_view_consistency/novel_view_consistency.py
index 45df02b..51ffccd 100644
--- a/worldbench/videogen/reconstruction/novel_view_consistency/novel_view_consistency.py
+++ b/worldbench/videogen/reconstruction/novel_view_consistency/novel_view_consistency.py
@@ -11,6 +11,10 @@
from torch.nn import functional as F
from ....utils import common, video_relative
from worldbench.videogen.generation.temporal_consistency.utils import clip_transform
+from worldbench.videogen.reconstruction._novel_view_contract import (
+ novel_metric_output_dir,
+ novel_view_groups,
+)
def optical_flow_warp(img1, img2):
"""warp img2 to img1 using optical flow"""
@@ -28,12 +32,16 @@ def optical_flow_warp(img1, img2):
class NOVEL_VIEW_CONSISTENCY:
def __init__(self, method_name,
generated_data_path="generated_results",
+ reconstruction_root=None,
+ clips=None,
repeat_times=1,
local_save_path=None,
repo_or_dir=None,
**kwargs):
self.method_name = method_name
self.generated_data_path = generated_data_path
+ self.reconstruction_root = reconstruction_root
+ self.clips = clips
self.repeat_times = repeat_times
self.local_save_path = local_save_path
self.repo_or_dir = repo_or_dir
@@ -106,14 +114,21 @@ def calculate_single_time(self, video_list):
return sim_per_frame, lpips_per_frame, video_results
def __call__(self):
-
- reconstruction_video_folder = os.path.join(self.generated_data_path, self.method_name, 'reconstruction')
- dims = os.listdir(reconstruction_video_folder)
- rel_save_path = os.path.join(self.generated_data_path, self.method_name, 'novel_view_consistency')
+ groups = novel_view_groups(
+ method_name=self.method_name,
+ generated_data_path=self.generated_data_path,
+ reconstruction_root=self.reconstruction_root,
+ clips=self.clips,
+ )
+ rel_save_path = novel_metric_output_dir(
+ method_name=self.method_name,
+ generated_data_path=self.generated_data_path,
+ reconstruction_root=self.reconstruction_root,
+ metric_name="novel_view_consistency",
+ )
os.makedirs(rel_save_path, exist_ok=True)
- for dim in dims:
- video_folder_path = os.path.join(reconstruction_video_folder, dim)
- video_list = common.find_videos_in_dir(video_folder_path, extension=".mp4")
+ all_results = {}
+ for dim, video_list in groups.items():
sim_per_frame, lpips_per_frame, video_results = self.calculate_single_time(video_list)
logger.info(f"Method: {self.method_name}, Dim: {dim}, Subject Consistency (per frame) by DINO: {sim_per_frame:.4f}, LPIPS: {lpips_per_frame:.4f}")
result = {
@@ -123,7 +138,9 @@ def __call__(self):
'lpips_per_frame': lpips_per_frame,
'video_results': video_results
}
+ all_results[dim] = result
save_path = os.path.join(rel_save_path, f'{dim}.json')
with open(save_path, 'w') as f:
json.dump(result, f, indent=4)
- logger.info(f"Results saved to {save_path}")
\ No newline at end of file
+ logger.info(f"Results saved to {save_path}")
+ return all_results
diff --git a/worldbench/videogen/reconstruction/novel_view_fidelity/novel_view_fidelity_iq.py b/worldbench/videogen/reconstruction/novel_view_fidelity/novel_view_fidelity_iq.py
index 63aec36..f1d0c9b 100644
--- a/worldbench/videogen/reconstruction/novel_view_fidelity/novel_view_fidelity_iq.py
+++ b/worldbench/videogen/reconstruction/novel_view_fidelity/novel_view_fidelity_iq.py
@@ -7,11 +7,14 @@
import subprocess
from pyiqa.archs.musiq_arch import MUSIQ
from loguru import logger
-from worldbench.utils import common
from worldbench.utils import video_relative
from tqdm import tqdm
import torch
from torchvision import transforms
+from worldbench.videogen.reconstruction._novel_view_contract import (
+ novel_metric_output_dir,
+ novel_view_groups,
+)
def transform(images, preprocess_mode='shorter'):
if preprocess_mode.startswith('shorter'):
@@ -37,14 +40,16 @@ def transform(images, preprocess_mode='shorter'):
class NOVEL_VIEW_FIDELITY_IQ:
def __init__(self, method_name,
- need_preprocessing = False,
generated_data_path = "generated_results",
+ reconstruction_root = None,
+ clips = None,
pretrained_model_path = "pretrained_models/musiq/musiq_spaq_ckpt-358bb6af.pth",
**kwargs):
self.method_name = method_name
- self.need_preprocessing = need_preprocessing
self.generated_data_path = generated_data_path
+ self.reconstruction_root = reconstruction_root
+ self.clips = clips
self.pretrained_model_path = pretrained_model_path
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.model = self.init_model()
@@ -81,17 +86,28 @@ def calculate_iq_metrics(self, video_list, **kwargs):
return average_score, video_results
def __call__(self, *args, **kwds):
- reconstruction_video_folder = os.path.join(self.generated_data_path, self.method_name, 'reconstruction')
- dims = os.listdir(reconstruction_video_folder)
+ groups = novel_view_groups(
+ method_name=self.method_name,
+ generated_data_path=self.generated_data_path,
+ reconstruction_root=self.reconstruction_root,
+ clips=self.clips,
+ )
+ result_dir = novel_metric_output_dir(
+ method_name=self.method_name,
+ generated_data_path=self.generated_data_path,
+ reconstruction_root=self.reconstruction_root,
+ metric_name="novel_view_fidelity",
+ )
result = {}
- for dim in dims:
- result_save_path = os.path.join(self.generated_data_path, self.method_name, 'novel_view_consistency', f'{dim}_image_quality.json')
- os.makedirs(os.path.dirname(result_save_path), exist_ok=True)
- video_list = common.find_videos_in_dir(os.path.join(reconstruction_video_folder, dim), extension=".mp4")
+ for dim, video_list in groups.items():
+ result_save_path = result_dir / f"{dim}_image_quality.json"
+ result_save_path.parent.mkdir(parents=True, exist_ok=True)
average_score, video_results = self.calculate_iq_metrics(video_list, **kwds)
result[dim] = {
'average_score': average_score,
'video_results': video_results
}
logger.info(f"Method: {self.method_name}, Dim: {dim}, Novel View Fidelity by MUSIQ: {average_score:.4f}")
- json.dump(result[dim], open(result_save_path, 'w'), indent=4)
\ No newline at end of file
+ with open(result_save_path, 'w') as f:
+ json.dump(result[dim], f, indent=4)
+ return result
diff --git a/worldbench/videogen/reconstruction/photometric_error/__init__.py b/worldbench/videogen/reconstruction/photometric_error/__init__.py
new file mode 100644
index 0000000..221bf0c
--- /dev/null
+++ b/worldbench/videogen/reconstruction/photometric_error/__init__.py
@@ -0,0 +1 @@
+from .photometric_error import PHOTOMETRIC_ERROR
diff --git a/worldbench/videogen/reconstruction/photometric_error/photometric_error.py b/worldbench/videogen/reconstruction/photometric_error/photometric_error.py
new file mode 100644
index 0000000..8929a6d
--- /dev/null
+++ b/worldbench/videogen/reconstruction/photometric_error/photometric_error.py
@@ -0,0 +1,5 @@
+from worldbench.videogen.reconstruction._asset_metrics import AssetPhotometricMetric
+
+
+class PHOTOMETRIC_ERROR(AssetPhotometricMetric):
+ pass
diff --git a/worldbench/videogen/reconstruction/prepare_assets.py b/worldbench/videogen/reconstruction/prepare_assets.py
new file mode 100644
index 0000000..3408940
--- /dev/null
+++ b/worldbench/videogen/reconstruction/prepare_assets.py
@@ -0,0 +1,108 @@
+#!/usr/bin/env python3
+"""Prepare WorldLens Reconstruction assets with DriveStudio."""
+
+from __future__ import annotations
+
+import argparse
+import os
+import subprocess
+import sys
+from pathlib import Path
+from typing import List, Mapping
+
+
+OPTION_FLAGS = {
+ "work_root": "--work-root",
+ "data_root": "--data-root",
+ "target_dir": "--target-dir",
+ "config_file": "--config-file",
+ "step": "--step",
+ "fps": "--fps",
+ "gt_method": "--gt-method",
+ "cameras": "--cameras",
+}
+
+PATH_OPTIONS = {"work_root", "data_root", "target_dir"}
+
+BOOLEAN_FLAGS = {
+ "skip_preprocess": "--skip-preprocess",
+ "skip_train": "--skip-train",
+ "skip_eval": "--skip-eval",
+ "dry_run": "--dry-run",
+}
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser("Prepare WorldLens Reconstruction assets with DriveStudio")
+ parser.add_argument("--method-name", required=True)
+ parser.add_argument("--clips", nargs="+", required=True)
+ parser.add_argument("--reconstruction-root", default="generated_results/reconstruction")
+ parser.add_argument(
+ "--drivestudio-root",
+ default=os.environ.get("DRIVESTUDIO_ROOT"),
+ help="DriveStudio source root. Defaults to worldbench/third_party/drivestudio.",
+ )
+ parser.add_argument(
+ "--python-bin",
+ default=os.environ.get("DRIVESTUDIO_PYTHON", sys.executable),
+ help="Python executable used to run the vendored DriveStudio backend. Defaults to the current WorldLens Python.",
+ )
+ parser.add_argument("--work-root", default=None)
+ parser.add_argument("--data-root", default=None)
+ parser.add_argument("--target-dir", default=None)
+ parser.add_argument("--config-file", default=None)
+ parser.add_argument("--step", default=None)
+ parser.add_argument("--fps", default=None)
+ parser.add_argument("--gt-method", default=None)
+ parser.add_argument("--cameras", default=None)
+ parser.add_argument("--skip-preprocess", action="store_true")
+ parser.add_argument("--skip-train", action="store_true")
+ parser.add_argument("--skip-eval", action="store_true")
+ parser.add_argument("--dry-run", action="store_true")
+ return parser.parse_args()
+
+
+def append_options(command: List[str], args: argparse.Namespace, option_flags: Mapping[str, str]) -> None:
+ for name, flag in option_flags.items():
+ value = getattr(args, name)
+ if value is not None:
+ if name in PATH_OPTIONS:
+ value = Path(value).resolve()
+ command.extend([flag, str(value)])
+
+
+def append_booleans(command: List[str], args: argparse.Namespace, boolean_flags: Mapping[str, str]) -> None:
+ for name, flag in boolean_flags.items():
+ if getattr(args, name):
+ command.append(flag)
+
+
+def main() -> None:
+ args = parse_args()
+ worldbench_root = Path(__file__).resolve().parents[2]
+ drivestudio_root = (
+ Path(args.drivestudio_root).resolve()
+ if args.drivestudio_root
+ else worldbench_root / "third_party" / "drivestudio"
+ )
+ backend_script = drivestudio_root / "tools" / "worldlens_reconstruction_backend.py"
+ if not backend_script.is_file():
+ raise SystemExit(f"Missing DriveStudio WorldLens backend script: {backend_script}")
+
+ command = [
+ str(args.python_bin),
+ str(backend_script.relative_to(drivestudio_root)),
+ "--method-name",
+ args.method_name,
+ "--clips",
+ *[str(clip) for clip in args.clips],
+ "--reconstruction-root",
+ str(Path(args.reconstruction_root).resolve()),
+ ]
+ append_options(command, args, OPTION_FLAGS)
+ append_booleans(command, args, BOOLEAN_FLAGS)
+ subprocess.run(command, cwd=str(drivestudio_root), check=True)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/worldbench/videogen/reconstruction/requirements-reconstruction-extra.txt b/worldbench/videogen/reconstruction/requirements-reconstruction-extra.txt
new file mode 100644
index 0000000..d65f595
--- /dev/null
+++ b/worldbench/videogen/reconstruction/requirements-reconstruction-extra.txt
@@ -0,0 +1,12 @@
+# Install after the base WorldLens environment.
+# Do not install the vendored DriveStudio requirements.txt directly here: it pins
+# torch/torchvision/CUDA versions that can conflict with WorldLens.
+gsplat==1.5.3
+pytorch-msssim==1.0.0
+kornia==0.7.2
+viser==0.2.1
+nerfview==0.0.3
+nvdiffrast @ git+https://github.com/NVlabs/nvdiffrast.git
+
+# PyTorch3D wheels are tied to the local torch/CUDA build. Install a matching
+# wheel separately if your WorldLens environment does not already provide it.
+ 
