VoMP: Predicting Volumetric Mechanical Properties
+Feed-forward, fine-grained, physically based volumetric material properties from Splats, Meshes, NeRFs, and more.
+
+
+
+
+
+
+
+
+
+
+ Feed-forward, fine-grained, physically based volumetric material properties from Splats, Meshes, NeRFs, and more.
+
+
+
+
+
+
+
+
+. Download the models and place them in the `weights/` directory. The checkpoints we will use are the `weights/matvae.safetensors` and `weights/geometry_transformer.pt` files.
+
+The above two files from the model repository contains the final checkpoint of the model. If you need to fine-tune the model, you can follow the same process but download the `ft` directory from the HuggingFace repo too and place them in the `weights/` directory.
+
+| **File** | **Model** |
+|------|----------------|
+| `matvae.safetensors` | MatVAE |
+| `geometry_transformer.pt` | Geometry Transformer |
+| `normalization_params.json` | Normalization Parameters |
+| `inference.json` | Inference Configuration |
+| `ft` | MatVAE and Geometry Transformer checkpoints (same as above but in a directory structure compatible for fine-tuning) |
+
+## ๐ Quickstart: Web Demo
+
+We provide a simple web demo to quickly test out VoMP in a GUI. The web-demo uses some additional dependecies over the base environment, see [`gradio/requirements.txt`](./gradio/requirements.txt). To start the web demo, run:
+
+```bash
+python gradio/app.py
+```
+
+Then, you can access the demo at the address shown in the terminal. The web demo allows you to run the model, visualize the outputs of the model and download an artifact which can directly be used for [๐ฆ Simulation](#-simulation).
+
+## ๐ฅ Loading the Model
+
+Before using any of the VoMP APIs, you need to load the model. We provide flexible loading options:
+
+### Using Inference Config (Recommended)
+
+The simplest way to load the model is using the inference configuration file:
+
+```python
+from vomp.inference import Vomp
+
+# Load model using inference config (recommended - uses final_ckpt.zip weights)
+model = Vomp.from_checkpoint(
+ config_path="weights/inference.json",
+ use_trt=False # Set to True to enable TensorRT acceleration (significantly faster but requires `torch-tensorrt`)
+)
+```
+
+> [!NOTE]
+> Using the `use_trt` flag will compile the DINO model with TensorRT. This makes the `from_checkpoint` function slower.
+
+### Using Direct File Paths
+
+For more control, you can specify exact checkpoint files, optionally overriding the inference config:
+
+```python
+# Load model using direct file paths
+model = Vomp.from_checkpoint(
+ config_path="weights/inference.json",
+ geometry_checkpoint_dir="weights/geometry_transformer.pt",
+ matvae_checkpoint_dir="weights/matvae.safetensors",
+ normalization_params_path="weights/normalization_params.json"
+)
+
+# Or override specific paths from inference config
+model = Vomp.from_checkpoint(
+ config_path="configs/materials/inference.json",
+ geometry_checkpoint_dir="custom/path/to/geometry_transformer.pt" # Override just this path
+)
+```
+
+### Using Directories (use for fine-tuning)
+
+Use this approach only if you are using the fine-tuning checkpoints i.e. the `ft/` directory in the model repository. This lets the model auto-find the latest checkpoints:
+
+```python
+from vomp.inference import Vomp
+
+# Load model using directories (auto-finds latest checkpoints)
+model = Vomp.from_checkpoint(
+ config_path="weights/inference.json",
+ geometry_checkpoint_dir="weights/ft/geometry_transformer",
+ matvae_checkpoint_dir="weights/ft/matvae",
+ normalization_params_path="weights/ft/matvae/normalization_params.json",
+ geometry_ckpt="latest" # Can also be a specific step number
+)
+```
+
+We provide a flexible Python API with both high-level and low-level interfaces for material property estimation.
+
+## ๐ฏ High-Level API
+
+### Gaussian Splats
+
+For Gaussian splats, use the high-level API for the easiest experience (see [Loading the Model](#-loading-the-model) section first):
+
+```python
+from vomp.inference.utils import save_materials
+
+# Get materials directly from PLY (auto-handles Gaussian loading)
+# By default, returns materials evaluated at each Gaussian splat center
+results = model.get_splat_materials("path/to/your/gaussian_splat.ply")
+
+# Or use Kaolin voxelizer for more accurate results
+# results = model.get_splat_materials("path/to/your/gaussian_splat.ply", voxel_method="kaolin")
+
+# Control where materials are evaluated using query_points:
+# results = model.get_splat_materials("path/to/your/gaussian_splat.ply", query_points="splat_centers") # Default
+# results = model.get_splat_materials("path/to/your/gaussian_splat.ply", query_points="voxel_centers") # Voxel centers (direct output of the model)
+# results = model.get_splat_materials("path/to/your/gaussian_splat.ply", query_points=custom_points) # Custom (N,3) array
+
+# Adjust DINO batch size for performance (higher values use more GPU memory)
+# results = model.get_splat_materials("path/to/your/gaussian_splat.ply", dino_batch_size=32)
+
+# Save results
+save_materials(results, "materials.npz")
+```
+
+### Meshes
+
+For mesh objects, use the equivalent high-level mesh API (see [Loading the Model](#-loading-the-model) section first):
+
+```python
+from vomp.inference.utils import save_materials
+
+# Get materials directly from mesh file (supports OBJ, PLY, STL, USD)
+# By default, returns materials evaluated at each mesh vertex (not recommended if you have vertices only on the surface)
+# Note: Requires Blender installation (see Dependencies section)
+results = model.get_mesh_materials(
+ "path/to/your/mesh.obj",
+ blender_path="/tmp/blender-3.0.1-linux-x64/blender" # Adjust path as needed
+)
+
+# Control where materials are evaluated using query_points:
+# results = model.get_mesh_materials("path/to/your/mesh.obj", query_points="mesh_vertices") # Default
+# results = model.get_mesh_materials("path/to/your/mesh.obj", query_points="voxel_centers") # Voxel centers (direct output of the model)
+# results = model.get_mesh_materials("path/to/your/mesh.obj", query_points=custom_points) # Custom (N,3) array
+
+# Use parallel rendering and adjust DINO batch size for better performance
+# results = model.get_mesh_materials("path/to/your/mesh.obj", num_render_jobs=4, dino_batch_size=32, blender_path="/path/to/blender")
+
+# Save results
+save_materials(results, "materials.npz")
+```
+
+### USD Assets (including meshes)
+
+USD files can come in many different formats with varying internal structures, materials, and organization. For USD assets, use the high-level USD API (see [Loading the Model](#-loading-the-model) section first):
+
+#### General USD Formats
+
+For USD files in any format, use [Isaac Sim Replicator](https://docs.isaacsim.omniverse.nvidia.com/5.1.0/replicator_tutorials/index.html) rendering with a separate mesh file for voxelization:
+
+```python
+from vomp.inference.utils import save_materials
+
+# For general USD files - requires Isaac Sim and separate mesh
+# Note: Requires Isaac Sim installation and a separate mesh file for voxelization
+# Isaac Sim renders the USD while the mesh is used for voxelization
+results = model.get_usd_materials(
+ usd_path="path/to/your/model.usd",
+ mesh_path="path/to/your/model.ply", # Mesh for voxelization (doesn't need to be normalized)
+ isaac_sim_path="~/isaac-sim/isaac-sim.sh", # Or set ISAAC_SIM_PATH environment variable
+ render_mode="path_tracing" # Options: "fast" or "path_tracing"
+)
+
+# Control where materials are evaluated using query_points:
+# results = model.get_usd_materials(..., query_points="voxel_centers") # Default (direct output)
+# results = model.get_usd_materials(..., query_points=custom_points) # Custom (N,3) array
+
+# Adjust DINO batch size for performance (higher values use more GPU memory):
+# results = model.get_usd_materials(..., dino_batch_size=32)
+
+# Save results
+save_materials(results, "materials.npz")
+```
+
+Isaac Sim Replicator provides flexible rendering modes:
+
+```python
+# Option 1: Fast Mode - Real-time ray tracing
+results = model.get_usd_materials(
+ usd_path="model.usd",
+ mesh_path="model.ply",
+ isaac_sim_path="~/isaac-sim/isaac-sim.sh",
+ render_mode="fast" # Real-time ray tracing
+)
+
+# Option 2: Path Tracing - Higher quality
+results = model.get_usd_materials(
+ usd_path="model.usd",
+ mesh_path="model.ply",
+ isaac_sim_path="~/isaac-sim/isaac-sim.sh",
+ render_mode="path_tracing" # 256 spp, 8 bounces, denoising enabled
+)
+
+# Option 3: start from a setting and override some RTX settings
+from vomp.inference import RTX_PRESETS
+print(RTX_PRESETS.keys()) # See available presets: ['fast', 'path_tracing']
+
+results = model.get_usd_materials(
+ usd_path="model.usd",
+ mesh_path="model.ply",
+ isaac_sim_path="~/isaac-sim/isaac-sim.sh",
+ render_mode="path_tracing",
+ rtx_settings_override={
+ # Enable path tracing renderer
+ "/rtx/rendermode": "PathTracing",
+
+ # Path tracing quality settings
+ "/rtx/pathtracing/spp": 256, # Samples per pixel (higher = better quality, slower)
+ "/rtx/pathtracing/totalSpp": 256, # Total samples per pixel
+ "/rtx/pathtracing/maxBounces": 8, # Maximum light bounces
+ "/rtx/pathtracing/maxSpecularAndTransmissionBounces": 8,
+
+ # Additional quality settings
+ "/rtx/pathtracing/fireflyFilter/enable": True, # Reduce fireflies (bright pixels)
+ "/rtx/pathtracing/optixDenoiser/enabled": True, # Enable denoiser for clean renders
+
+ # ... other RTX settings you want to override
+ }
+)
+```
+
+> [!WARNING]
+> Please do not override the following RTX settings, as they are required for the model to work correctly:
+> - "/rtx/post/backgroundZeroAlpha/enabled": True,
+> - "/rtx/post/backgroundZeroAlpha/backgroundComposite": False,
+> - "/rtx/post/backgroundZeroAlpha/outputAlphaInComposite": True,
+
+#### SimReady Format USD
+
+If your USD file is in the **SimReady format** (like the USD files in our dataset), you can use the following arguments:
+
+```python
+from vomp.inference.utils import save_materials
+
+results = model.get_usd_materials(
+ usd_path="model.usd",
+ use_simready_usd_format=True,
+ blender_path="/path/to/blender",
+ seed=42
+)
+
+# Control where materials are evaluated using query_points:
+# results = model.get_usd_materials(..., query_points="voxel_centers") # Default (direct output)
+# results = model.get_usd_materials(..., query_points=custom_points) # Custom (N,3) array
+
+# Save results
+save_materials(results, "materials.npz")
+```
+
+## ๐จ Visualizing Material Results
+
+After estimating material properties, you can visualize them using our interactive `polyscope`-based viewer.
+
+```python
+# After getting results from any API
+from vomp.inference.utils import save_materials
+
+# Save your results
+save_materials(results, "my_materials.npz")
+```
+
+```bash
+# Launch interactive property viewer
+python scripts/viewer.py my_materials.npz
+```
+
+The viewer also saves the colorbars for visualizations as PNG images that look like this:
+
+
+
+## ๐ง Low-Level API
+
+### Gaussian Splats
+
+For fine-grained control with Gaussian splats, use the low-level API (see [Loading the Model](#-loading-the-model) section first):
+
+```python
+from vomp.representations.gaussian import Gaussian
+from vomp.inference.utils import save_materials
+
+# Load Gaussian splat
+gaussian = Gaussian(sh_degree=3, aabb=[-1,-1,-1,2,2,2], device="cuda")
+gaussian.load_ply("path/to/your/gaussian_splat.ply")
+
+# Step-by-step pipeline
+output_dir = "outputs/materials"
+renders_metadata = model.render_sampled_views(gaussian, output_dir, num_views=150)
+voxel_centers = model._voxelize_gaussian(gaussian, output_dir)
+coords, features = model._extract_dino_features(output_dir, voxel_centers, renders_metadata, save_features=True)
+results = model.predict_materials(coords, features)
+save_materials(results, "materials.npz")
+```
+
+### Meshes
+
+For fine-grained control with meshes, use the equivalent low-level mesh API (see [Loading the Model](#-loading-the-model) section first):
+
+```python
+from vomp.inference.utils import save_materials
+
+# Step-by-step pipeline for meshes
+output_dir = "outputs/materials"
+mesh_path = "path/to/your/mesh.obj"
+blender_path = "/tmp/blender-3.0.1-linux-x64/blender" # Adjust for your installation
+renders_metadata = model.render_mesh_views(mesh_path, output_dir, num_views=150, blender_path=blender_path)
+voxel_centers = model._voxelize_mesh(mesh_path, output_dir)
+coords, features = model._extract_dino_features(output_dir, voxel_centers, renders_metadata, save_features=True)
+results = model.predict_materials(coords, features)
+save_materials(results, "materials.npz")
+```
+
+### USD Assets
+
+For fine-grained control with USD assets using Replicator rendering (see [Loading the Model](#-loading-the-model) section first):
+
+```python
+from vomp.inference.utils import save_materials
+
+# Step-by-step pipeline for USD assets with Replicator
+output_dir = "outputs/materials"
+usd_path = "path/to/your/model.usd"
+mesh_path = "path/to/your/model.ply" # For voxelization
+isaac_sim_path = "~/isaac-sim/isaac-sim.sh"
+
+# Render using Replicator (with custom settings)
+renders_metadata = model.render_views_replicator(
+ asset_path=usd_path,
+ output_dir=output_dir,
+ num_views=150,
+ isaac_sim_path=isaac_sim_path,
+ render_mode="path_tracing", # or "fast"
+ rtx_settings_override={
+ "/rtx/pathtracing/spp": 512 # Optional: custom settings
+ }
+)
+
+# Voxelize and extract features
+voxel_centers = model._voxelize_mesh(mesh_path, output_dir)
+coords, features = model._extract_dino_features(output_dir, voxel_centers, renders_metadata, save_features=True)
+results = model.predict_materials(coords, features)
+save_materials(results, "materials.npz")
+```
+
+## ๐งฉ Custom 3D Representations
+
+Bring your own 3D representation with custom render/voxelize functions (see [Loading the Model](#-loading-the-model) section first):
+
+```python
+from vomp.inference.utils import save_materials
+
+def my_render_func(obj, output_dir, num_views, image_size, **kwargs):
+ # Your rendering code here - save images to output_dir/renders/
+ frames_metadata = []
+ for i in range(num_views):
+ # Your custom rendering logic
+ frames_metadata.append({
+ "file_path": f"frame_{i:04d}.png",
+ "transform_matrix": camera_matrix.tolist(), # 4x4 matrix
+ "camera_angle_x": fov_radians
+ })
+ return frames_metadata
+
+def my_voxelize_func(obj, output_dir, **kwargs):
+ # Your voxelization code here
+ voxel_centers = your_voxelization_method(obj) # (N, 3) array
+ return voxel_centers
+
+# Use with any 3D representation
+coords, features = model.get_features(
+ obj_3d=your_mesh,
+ render_func=my_render_func,
+ voxelize_func=my_voxelize_func,
+ num_views=150
+)
+
+# Get materials
+results = model.predict_materials(coords, features)
+save_materials(results, "materials.npz")
+```
+
+## ๐งฌ Material Upsampler
+
+The high-level splat API (`get_splat_materials`) automatically returns materials interpolated to splat centers. However, you may want to upsample materials to other locations like higher resolution grids or custom query points. We provide a utility class for these cases (see [Loading the Model](#-loading-the-model) section first).
+
+```python
+import numpy as np
+from vomp.inference.utils import MaterialUpsampler
+from vomp.representations.gaussian import Gaussian
+
+# Get voxel-level materials (needed for upsampling to custom locations)
+# Note: Use query_points="voxel_centers" to get voxel-level results
+voxel_results = model.get_splat_materials("path/to/your/gaussian_splat.ply", query_points="voxel_centers")
+# OR for meshes
+# voxel_results = model.get_mesh_materials("path/to/your/mesh.obj", query_points="voxel_centers", blender_path="/path/to/blender")
+
+# Create upsampler from voxel-level results
+upsampler = MaterialUpsampler(
+ voxel_coords=voxel_results["voxel_coords_world"],
+ voxel_materials=np.column_stack([
+ voxel_results["youngs_modulus"],
+ voxel_results["poisson_ratio"],
+ voxel_results["density"]
+ ])
+)
+
+# Example 1: Interpolate to Gaussian centers manually (usually not needed - high-level API does this)
+gaussian = Gaussian(sh_degree=3, aabb=[-1,-1,-1,2,2,2], device="cuda")
+gaussian.load_ply("path/to/your/gaussian_splat.ply")
+gaussian_materials, gaussian_distances = upsampler.interpolate_to_gaussians(gaussian)
+
+# Example 2: Interpolate to higher resolution grid (128x128x128) - main use case for manual upsampling
+x = np.linspace(-0.5, 0.5, 128)
+xx, yy, zz = np.meshgrid(x, x, x)
+high_res_points = np.column_stack([xx.ravel(), yy.ravel(), zz.ravel()])
+high_res_materials, high_res_distances = upsampler.interpolate(high_res_points)
+
+# Example 3: Interpolate to custom query points - another main use case for manual upsampling
+query_points = np.random.uniform(-0.4, 0.4, size=(1000, 3))
+query_materials, query_distances = upsampler.interpolate(query_points)
+
+# Save results
+upsampler.save_results(gaussian.get_xyz.detach().cpu().numpy(), gaussian_materials,
+ gaussian_distances, "gaussian_materials.npz")
+upsampler.save_results(high_res_points, high_res_materials, high_res_distances, "high_res_materials.npz")
+upsampler.save_results(query_points, query_materials, query_distances, "custom_materials.npz")
+```
+
+## ๐พ Using our Benchmark
+
+> [!NOTE]
+> Due to licenses we are unable to make the vegetation subset of the dataset public. Thus, when you compare outputs to the paper make sure to compare them to the listed results on the "public dataset" (Table 2 and Table 3).
+
+We provide a dataset and a benchmark with fine-grained volumetric mechanical properties (65.9 GB) at 
(or preprocess it yourself using the instructions in [README_train.md](./README_train.md)). We also provide code allowing the evaluation of models on this dataset.
+
+### Reproducing results from the paper
+
+Since our dataset is quite large, we provide a way to download only the test set by running the following command:
+
+```bash
+huggingface-cli download nvidia/PhysicalAI-Robotics-PhysicalAssets-VoMP-Eval --repo-type dataset --local-dir datasets/simready
+```
+
+We can now run VoMP on the test set:
+
+```bash
+python scripts/evaluate_geometry_encoder.py \
+ --config weights/inference.json \ # replace with your own config file
+ --checkpoint_dir weights/ft/geometry_transformer \ # replace with your own checkpoint directory
+ --data_dir datasets/simready \ # replace with your own data directory
+ # --ckpt latest \
+ # --results
+```
+
+This script requires loading the model in the ["Using Directories" method](#using-directories).
+
+This prints out many detailed metrics. Particularly, you can also make sure you can reproduce the main results from the paper by comparing Table 2 and Appendix Table 3 from the paper with the outputs from Section 5 (SUMMARY TABLES) of the results printed out.
+
+To build on top of our benchmark, you can replace the `load_model` and `evaluate_model` functions in the `scripts/evaluate_geometry_encoder.py` script with your own model and evaluation code.
+
+## ๐ฆ Simulation
+
+Our properties are compatible with all simulators. We provide instructions to run a few kinds of simulations with the properties.
+
+### Simplicits simulation
+
+For the large-scale simulations that we perform with [Simplicits](https://research.nvidia.com/labs/toronto-ai/simplicits/), refer to the [Simplicits](https://kaolin.readthedocs.io/en/latest/notes/simplicits.html) documentation.
+
+### FEM simulation using warp.fem
+
+We provide a way to run FEM simulations using [`warp.fem`](https://nvidia.github.io/warp/modules/fem.html).
+
+```bash
+cd simulation/warp.fem
+PYTHONPATH=./ python drop_tetmesh.py --mesh assets/cube_res20.msh --materials assets/cube_materials_two_halves.npz
+```
+
+This simple example has an artificially constructed NPZ file which can be used in `warp.fem`. This requires installing [`warp`](https://nvidia.github.io/warp/) and `meshio`.
+
+### FEM simulation using libuipc
+
+We provide a way to run FEM simulations using [`libuipc`](https://github.com/spiriMirror/libuipc/). These simulations use the config files in the `configs/sim/` directory and they can be run as,
+
+```bash
+python vomp/sim/main.py configs/sim/falling_oranges.json
+```
+
+This config runs a simulation of falling oranges (Figure 5 from the paper) with the NPZ files we generated from the VoMP model.
+
+These simulations require a `.npz` file with the estimated mechanical properties of the object. This requires installing the Python version of `libuipc` using the instructions in the [`libuipc`](https://github.com/spiriMirror/libuipc/) repository. The command above will run the simulation, show it in a GUI, and save framewise surface meshes in the `outputs/simulation_output/falling_oranges` directory. The config also specifies a visual textured surface mesh so the per frame visualizations will use the high resolution visual mesh and also have textures.
+
+### Newton simulation
+
+We provide a way to run [Newton](https://github.com/newton-physics/newton/) simulations. Run an example simulation of a soft body cube with the NPZ files we generated from the VoMP model by running the following command:
+
+```bash
+cd simulation/newton
+python mesh_falling_sim.py --grid_dim 16 --materials cube_high_E.npz
+python mesh_falling_sim.py --grid_dim 16 --materials cube_low_E.npz
+```
+
+This simple example has two artificially constructed NPZ files which can be used in Newton. Observe the difference in simulation showing all Young's modulus, Poisson's ratio, and density values were properly applied. This requires installing [`newton`](https://github.com/newton-physics/newton/) and `meshio`.
+
+> [!NOTE]
+> Our properties are also compatible with [PhysX](https://developer.nvidia.com/physx-sdk) and rigid-body simulators. We plan to release some example code to do so at a later date. Until then, if you want to use our properties in PhysX, we recommend clustering the properties we produce, split the underlying meshes based on the clusters, and then add the averaged property for each such "connected part".
+
+## ๐ค Credits
+
+We are also grateful to several other open-source repositories that we drew inspiration from or built upon during the development of our pipeline:
+
+- [DINOv2](https://github.com/facebookresearch/dinov2)
+- [fTetWild](https://github.com/wildmeshing/fTetWild)
+- [gaussian-splatting](https://github.com/graphdeco-inria/gaussian-splatting)
+- [Isaac Sim](https://developer.nvidia.com/isaac/sim)
+- [kaolin](https://github.com/NVIDIAGameWorks/kaolin)
+- [libuipc](https://github.com/spiriMirror/libuipc)
+- [newton](https://github.com/newton-physics/newton)
+- [Simplicits](https://research.nvidia.com/labs/toronto-ai/simplicits/)
+- [textgrad](https://github.com/zou-group/textgrad)
+- [TRELLIS](https://github.com/microsoft/TRELLIS)
+- [Warp](https://nvidia.github.io/warp/)
+
+## ๐ Citation
+
+If you find VoMP helpful, please consider citing:
+
+```bibtex
+@inproceedings{dagli2026vomp,
+ title={Vo{MP}: Predicting Volumetric Mechanical Property Fields},
+ author={Rishit Dagli and Donglai Xiang and Vismay Modi and Charles Loop and Clement Fuji Tsang and Anka He Chen and Anita Hu and Gavriel State and David Levin I.W. and Maria Shugrina},
+ booktitle={The Fourteenth International Conference on Learning Representations},
+ year={2026},
+url={https://openreview.net/forum?id=aTP1IM6alo}
+}
+```
+
+## License and Contact
+
+This project will download and install additional third-party open source software projects. Review the license terms of these open source projects before use.
+
+VoMP source code is released under the [Apache 2 License](https://www.apache.org/licenses/LICENSE-2.0).
+
+VoMP models are released under the [NVIDIA Open Model License](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license). For a custom license, please visit our website and submit the form: [NVIDIA Research Licensing](https://www.nvidia.com/en-us/research/inquiries/).
\ No newline at end of file
diff --git a/deps/vomp/README_train.md b/deps/vomp/README_train.md
new file mode 100644
index 0000000000000000000000000000000000000000..15e4c81382d6a719962965920ca11f41da858a35
--- /dev/null
+++ b/deps/vomp/README_train.md
@@ -0,0 +1,285 @@
+
+
+
+
+. We are unable to make the MTD dataset public due to licenses.
+
+### Material Triplet Dataset (MTD)
+
+First compile the `material_ranges.csv` file by extracting data from the following sources (and deduplicate the data):
+
+- [MatWeb](https://matweb.com/)
+- [Engineering Toolbox](https://www.engineeringtoolbox.com/engineering-materials-properties-d_1225.html)
+- [Cambridge University Press](https://teaching.eng.cam.ac.uk/sites/teaching.eng.cam.ac.uk/files/Documents/Databooks/MATERIALS%20DATABOOK%20(2011)%20version%20for%20Moodle.pdf)
+
+The Material Triplet Dataset (MTD) is used to train the MatVAE. Assuming you have the `material_ranges.csv` file in the `datasets/latent_space/` directory, you can create the MTD by running the following command:
+
+```bash
+python dataset_toolkits/latent_space/make_csv.py datasets/latent_space/
+```
+
+Due to the dataset licenses, we cannot provide the `material_ranges.csv` file.
+
+### Geometry with Volumetric Materials (GVM)
+
+The Geometry with Volumetric Materials (GVM) is used to train the Geometry Transformer. First, download the following datasets to `datasets/raw/`:
+
+- [SimReady (13.9 GB + 20.5 GB + 9.4 GB + 21.4 GB + 20.6 GB)](https://docs.omniverse.nvidia.com/usd/latest/usd_content_samples/downloadable_packs.html#simready-warehouse-01-assets-pack)
+- [Commercial (5.8 GB)](https://docs.omniverse.nvidia.com/usd/latest/usd_content_samples/downloadable_packs.html#commercial-assets-pack)
+- [Residential (22.5 GB)](https://docs.omniverse.nvidia.com/usd/latest/usd_content_samples/downloadable_packs.html#residential-assets-pack)
+- [Vegetation (2.7 GB)](https://docs.omniverse.nvidia.com/usd/latest/usd_content_samples/downloadable_packs.html#vegetation-assets-pack)
+
+> [!NOTE]
+> The SimReady dataset is split into 5 parts. You can download them all from the aforementioned URL.
+
+Next, unzip these datasets to `datasets/raw/`, to create a directory structure like:
+
+```
+datasets/raw/
+โโโ simready/
+โโโ commercial/
+โโโ residential/
+โโโ vegetation/
+```
+
+Then, run the following command to create the GVM. This step takes ~2.5 days on 2 A100 GPUs, assuming you have enough CPU resources, as we perform significant CPU rendering.
+
+```bash
+mkdir -p /tmp/vlm
+
+python dataset_toolkits/material_objects/vlm_annotations/main.py \
+ --dataset simready residential commercial vegetation \
+ -o datasets/raw/material_annotations.json \
+ --verbose
+```
+
+The VLM prompt is optimized using the `scripts/optimize_prompt.py` script which requires installing [textgrad](https://github.com/zou-group/textgrad).
+
+This saves the annotations to `datasets/raw/material_annotations.json` in the following format.
+
+```json
+[
+ {
+ "object_name": "aluminumpallet_a01",
+ "category": "pallet",
+ "dataset_type": "simready",
+ "segments": {
+ "SM_AluminumPallet_A01_01": {
+ "name": "default__metal__aluminumpallet_a01",
+ "opacity": "opaque",
+ "material_type": "metal",
+ "semantic_usage": "aluminumpallet_a01",
+ "density": 2700.0,
+ "dynamic_friction": 0.1,
+ "static_friction": 0.1,
+ "restitution": 0.1,
+ "textures": {
+ "albedo": "datasets/raw/simready/common_assets/props/aluminumpallet_a01/textures/T_Aluminium_Brushed_A1_Albedo.png",
+ "orm": "datasets/raw/simready/common_assets/props/aluminumpallet_a01/textures/T_Aluminium_Brushed_A1_ORM.png",
+ "normal": "datasets/raw/simready/common_assets/props/aluminumpallet_a01/textures/T_Aluminium_Brushed_A1_Normal.png"
+ },
+ "vlm_analysis": "...",
+ "youngs_modulus": 70000000000.0,
+ "poissons_ratio": 0.33
+ }
+ },
+ "file_path": "datasets/raw/simready/common_assets/props/aluminumpallet_a01/aluminumpallet_a01_inst_base.usd"
+ },
+ ...
+]
+```
+
+### Preparing your own data for training the Geometry Transformer
+
+To train VoMP on your own data, you need to prepare a dataset of 3D objects with volumetric materials. Particularly, you need to prepare a JSON file and USD files with the following format:
+
+```json
+[
+ {
+ "object_name": "[object name]",
+ "segments": {
+ "[segment name that matches the segment name in the USD file]": {
+ "density": 2700.0,
+ "youngs_modulus": 70000000000.0,
+ "poissons_ratio": 0.33
+ }
+ },
+ "file_path": "path/to/your/object.usd"
+ }
+ ...
+]
+```
+
+If you are preparing your own dataset make sure the individual segments you list in the JSON file match the segment names in the USD file and each segment is a mesh. Also make sure the object has appearance properties. The workflow would work even if you do not have appearance properties, but the estimated properties would be significantly worse.
+
+## ๐ป Training
+
+### Training the MatVAE
+
+First run `accelerate` config to create a config file, setting your hardware details and if you want to do distributed training. We highly recommend using a single GPU for training MatVAE. This step takes ~12 hours on a single A100 GPU.
+
+Training hyperparameters and model architectures are defined in configuration files under the `configs/` directory. Example configuration files include:
+
+| **Config** | **Description** |
+|------------|-----------------|
+| `configs/materials/material_vae/matvae.json` | Training configuration for MatVAE. |
+| ... | Training configuration for ablations. |
+
+Any configuration file can be used to start training (use `accelerate launch` instead of `python` if you want to do distributed training),
+
+```bash
+python train_material_vae.py --config ...
+```
+
+Train the MatVAE by running the following command:
+
+```bash
+python train_material_vae.py --config configs/materials/material_vae/matvae.json
+```
+
+This creates the `outputs/matvae/` directory, which contains the trained model and tensorboard logs.
+
+### Training the Geometry Transformer
+
+First, start by performing data preprocessing. This step takes ~2 days on an A100 GPU + ~1.5 days on an RTX6000 GPU (used for rendering).
+
+```bash
+# python dataset_toolkits/build_metadata.py simready --output_dir datasets/simready
+python dataset_toolkits/build_metadata.py allmats --output_dir datasets/simready
+
+# Render USD files to images (can be parallelized across GPUs)
+# For multi-GPU: use --rank and --world_size arguments
+# Example: python ... --rank 0 --world_size 4 (run on GPU 0)
+# python ... --rank 1 --world_size 4 (run on GPU 1), etc.
+python dataset_toolkits/material_objects/render_usd.py allmats --output_dir datasets/simready --quiet --max_workers 3
+
+python dataset_toolkits/build_metadata.py allmats --output_dir datasets/simready --from_file
+python dataset_toolkits/material_objects/voxelize.py --output_dir datasets/simready --max_voxels 72000 --force
+python dataset_toolkits/build_metadata.py allmats --output_dir datasets/simready --from_file
+
+python dataset_toolkits/extract_feature.py --output_dir datasets/simready --force
+python dataset_toolkits/build_metadata.py allmats --output_dir datasets/simready
+```
+
+This creates the `datasets/simready/` directory, which contains the preprocessed data.
+
+```bash
+datasets/simready
+โโโ features (outputs from DINOv2 feature aggregation)
+โโโ merged_records
+โโโ metadata.csv
+โโโ renders (150 rendered images per object with camera poses)
+โโโ splits (train/val/test splits)
+โโโ statistics.txt (statistics of the dataset)
+โโโ voxels (voxelized meshes and voxel-wise mechanical properties)
+```
+
+Next, run the following command to train the Geometry Transformer. This step takes ~5 days on 4 A100 GPUs.
+
+```bash
+python train_geometry_encoder.py --config configs/materials/geometry_encoder/train.json --output_dir outputs/geometry_encoder
+```
+
+This creates the `outputs/geometry_encoder/` directory, which contains the trained model and tensorboard logs.
+
+### Training on your own data
+
+Once you have prepared your dataset following the format above, training is straightforward.
+
+```bash
+python train_geometry_encoder.py --config ... --output_dir ...
+```
+
+Replace the config and output directory with your own. You can make a new config file by copying one of the existing ones in the `configs/` directory and modifying the hyperparameters and dataset paths.
+
+### Fine-tuning
+
+Fine-tuning from pre-trained checkpoints is built into the training pipeline, simply run the following command:
+
+```bash
+python train_geometry_encoder.py --config ... --output_dir ...
+```
+
+It searches for models in the `outputs/geometry_encoder/ckpts/` directory in the following format `geometry_encoder_step[0-9]+.pt` and uses it to continue training.
+
+```bash
+โโโ geometry_encoder_ema0.9999_step0060000.pt
+โโโ geometry_encoder_ema0.9999_step0200000.pt
+โโโ geometry_encoder_step0060000.pt
+โโโ geometry_encoder_step0200000.pt
+โโโ misc_step0060000.pt
+โโโ misc_step0200000.pt
+```
+
+It also optionally searches for the `misc_step[0-9]+.pt` file to restore the optimizer state and scheduler state as well as `geometry_encoder_ema0.9999_step[0-9]+.pt` to restore the EMA model weights.
+
+## ๐ก Tips
+
+- Running the model requires 40 GB VRAM. If you often run into out of memory errors, you can reduce the amount of voxels we use for the object.
+- Dataset annotation with a VLM uses Qwen2.5-VL-72B which requires ~138 GB VRAM even when you load it in BF16 precision. The dataset annotation was done on 2 A100 GPUs. If you often run into out of memory errors, you can swap for a smaller version of Qwen2.5-VL or some other model, though the annotation would likely be degraded.
\ No newline at end of file
diff --git a/deps/vomp/configs/materials/geometry_encoder/train.json b/deps/vomp/configs/materials/geometry_encoder/train.json
new file mode 100644
index 0000000000000000000000000000000000000000..9b0b1433b39e3c7f68628aecb8e2f023387b637c
--- /dev/null
+++ b/deps/vomp/configs/materials/geometry_encoder/train.json
@@ -0,0 +1,83 @@
+{
+ "matvae_checkpoint": "outputs/matvae2/checkpoints/checkpoint_853/model.safetensors",
+ "trellis_weights_path": "weights/TRELLIS-image-large",
+ "models": {
+ "geometry_encoder": {
+ "name": "geometry_encoder",
+ "args": {
+ "resolution": 64,
+ "in_channels": 1024,
+ "model_channels": 768,
+ "latent_channels": 2,
+ "num_blocks": 12,
+ "num_heads": 12,
+ "mlp_ratio": 4,
+ "attn_mode": "swin",
+ "window_size": 8,
+ "use_fp16": true
+ }
+ },
+ "matvae": {
+ "name": "matvae",
+ "args": {
+ "width": 256,
+ "depth": 3,
+ "z_dim": 2,
+ "p_drop": 0.05,
+ "use_learned_variances": false,
+ "use_additional_losses": true
+ }
+ }
+ },
+ "dataset": {
+ "name": "SparseVoxelMaterials",
+ "normalization_type": "log_minmax",
+ "args": {
+ "roots": "datasets/simready",
+ "image_size": 512,
+ "model": "dinov2_vitl14_reg",
+ "resolution": 64,
+ "min_aesthetic_score": 0.0,
+ "max_num_voxels": 32768,
+ "compute_material_stats": false
+ }
+ },
+ "trainer": {
+ "name": "SLatVaeMaterialsTrainer",
+ "args": {
+ "max_steps": 1000000,
+ "batch_size_per_gpu": 16,
+ "batch_split": 1,
+ "optimizer": {
+ "name": "AdamW",
+ "args": {
+ "lr": 1e-4,
+ "weight_decay": 0.0
+ }
+ },
+ "ema_rate": [
+ 0.9999
+ ],
+ "fp16_mode": "inflat_all",
+ "fp16_scale_growth": 0.001,
+ "elastic": {
+ "name": "LinearMemoryController",
+ "args": {
+ "target_ratio": 0.75,
+ "max_mem_ratio_start": 0.5
+ }
+ },
+ "grad_clip": {
+ "name": "AdaptiveGradClipper",
+ "args": {
+ "max_norm": 1.0,
+ "clip_percentile": 95
+ }
+ },
+ "i_log": 10,
+ "i_save": 2000,
+ "i_eval": 1000,
+ "loss_type": "l1"
+ }
+ }
+}
\ No newline at end of file
diff --git a/deps/vomp/configs/materials/geometry_encoder/train_encoder_decoder_direct.json b/deps/vomp/configs/materials/geometry_encoder/train_encoder_decoder_direct.json
new file mode 100644
index 0000000000000000000000000000000000000000..3f8d2d4684f7f89d941aebf64d1ff5747b5a5932
--- /dev/null
+++ b/deps/vomp/configs/materials/geometry_encoder/train_encoder_decoder_direct.json
@@ -0,0 +1,99 @@
+{
+ "training_mode": "encoder_decoder_direct",
+ "matvae_checkpoint": "outputs/matvae/checkpoints/checkpoint_821/model.safetensors",
+ "trellis_weights_path": "weights/TRELLIS-image-large",
+ "models": {
+ "geometry_encoder": {
+ "name": "geometry_encoder",
+ "args": {
+ "resolution": 64,
+ "in_channels": 1024,
+ "model_channels": 768,
+ "latent_channels": 8,
+ "num_blocks": 12,
+ "num_heads": 12,
+ "mlp_ratio": 4,
+ "attn_mode": "swin",
+ "window_size": 8,
+ "use_fp16": true
+ }
+ },
+ "decoder": {
+ "name": "decoder",
+ "args": {
+ "resolution": 64,
+ "model_channels": 768,
+ "latent_channels": 8,
+ "num_blocks": 12,
+ "out_channels": 3,
+ "num_heads": 12,
+ "mlp_ratio": 4,
+ "attn_mode": "swin",
+ "window_size": 8,
+ "use_fp16": true
+ }
+ },
+ "matvae": {
+ "name": "matvae",
+ "args": {
+ "width": 256,
+ "depth": 3,
+ "z_dim": 2,
+ "p_drop": 0.05,
+ "use_learned_variances": false,
+ "use_additional_losses": true
+ }
+ }
+ },
+ "dataset": {
+ "name": "SparseVoxelMaterials",
+ "normalization_type": "log_minmax",
+ "args": {
+ "roots": "datasets/simready",
+ "image_size": 512,
+ "model": "dinov2_vitl14_reg",
+ "resolution": 64,
+ "min_aesthetic_score": 0.0,
+ "max_num_voxels": 32768,
+ "compute_material_stats": false
+ }
+ },
+ "trainer": {
+ "name": "SLatVaeMaterialsTrainer",
+ "args": {
+ "max_steps": 1000000,
+ "batch_size_per_gpu": 16,
+ "batch_split": 1,
+ "optimizer": {
+ "name": "AdamW",
+ "args": {
+ "lr": 1e-4,
+ "weight_decay": 0.0
+ }
+ },
+ "ema_rate": [
+ 0.9999
+ ],
+ "fp16_mode": "inflat_all",
+ "fp16_scale_growth": 0.001,
+ "elastic": {
+ "name": "LinearMemoryController",
+ "args": {
+ "target_ratio": 0.75,
+ "max_mem_ratio_start": 0.5
+ }
+ },
+ "grad_clip": {
+ "name": "AdaptiveGradClipper",
+ "args": {
+ "max_norm": 1.0,
+ "clip_percentile": 95
+ }
+ },
+ "i_log": 10,
+ "i_save": 1000,
+ "i_eval": 1000,
+ "loss_type": "l1"
+ }
+ }
+}
\ No newline at end of file
diff --git a/deps/vomp/configs/materials/geometry_encoder/train_encoder_decoder_matvae.json b/deps/vomp/configs/materials/geometry_encoder/train_encoder_decoder_matvae.json
new file mode 100644
index 0000000000000000000000000000000000000000..d05339b48b8077055b3e0629d88011c9a9e916a3
--- /dev/null
+++ b/deps/vomp/configs/materials/geometry_encoder/train_encoder_decoder_matvae.json
@@ -0,0 +1,99 @@
+{
+ "training_mode": "encoder_decoder_matvae",
+ "matvae_checkpoint": "outputs/matvae/checkpoints/checkpoint_821/model.safetensors",
+ "trellis_weights_path": "weights/TRELLIS-image-large",
+ "models": {
+ "geometry_encoder": {
+ "name": "geometry_encoder",
+ "args": {
+ "resolution": 64,
+ "in_channels": 1024,
+ "model_channels": 768,
+ "latent_channels": 8,
+ "num_blocks": 12,
+ "num_heads": 12,
+ "mlp_ratio": 4,
+ "attn_mode": "swin",
+ "window_size": 8,
+ "use_fp16": true
+ }
+ },
+ "decoder": {
+ "name": "decoder",
+ "args": {
+ "resolution": 64,
+ "model_channels": 768,
+ "latent_channels": 8,
+ "num_blocks": 12,
+ "out_channels": 2,
+ "num_heads": 12,
+ "mlp_ratio": 4,
+ "attn_mode": "swin",
+ "window_size": 8,
+ "use_fp16": true
+ }
+ },
+ "matvae": {
+ "name": "matvae",
+ "args": {
+ "width": 256,
+ "depth": 3,
+ "z_dim": 2,
+ "p_drop": 0.05,
+ "use_learned_variances": false,
+ "use_additional_losses": true
+ }
+ }
+ },
+ "dataset": {
+ "name": "SparseVoxelMaterials",
+ "normalization_type": "log_minmax",
+ "args": {
+ "roots": "datasets/simready",
+ "image_size": 512,
+ "model": "dinov2_vitl14_reg",
+ "resolution": 64,
+ "min_aesthetic_score": 0.0,
+ "max_num_voxels": 32768,
+ "compute_material_stats": false
+ }
+ },
+ "trainer": {
+ "name": "SLatVaeMaterialsTrainer",
+ "args": {
+ "max_steps": 1000000,
+ "batch_size_per_gpu": 16,
+ "batch_split": 1,
+ "optimizer": {
+ "name": "AdamW",
+ "args": {
+ "lr": 1e-4,
+ "weight_decay": 0.0
+ }
+ },
+ "ema_rate": [
+ 0.9999
+ ],
+ "fp16_mode": "inflat_all",
+ "fp16_scale_growth": 0.001,
+ "elastic": {
+ "name": "LinearMemoryController",
+ "args": {
+ "target_ratio": 0.75,
+ "max_mem_ratio_start": 0.5
+ }
+ },
+ "grad_clip": {
+ "name": "AdaptiveGradClipper",
+ "args": {
+ "max_norm": 1.0,
+ "clip_percentile": 95
+ }
+ },
+ "i_log": 10,
+ "i_save": 1000,
+ "i_eval": 1000,
+ "loss_type": "l1"
+ }
+ }
+}
\ No newline at end of file
diff --git a/deps/vomp/configs/materials/geometry_encoder/train_standard.json b/deps/vomp/configs/materials/geometry_encoder/train_standard.json
new file mode 100644
index 0000000000000000000000000000000000000000..dbcbc7fd5b519663c7eb699ce9e657e1eb18796e
--- /dev/null
+++ b/deps/vomp/configs/materials/geometry_encoder/train_standard.json
@@ -0,0 +1,83 @@
+{
+ "matvae_checkpoint": "outputs/matvae2/checkpoints/checkpoint_853/model.safetensors",
+ "trellis_weights_path": "weights/TRELLIS-image-large",
+ "models": {
+ "geometry_encoder": {
+ "name": "geometry_encoder",
+ "args": {
+ "resolution": 64,
+ "in_channels": 1024,
+ "model_channels": 768,
+ "latent_channels": 2,
+ "num_blocks": 12,
+ "num_heads": 12,
+ "mlp_ratio": 4,
+ "attn_mode": "swin",
+ "window_size": 8,
+ "use_fp16": true
+ }
+ },
+ "matvae": {
+ "name": "matvae",
+ "args": {
+ "width": 256,
+ "depth": 3,
+ "z_dim": 2,
+ "p_drop": 0.05,
+ "use_learned_variances": false,
+ "use_additional_losses": true
+ }
+ }
+ },
+ "dataset": {
+ "name": "SparseVoxelMaterials",
+ "normalization_type": "standard",
+ "args": {
+ "roots": "datasets/simready",
+ "image_size": 512,
+ "model": "dinov2_vitl14_reg",
+ "resolution": 64,
+ "min_aesthetic_score": 0.0,
+ "max_num_voxels": 32768,
+ "compute_material_stats": false
+ }
+ },
+ "trainer": {
+ "name": "SLatVaeMaterialsTrainer",
+ "args": {
+ "max_steps": 1000000,
+ "batch_size_per_gpu": 16,
+ "batch_split": 1,
+ "optimizer": {
+ "name": "AdamW",
+ "args": {
+ "lr": 1e-4,
+ "weight_decay": 0.0
+ }
+ },
+ "ema_rate": [
+ 0.9999
+ ],
+ "fp16_mode": "inflat_all",
+ "fp16_scale_growth": 0.001,
+ "elastic": {
+ "name": "LinearMemoryController",
+ "args": {
+ "target_ratio": 0.75,
+ "max_mem_ratio_start": 0.5
+ }
+ },
+ "grad_clip": {
+ "name": "AdaptiveGradClipper",
+ "args": {
+ "max_norm": 1.0,
+ "clip_percentile": 95
+ }
+ },
+ "i_log": 10,
+ "i_save": 2000,
+ "i_eval": 1000,
+ "loss_type": "l1"
+ }
+ }
+}
\ No newline at end of file
diff --git a/deps/vomp/configs/materials/inference.json b/deps/vomp/configs/materials/inference.json
new file mode 100644
index 0000000000000000000000000000000000000000..577b7966500995a5457290dd560bd635d9429b54
--- /dev/null
+++ b/deps/vomp/configs/materials/inference.json
@@ -0,0 +1,86 @@
+{
+ "geometry_checkpoint_dir": "weights/geometry_transformer.pt",
+ "matvae_checkpoint_dir": "weights/matvae.safetensors",
+ "normalization_params_path": "weights/normalization_params.json",
+ "matvae_checkpoint": "weights/matvae.safetensors",
+ "trellis_weights_path": "weights/TRELLIS-image-large",
+ "models": {
+ "geometry_encoder": {
+ "name": "geometry_encoder",
+ "args": {
+ "resolution": 64,
+ "in_channels": 1024,
+ "model_channels": 768,
+ "latent_channels": 2,
+ "num_blocks": 12,
+ "num_heads": 12,
+ "mlp_ratio": 4,
+ "attn_mode": "swin",
+ "window_size": 8,
+ "use_fp16": true
+ }
+ },
+ "matvae": {
+ "name": "matvae",
+ "args": {
+ "width": 256,
+ "depth": 3,
+ "z_dim": 2,
+ "p_drop": 0.05,
+ "use_learned_variances": false,
+ "use_additional_losses": true
+ }
+ }
+ },
+ "dataset": {
+ "name": "SparseVoxelMaterials",
+ "normalization_type": "log_minmax",
+ "args": {
+ "roots": "datasets/simready",
+ "image_size": 512,
+ "model": "dinov2_vitl14_reg",
+ "resolution": 64,
+ "min_aesthetic_score": 0.0,
+ "max_num_voxels": 32768,
+ "compute_material_stats": false
+ }
+ },
+ "trainer": {
+ "name": "SLatVaeMaterialsTrainer",
+ "args": {
+ "max_steps": 1000000,
+ "batch_size_per_gpu": 16,
+ "batch_split": 1,
+ "optimizer": {
+ "name": "AdamW",
+ "args": {
+ "lr": 1e-4,
+ "weight_decay": 0.0
+ }
+ },
+ "ema_rate": [
+ 0.9999
+ ],
+ "fp16_mode": "inflat_all",
+ "fp16_scale_growth": 0.001,
+ "elastic": {
+ "name": "LinearMemoryController",
+ "args": {
+ "target_ratio": 0.75,
+ "max_mem_ratio_start": 0.5
+ }
+ },
+ "grad_clip": {
+ "name": "AdaptiveGradClipper",
+ "args": {
+ "max_norm": 1.0,
+ "clip_percentile": 95
+ }
+ },
+ "i_log": 10,
+ "i_save": 2000,
+ "i_eval": 1000,
+ "loss_type": "l1"
+ }
+ }
+}
\ No newline at end of file
diff --git a/deps/vomp/configs/materials/material_vae/beta_tc_final.json b/deps/vomp/configs/materials/material_vae/beta_tc_final.json
new file mode 100644
index 0000000000000000000000000000000000000000..d9d0a6137f19bd61b299771a97667cc39a6713c7
--- /dev/null
+++ b/deps/vomp/configs/materials/material_vae/beta_tc_final.json
@@ -0,0 +1,68 @@
+{
+ "dry_run": false,
+ "standard_vae": false,
+
+ "data_csv": "datasets/latent_space/materials_filtered.csv",
+ "dataloader": {
+ "batch_size": 256,
+ "num_workers": 8,
+ "pin_memory": true,
+ "prefetch_factor": 4,
+ "persistent_workers": true
+ },
+
+ "project_dir": "./outputs/beta_tc",
+ "tracker_name": "tb_logs",
+ "log_with": "tensorboard",
+
+ "epochs": 25000,
+ "gradient_accumulation_steps": 1,
+ "keep_last_checkpoints": 3,
+
+ "mixed_precision": "no",
+ "use_stateful_dataloader": false,
+ "find_unused_parameters": false,
+
+ "compile": {
+ "enabled": false,
+ "backend": "inductor",
+ "mode": "default",
+ "fullgraph": true,
+ "dynamic": false
+ },
+
+ "optimizer": {
+ "lr": 5e-4,
+ "weight_decay": 1e-4,
+ "grad_clip_norm": 5.0
+ },
+
+ "lr_scheduler": {
+ "type": "cosine",
+ "eta_min": 1e-5
+ },
+
+ "free_nats": 0.1,
+ "alpha": 1.0,
+ "beta": 2.0,
+ "gamma": 1.0,
+ "iwae_K": 50,
+
+ "eval_interval": 1,
+ "save_interval": 1,
+ "visualization_interval": 1000,
+ "n_vis_samples": 5,
+ "n_vis_steps": 10,
+
+ "model": {
+ "width": 512,
+ "depth": 4,
+ "z_dim": 2,
+ "p_drop": 0.05,
+ "use_flow": false
+ },
+
+ "seed": 42,
+
+ "resume_from_checkpoint": null
+}
\ No newline at end of file
diff --git a/deps/vomp/configs/materials/material_vae/matvae.json b/deps/vomp/configs/materials/material_vae/matvae.json
new file mode 100644
index 0000000000000000000000000000000000000000..83f54c79a815b64e8553360d255bcfc14f63e5a0
--- /dev/null
+++ b/deps/vomp/configs/materials/material_vae/matvae.json
@@ -0,0 +1,76 @@
+{
+ "dry_run": false,
+ "standard_vae": false,
+
+ "data_csv": "datasets/latent_space/materials_filtered.csv",
+ "dataloader": {
+ "batch_size": 256,
+ "num_workers": 8,
+ "pin_memory": true,
+ "prefetch_factor": 4,
+ "persistent_workers": true
+ },
+
+ "project_dir": "./outputs/matvae",
+ "tracker_name": "tb_logs",
+ "log_with": "tensorboard",
+
+ "epochs": 850,
+ "gradient_accumulation_steps": 1,
+ "keep_last_checkpoints": 3,
+
+ "mixed_precision": "no",
+ "use_stateful_dataloader": false,
+ "find_unused_parameters": false,
+
+ "compile": {
+ "enabled": false,
+ "backend": "inductor",
+ "mode": "default",
+ "fullgraph": true,
+ "dynamic": false
+ },
+
+ "optimizer": {
+ "lr": 1e-4,
+ "weight_decay": 1e-4,
+ "grad_clip_norm": 5.0
+ },
+
+ "lr_scheduler": {
+ "type": "cosine",
+ "eta_min": 1e-5
+ },
+
+ "free_nats": 0.1,
+ "kl_annealing": true,
+ "kl_annealing_epochs": 200,
+ "recon_scale": 1.0,
+ "kl_weight": 1.0,
+ "iwae_K": 50,
+
+ "alpha": 1.0,
+ "beta": 2.0,
+ "gamma": 1.0,
+
+ "normalization_type": "log_minmax",
+
+ "eval_interval": 1,
+ "save_interval": 1,
+ "visualization_interval": 1000,
+ "n_vis_samples": 5,
+ "n_vis_steps": 10,
+
+ "model": {
+ "width": 256,
+ "depth": 3,
+ "z_dim": 2,
+ "p_drop": 0.05,
+ "use_learned_variances": false,
+ "use_additional_losses": true
+ },
+
+ "seed": 42,
+
+ "resume_from_checkpoint": null
+}
\ No newline at end of file
diff --git a/deps/vomp/configs/materials/material_vae/matvae_log_minmax_no_density.json b/deps/vomp/configs/materials/material_vae/matvae_log_minmax_no_density.json
new file mode 100644
index 0000000000000000000000000000000000000000..15aae69fa68470a568f23ae2a9cf2781cd96deb3
--- /dev/null
+++ b/deps/vomp/configs/materials/material_vae/matvae_log_minmax_no_density.json
@@ -0,0 +1,76 @@
+{
+ "dry_run": false,
+ "standard_vae": false,
+
+ "data_csv": "datasets/latent_space/materials_filtered.csv",
+ "dataloader": {
+ "batch_size": 256,
+ "num_workers": 8,
+ "pin_memory": true,
+ "prefetch_factor": 4,
+ "persistent_workers": true
+ },
+
+ "project_dir": "./outputs/matvae",
+ "tracker_name": "tb_logs",
+ "log_with": "tensorboard",
+
+ "epochs": 850,
+ "gradient_accumulation_steps": 1,
+ "keep_last_checkpoints": 3,
+
+ "mixed_precision": "no",
+ "use_stateful_dataloader": false,
+ "find_unused_parameters": false,
+
+ "compile": {
+ "enabled": false,
+ "backend": "inductor",
+ "mode": "default",
+ "fullgraph": true,
+ "dynamic": false
+ },
+
+ "optimizer": {
+ "lr": 1e-4,
+ "weight_decay": 1e-4,
+ "grad_clip_norm": 5.0
+ },
+
+ "lr_scheduler": {
+ "type": "cosine",
+ "eta_min": 1e-5
+ },
+
+ "free_nats": 0.1,
+ "kl_annealing": true,
+ "kl_annealing_epochs": 200,
+ "recon_scale": 1.0,
+ "kl_weight": 1.0,
+ "iwae_K": 50,
+
+ "alpha": 1.0,
+ "beta": 2.0,
+ "gamma": 1.0,
+
+ "normalization_type": "log_minmax_no_density",
+
+ "eval_interval": 1,
+ "save_interval": 1,
+ "visualization_interval": 1000,
+ "n_vis_samples": 5,
+ "n_vis_steps": 10,
+
+ "model": {
+ "width": 256,
+ "depth": 3,
+ "z_dim": 2,
+ "p_drop": 0.05,
+ "use_learned_variances": false,
+ "use_additional_losses": true
+ },
+
+ "seed": 42,
+
+ "resume_from_checkpoint": null
+}
diff --git a/deps/vomp/configs/materials/material_vae/matvae_no_beta_tc.json b/deps/vomp/configs/materials/material_vae/matvae_no_beta_tc.json
new file mode 100644
index 0000000000000000000000000000000000000000..f8503782610c7bed502fb0cee9a905bce3bb5094
--- /dev/null
+++ b/deps/vomp/configs/materials/material_vae/matvae_no_beta_tc.json
@@ -0,0 +1,77 @@
+{
+ "dry_run": false,
+ "standard_vae": true,
+
+ "data_csv": "datasets/latent_space/materials_filtered.csv",
+ "dataloader": {
+ "batch_size": 256,
+ "num_workers": 8,
+ "pin_memory": true,
+ "prefetch_factor": 4,
+ "persistent_workers": true
+ },
+
+ "project_dir": "./outputs/matvae_ablations/no_beta_tc",
+ "tracker_name": "tb_logs",
+ "log_with": "tensorboard",
+
+ "epochs": 850,
+ "gradient_accumulation_steps": 1,
+ "keep_last_checkpoints": 3,
+
+ "mixed_precision": "no",
+ "use_stateful_dataloader": false,
+ "find_unused_parameters": false,
+
+ "compile": {
+ "enabled": false,
+ "backend": "inductor",
+ "mode": "default",
+ "fullgraph": true,
+ "dynamic": false
+ },
+
+ "optimizer": {
+ "lr": 1e-4,
+ "weight_decay": 1e-4,
+ "grad_clip_norm": 5.0
+ },
+
+ "lr_scheduler": {
+ "type": "cosine",
+ "eta_min": 1e-5
+ },
+
+ "free_nats": 0.1,
+ "kl_annealing": true,
+ "kl_annealing_epochs": 200,
+ "recon_scale": 1.0,
+ "kl_weight": 1.0,
+ "iwae_K": 50,
+
+ "alpha": 1.0,
+ "beta": 2.0,
+ "gamma": 1.0,
+
+ "normalization_type": "log_minmax",
+
+ "eval_interval": 1,
+ "save_interval": 1,
+ "visualization_interval": 1000,
+ "n_vis_samples": 5,
+ "n_vis_steps": 10,
+
+ "model": {
+ "width": 256,
+ "depth": 3,
+ "z_dim": 2,
+ "p_drop": 0.05,
+ "use_learned_variances": false,
+ "use_additional_losses": true,
+ "use_flow": true
+ },
+
+ "seed": 42,
+
+ "resume_from_checkpoint": null
+}
\ No newline at end of file
diff --git a/deps/vomp/configs/materials/material_vae/matvae_no_flow.json b/deps/vomp/configs/materials/material_vae/matvae_no_flow.json
new file mode 100644
index 0000000000000000000000000000000000000000..8dedf8dbb145ba5dafc839e999ab680e84bb272a
--- /dev/null
+++ b/deps/vomp/configs/materials/material_vae/matvae_no_flow.json
@@ -0,0 +1,77 @@
+{
+ "dry_run": false,
+ "standard_vae": false,
+
+ "data_csv": "datasets/latent_space/materials_filtered.csv",
+ "dataloader": {
+ "batch_size": 256,
+ "num_workers": 8,
+ "pin_memory": true,
+ "prefetch_factor": 4,
+ "persistent_workers": true
+ },
+
+ "project_dir": "./outputs/matvae_ablations/no_flow",
+ "tracker_name": "tb_logs",
+ "log_with": "tensorboard",
+
+ "epochs": 850,
+ "gradient_accumulation_steps": 1,
+ "keep_last_checkpoints": 3,
+
+ "mixed_precision": "no",
+ "use_stateful_dataloader": false,
+ "find_unused_parameters": false,
+
+ "compile": {
+ "enabled": false,
+ "backend": "inductor",
+ "mode": "default",
+ "fullgraph": true,
+ "dynamic": false
+ },
+
+ "optimizer": {
+ "lr": 1e-4,
+ "weight_decay": 1e-4,
+ "grad_clip_norm": 5.0
+ },
+
+ "lr_scheduler": {
+ "type": "cosine",
+ "eta_min": 1e-5
+ },
+
+ "free_nats": 0.1,
+ "kl_annealing": true,
+ "kl_annealing_epochs": 200,
+ "recon_scale": 1.0,
+ "kl_weight": 1.0,
+ "iwae_K": 50,
+
+ "alpha": 1.0,
+ "beta": 2.0,
+ "gamma": 1.0,
+
+ "normalization_type": "log_minmax",
+
+ "eval_interval": 1,
+ "save_interval": 1,
+ "visualization_interval": 1000,
+ "n_vis_samples": 5,
+ "n_vis_steps": 10,
+
+ "model": {
+ "width": 256,
+ "depth": 3,
+ "z_dim": 2,
+ "p_drop": 0.05,
+ "use_learned_variances": false,
+ "use_additional_losses": true,
+ "use_flow": false
+ },
+
+ "seed": 42,
+
+ "resume_from_checkpoint": null
+}
\ No newline at end of file
diff --git a/deps/vomp/configs/materials/material_vae/matvae_no_free_nats.json b/deps/vomp/configs/materials/material_vae/matvae_no_free_nats.json
new file mode 100644
index 0000000000000000000000000000000000000000..0b1d7bcefaed996aee6b67340c9c8fd02096a943
--- /dev/null
+++ b/deps/vomp/configs/materials/material_vae/matvae_no_free_nats.json
@@ -0,0 +1,76 @@
+{
+ "dry_run": false,
+ "standard_vae": false,
+
+ "data_csv": "datasets/latent_space/materials_filtered.csv",
+ "dataloader": {
+ "batch_size": 256,
+ "num_workers": 8,
+ "pin_memory": true,
+ "prefetch_factor": 4,
+ "persistent_workers": true
+ },
+
+ "project_dir": "./outputs/matvae_ablations/no_free_nats",
+ "tracker_name": "tb_logs",
+ "log_with": "tensorboard",
+
+ "epochs": 850,
+ "gradient_accumulation_steps": 1,
+ "keep_last_checkpoints": 3,
+
+ "mixed_precision": "no",
+ "use_stateful_dataloader": false,
+ "find_unused_parameters": false,
+
+ "compile": {
+ "enabled": false,
+ "backend": "inductor",
+ "mode": "default",
+ "fullgraph": true,
+ "dynamic": false
+ },
+
+ "optimizer": {
+ "lr": 1e-4,
+ "weight_decay": 1e-4,
+ "grad_clip_norm": 5.0
+ },
+
+ "lr_scheduler": {
+ "type": "cosine",
+ "eta_min": 1e-5
+ },
+
+ "free_nats": 0.0,
+ "kl_annealing": true,
+ "kl_annealing_epochs": 200,
+ "recon_scale": 1.0,
+ "kl_weight": 1.0,
+ "iwae_K": 50,
+
+ "alpha": 1.0,
+ "beta": 2.0,
+ "gamma": 1.0,
+
+ "normalization_type": "log_minmax",
+
+ "eval_interval": 1,
+ "save_interval": 1,
+ "visualization_interval": 1000,
+ "n_vis_samples": 5,
+ "n_vis_steps": 10,
+
+ "model": {
+ "width": 256,
+ "depth": 3,
+ "z_dim": 2,
+ "p_drop": 0.05,
+ "use_learned_variances": false,
+ "use_additional_losses": true
+ },
+
+ "seed": 42,
+
+ "resume_from_checkpoint": null
+}
\ No newline at end of file
diff --git a/deps/vomp/configs/materials/material_vae/matvae_standard.json b/deps/vomp/configs/materials/material_vae/matvae_standard.json
new file mode 100644
index 0000000000000000000000000000000000000000..62ed067e95878c52dffb88ed2cdd11d5a971985e
--- /dev/null
+++ b/deps/vomp/configs/materials/material_vae/matvae_standard.json
@@ -0,0 +1,76 @@
+{
+ "dry_run": false,
+ "standard_vae": true,
+
+ "data_csv": "datasets/latent_space/materials_filtered.csv",
+ "dataloader": {
+ "batch_size": 256,
+ "num_workers": 8,
+ "pin_memory": true,
+ "prefetch_factor": 4,
+ "persistent_workers": true
+ },
+
+ "project_dir": "./outputs/matvae",
+ "tracker_name": "tb_logs",
+ "log_with": "tensorboard",
+
+ "epochs": 850,
+ "gradient_accumulation_steps": 1,
+ "keep_last_checkpoints": 3,
+
+ "mixed_precision": "no",
+ "use_stateful_dataloader": false,
+ "find_unused_parameters": false,
+
+ "compile": {
+ "enabled": false,
+ "backend": "inductor",
+ "mode": "default",
+ "fullgraph": true,
+ "dynamic": false
+ },
+
+ "optimizer": {
+ "lr": 1e-4,
+ "weight_decay": 1e-4,
+ "grad_clip_norm": 5.0
+ },
+
+ "lr_scheduler": {
+ "type": "cosine",
+ "eta_min": 1e-5
+ },
+
+ "free_nats": 0.1,
+ "kl_annealing": true,
+ "kl_annealing_epochs": 200,
+ "recon_scale": 1.0,
+ "kl_weight": 1.0,
+ "iwae_K": 50,
+
+ "alpha": 1.0,
+ "beta": 2.0,
+ "gamma": 1.0,
+
+ "normalization_type": "log_minmax",
+
+ "eval_interval": 1,
+ "save_interval": 1,
+ "visualization_interval": 1000,
+ "n_vis_samples": 5,
+ "n_vis_steps": 10,
+
+ "model": {
+ "width": 256,
+ "depth": 3,
+ "z_dim": 2,
+ "p_drop": 0.05,
+ "use_learned_variances": false,
+ "use_additional_losses": true
+ },
+
+ "seed": 42,
+
+ "resume_from_checkpoint": null
+ }
\ No newline at end of file
diff --git a/deps/vomp/configs/materials/material_vae/matvae_standard_norm.json b/deps/vomp/configs/materials/material_vae/matvae_standard_norm.json
new file mode 100644
index 0000000000000000000000000000000000000000..e29928a8bc42af6b25bfe12ff811566620743f06
--- /dev/null
+++ b/deps/vomp/configs/materials/material_vae/matvae_standard_norm.json
@@ -0,0 +1,76 @@
+{
+ "dry_run": false,
+ "standard_vae": true,
+
+ "data_csv": "datasets/latent_space/materials_filtered.csv",
+ "dataloader": {
+ "batch_size": 256,
+ "num_workers": 8,
+ "pin_memory": true,
+ "prefetch_factor": 4,
+ "persistent_workers": true
+ },
+
+ "project_dir": "./outputs/matvae",
+ "tracker_name": "tb_logs",
+ "log_with": "tensorboard",
+
+ "epochs": 850,
+ "gradient_accumulation_steps": 1,
+ "keep_last_checkpoints": 3,
+
+ "mixed_precision": "no",
+ "use_stateful_dataloader": false,
+ "find_unused_parameters": false,
+
+ "compile": {
+ "enabled": false,
+ "backend": "inductor",
+ "mode": "default",
+ "fullgraph": true,
+ "dynamic": false
+ },
+
+ "optimizer": {
+ "lr": 1e-4,
+ "weight_decay": 1e-4,
+ "grad_clip_norm": 5.0
+ },
+
+ "lr_scheduler": {
+ "type": "cosine",
+ "eta_min": 1e-5
+ },
+
+ "free_nats": 0.1,
+ "kl_annealing": true,
+ "kl_annealing_epochs": 200,
+ "recon_scale": 1.0,
+ "kl_weight": 1.0,
+ "iwae_K": 50,
+
+ "alpha": 1.0,
+ "beta": 2.0,
+ "gamma": 1.0,
+
+ "normalization_type": "standard",
+
+ "eval_interval": 1,
+ "save_interval": 1,
+ "visualization_interval": 1000,
+ "n_vis_samples": 5,
+ "n_vis_steps": 10,
+
+ "model": {
+ "width": 256,
+ "depth": 3,
+ "z_dim": 2,
+ "p_drop": 0.05,
+ "use_learned_variances": false,
+ "use_additional_losses": true
+ },
+
+ "seed": 42,
+
+ "resume_from_checkpoint": null
+ }
\ No newline at end of file
diff --git a/deps/vomp/configs/materials/material_vae/standard_vae_final.json b/deps/vomp/configs/materials/material_vae/standard_vae_final.json
new file mode 100644
index 0000000000000000000000000000000000000000..c8f296484d20ef37b083a23ea679cf2160e09a41
--- /dev/null
+++ b/deps/vomp/configs/materials/material_vae/standard_vae_final.json
@@ -0,0 +1,67 @@
+{
+ "dry_run": false,
+ "standard_vae": true,
+
+ "data_csv": "datasets/latent_space/materials_filtered.csv",
+ "dataloader": {
+ "batch_size": 256,
+ "num_workers": 8,
+ "pin_memory": true,
+ "prefetch_factor": 4,
+ "persistent_workers": true
+ },
+
+ "project_dir": "./outputs/standard_vae",
+ "tracker_name": "tb_logs",
+ "log_with": "tensorboard",
+
+ "epochs": 25000,
+ "gradient_accumulation_steps": 1,
+ "keep_last_checkpoints": 3,
+
+ "mixed_precision": "no",
+ "use_stateful_dataloader": false,
+ "find_unused_parameters": false,
+
+ "compile": {
+ "enabled": false,
+ "backend": "inductor",
+ "mode": "default",
+ "fullgraph": true,
+ "dynamic": false
+ },
+
+ "optimizer": {
+ "lr": 1e-3,
+ "weight_decay": 1e-4,
+ "grad_clip_norm": 5.0
+ },
+
+ "lr_scheduler": {
+ "type": "cosine",
+ "eta_min": 1e-5
+ },
+
+ "free_nats": 0.1,
+ "kl_annealing": true,
+ "kl_annealing_epochs": 200,
+ "recon_scale": 1.0,
+ "iwae_K": 50,
+
+ "eval_interval": 1,
+ "save_interval": 1,
+ "visualization_interval": 1000,
+ "n_vis_samples": 5,
+ "n_vis_steps": 10,
+
+ "model": {
+ "width": 256,
+ "depth": 3,
+ "z_dim": 2,
+ "p_drop": 0.05
+ },
+
+ "seed": 42,
+
+ "resume_from_checkpoint": null
+}
\ No newline at end of file
diff --git a/deps/vomp/configs/sim/armchair_and_orange.json b/deps/vomp/configs/sim/armchair_and_orange.json
new file mode 100644
index 0000000000000000000000000000000000000000..f688dd52c0033b9c3741fbdcccfaf033f0a0e910
--- /dev/null
+++ b/deps/vomp/configs/sim/armchair_and_orange.json
@@ -0,0 +1,59 @@
+{
+ "simulation": {
+ "dt": 0.01,
+ "gravity": [0.0, -9.8, 0.0],
+ "max_frames": 500,
+ "auto_start": false,
+ "contact": {
+ "friction_enable": true,
+ "d_hat": 0.005
+ }
+ },
+ "engine": {
+ "type": "cuda"
+ },
+ "contact_model": {
+ "friction": 0.5,
+ "contact_resistance": 2.0
+ },
+ "ground": {
+ "enable": true,
+ "height": 0.0
+ },
+ "objects": [
+ {
+ "name": "armchair_1",
+ "type": "msh",
+ "msh_path": "assets/armchair/armchair.msh",
+ "normalize_visual_mesh": false,
+ "scale": 1.0,
+ "translation": [0.0, 0.004, 0.0],
+ "rotation": [270.0, 0.0, 0.0],
+ "material": {
+ "file": "assets/armchair/materials_aligned.npz"
+ }
+ },
+ {
+ "name": "orange_1",
+ "type": "msh",
+ "msh_path": "assets/orange/orange_02_inst_base.msh",
+ "normalize_visual_mesh": false,
+ "scale": 3.0,
+ "translation": [0.0, 2.0, 0.0],
+ "rotation": [0.0, 0.0, 0.0],
+ "material": {
+ "file": "assets/orange/materials.npz"
+ }
+ }
+ ],
+ "gui": {
+ "enable": true
+ },
+ "output": {
+ "directory": "./outputs/simulation_output/armchair_and_orange",
+ "save_meshes": true
+ },
+ "logging": {
+ "level": "warn"
+ }
+}
\ No newline at end of file
diff --git a/deps/vomp/configs/sim/falling_armchair.json b/deps/vomp/configs/sim/falling_armchair.json
new file mode 100644
index 0000000000000000000000000000000000000000..5a18314a80f04566fb85a545e0752834cb8323c4
--- /dev/null
+++ b/deps/vomp/configs/sim/falling_armchair.json
@@ -0,0 +1,48 @@
+{
+ "simulation": {
+ "dt": 0.01,
+ "gravity": [0.0, -9.8, 0.0],
+ "max_frames": 500,
+ "auto_start": false,
+ "contact": {
+ "friction_enable": true,
+ "d_hat": 0.01
+ }
+ },
+ "engine": {
+ "type": "cuda"
+ },
+ "contact_model": {
+ "friction": 0.5,
+ "contact_resistance": 1.0
+ },
+ "ground": {
+ "enable": true,
+ "height": 0.0
+ },
+ "objects": [
+ {
+ "name": "armchair_1",
+ "type": "msh",
+ "msh_path": "assets/armchair/armchair.msh",
+ "visual_mesh": "assets/armchair/armchair_inst_base.obj",
+ "normalize_visual_mesh": false,
+ "scale": 1.0,
+ "translation": [0.0, 1.0, 0.0],
+ "rotation": [0.0, 0.0, 0.0],
+ "material": {
+ "file": "assets/armchair/materials_aligned.npz"
+ }
+ }
+ ],
+ "gui": {
+ "enable": true
+ },
+ "output": {
+ "directory": "./outputs/simulation_output/falling_armchair",
+ "save_meshes": true
+ },
+ "logging": {
+ "level": "warn"
+ }
+}
\ No newline at end of file
diff --git a/deps/vomp/configs/sim/falling_bar_stool.json b/deps/vomp/configs/sim/falling_bar_stool.json
new file mode 100644
index 0000000000000000000000000000000000000000..028b1628f5199617b62f9b104408e9fee3ddf19a
--- /dev/null
+++ b/deps/vomp/configs/sim/falling_bar_stool.json
@@ -0,0 +1,50 @@
+{
+ "simulation": {
+ "dt": 0.01,
+ "gravity": [0.0, -9.8, 0.0],
+ "max_frames": 500,
+ "auto_start": false,
+ "contact": {
+ "friction_enable": true,
+ "d_hat": 0.01
+ }
+ },
+ "engine": {
+ "type": "cuda"
+ },
+ "contact_model": {
+ "friction": 0.5,
+ "contact_resistance": 1.0
+ },
+ "ground": {
+ "enable": true,
+ "height": 0.0
+ },
+ "objects": [
+ {
+ "name": "bar_stool_1",
+ "type": "msh",
+ "msh_path": "assets/bar_stool/bar_stool_inst_base.msh",
+ "visual_mesh": "assets/bar_stool/bar_stool_inst_base.obj",
+ "normalize_visual_mesh": false,
+ "scale": 1.0,
+ "translation": [0.0, 2.0, 0.0],
+ "rotation": [0.0, 0.0, 0.0],
+ "material": {
+ "youngs_modulus": 1e5,
+ "density": 900.0,
+ "poisson_ratio": 0.3
+ }
+ }
+ ],
+ "gui": {
+ "enable": true
+ },
+ "output": {
+ "directory": "./outputs/simulation_output/falling_bar_stool",
+ "save_meshes": true
+ },
+ "logging": {
+ "level": "warn"
+ }
+}
\ No newline at end of file
diff --git a/deps/vomp/configs/sim/falling_birch.json b/deps/vomp/configs/sim/falling_birch.json
new file mode 100644
index 0000000000000000000000000000000000000000..51701c91d45fc9f77f9b0d051c3654c157f2da40
--- /dev/null
+++ b/deps/vomp/configs/sim/falling_birch.json
@@ -0,0 +1,50 @@
+{
+ "simulation": {
+ "dt": 0.005,
+ "gravity": [0.0, -9.8, 0.0],
+ "max_frames": 500,
+ "auto_start": false,
+ "contact": {
+ "friction_enable": true,
+ "d_hat": 0.01
+ }
+ },
+ "engine": {
+ "type": "cuda"
+ },
+ "contact_model": {
+ "friction": 0.5,
+ "contact_resistance": 1.0
+ },
+ "ground": {
+ "enable": true,
+ "height": 0.0
+ },
+ "objects": [
+ {
+ "name": "birch_1",
+ "type": "msh",
+ "msh_path": "assets/birch/birch.msh",
+ "visual_mesh": "assets/birch/birch_lowbackseat_inst_base.obj",
+ "normalize_visual_mesh": false,
+ "scale": 1.0,
+ "translation": [0.0, 2.0, 0.0],
+ "rotation": [0.0, 0.0, 0.0],
+ "material": {
+ "youngs_modulus": 1e6,
+ "density": 1000.0,
+ "poisson_ratio": 0.45
+ }
+ }
+ ],
+ "gui": {
+ "enable": true
+ },
+ "output": {
+ "directory": "./outputs/simulation_output/falling_birch",
+ "save_meshes": true
+ },
+ "logging": {
+ "level": "warn"
+ }
+}
\ No newline at end of file
diff --git a/deps/vomp/configs/sim/falling_oranges.json b/deps/vomp/configs/sim/falling_oranges.json
new file mode 100644
index 0000000000000000000000000000000000000000..bf6ed0a5e49b9ff4f790b9fd6f18b61f6450e317
--- /dev/null
+++ b/deps/vomp/configs/sim/falling_oranges.json
@@ -0,0 +1,80 @@
+{
+ "simulation": {
+ "dt": 0.02,
+ "gravity": [0.0, -9.8, 0.0],
+ "max_frames": 500,
+ "auto_start": false,
+ "contact": {
+ "friction_enable": true,
+ "d_hat": 0.01
+ }
+ },
+ "engine": {
+ "type": "cuda"
+ },
+ "contact_model": {
+ "friction": 0.5,
+ "contact_resistance": 1.0
+ },
+ "ground": {
+ "enable": true,
+ "height": 0.0
+ },
+ "objects": [
+ {
+ "name": "orange_1",
+ "type": "voxel",
+ "voxel_path": "assets/orange/voxels.ply",
+ "visual_mesh": "assets/orange/orange_02_inst_base.obj",
+ "normalize_visual_mesh": true,
+ "voxel_size": 1.0,
+ "scale": 1.0,
+ "max_voxels": 32000,
+ "translation": [0.0, 2.0, 0.0],
+ "rotation": [0.0, 0.0, 0.0],
+ "material": {
+ "file": "assets/orange/materials.npz"
+ }
+ },
+ {
+ "name": "orange_2",
+ "type": "voxel",
+ "voxel_path": "assets/orange/voxels.ply",
+ "visual_mesh": "assets/orange/orange_02_inst_base.obj",
+ "normalize_visual_mesh": true,
+ "voxel_size": 1.0,
+ "scale": 1.0,
+ "max_voxels": 32000,
+ "translation": [0.0, 3.5, 0.0],
+ "rotation": [0.0, 0.0, 0.0],
+ "material": {
+ "file": "assets/orange/materials.npz"
+ }
+ },
+ {
+ "name": "orange_3",
+ "type": "voxel",
+ "voxel_path": "assets/orange/voxels.ply",
+ "visual_mesh": "assets/orange/orange_02_inst_base.obj",
+ "normalize_visual_mesh": true,
+ "voxel_size": 1.0,
+ "scale": 1.0,
+ "max_voxels": 32000,
+ "translation": [0.0, 5.0, 0.0],
+ "rotation": [0.0, 0.0, 0.0],
+ "material": {
+ "file": "assets/orange/materials.npz"
+ }
+ }
+ ],
+ "gui": {
+ "enable": true
+ },
+ "output": {
+ "directory": "./outputs/simulation_output/falling_oranges",
+ "save_meshes": true
+ },
+ "logging": {
+ "level": "warn"
+ }
+}
\ No newline at end of file
diff --git a/deps/vomp/configs/sim/falling_sphere_soft.json b/deps/vomp/configs/sim/falling_sphere_soft.json
new file mode 100644
index 0000000000000000000000000000000000000000..ccb478b2de136f5b6eeb8ae95e8d0cabe20f3fc6
--- /dev/null
+++ b/deps/vomp/configs/sim/falling_sphere_soft.json
@@ -0,0 +1,51 @@
+{
+ "simulation": {
+ "dt": 0.005,
+ "gravity": [0.0, -9.8, 0.0],
+ "max_frames": 500,
+ "auto_start": false,
+ "contact": {
+ "friction_enable": false,
+ "d_hat": 0.005
+ }
+ },
+ "engine": {
+ "type": "cuda"
+ },
+ "contact_model": {
+ "friction": 0.5,
+ "contact_resistance": 0.01
+ },
+ "ground": {
+ "enable": true,
+ "height": 0.0
+ },
+ "objects": [
+ {
+ "name": "soft_sphere",
+ "type": "msh",
+ "msh_path": "assets/sphere/sphere_tetrahedral.msh",
+ "visual_mesh": "assets/sphere/sphere_visual.obj",
+ "normalize_visual_mesh": false,
+ "scale": 1.0,
+ "translation": [0.0, 0.3, 0.0],
+ "rotation": [0.0, 0.0, 0.0],
+ "material": {
+ "type": "StableNeoHookean",
+ "young_modulus": 1e4,
+ "poisson_ratio": 0.3,
+ "density": 1000
+ }
+ }
+ ],
+ "gui": {
+ "enable": true
+ },
+ "output": {
+ "directory": "./outputs/simulation_output/fem",
+ "save_meshes": true
+ },
+ "logging": {
+ "level": "info"
+ }
+}
\ No newline at end of file
diff --git a/deps/vomp/configs/sim/zag_and_falling_orange.json b/deps/vomp/configs/sim/zag_and_falling_orange.json
new file mode 100644
index 0000000000000000000000000000000000000000..8429eb06d3e58f3d92622b45f46afac33f4b90f0
--- /dev/null
+++ b/deps/vomp/configs/sim/zag_and_falling_orange.json
@@ -0,0 +1,59 @@
+{
+ "simulation": {
+ "dt": 0.01,
+ "gravity": [0.0, -9.8, 0.0],
+ "max_frames": 500,
+ "auto_start": false,
+ "contact": {
+ "friction_enable": true,
+ "d_hat": 0.005
+ }
+ },
+ "engine": {
+ "type": "cuda"
+ },
+ "contact_model": {
+ "friction": 0.5,
+ "contact_resistance": 2.0
+ },
+ "ground": {
+ "enable": true,
+ "height": 0.0
+ },
+ "objects": [
+ {
+ "name": "zag_middle_base",
+ "type": "msh",
+ "msh_path": "assets/zag/zag_middle_inst_base.msh",
+ "normalize_visual_mesh": false,
+ "scale": 1.0,
+ "translation": [0.0, 0.001, 0.0],
+ "rotation": [270.0, 0.0, 0.0],
+ "material": {
+ "file": "assets/zag/materials.npz"
+ }
+ },
+ {
+ "name": "orange_1",
+ "type": "msh",
+ "msh_path": "assets/orange/orange_02_inst_base.msh",
+ "normalize_visual_mesh": false,
+ "scale": 3.0,
+ "translation": [0.0, 2.0, 0.0],
+ "rotation": [0.0, 0.0, 0.0],
+ "material": {
+ "file": "assets/orange/materials.npz"
+ }
+ }
+ ],
+ "gui": {
+ "enable": true
+ },
+ "output": {
+ "directory": "./outputs/simulation_output/zag_and_falling_orange",
+ "save_meshes": true
+ },
+ "logging": {
+ "level": "warn"
+ }
+}
\ No newline at end of file
diff --git a/deps/vomp/configs/sim/zag_and_falling_oranges.json b/deps/vomp/configs/sim/zag_and_falling_oranges.json
new file mode 100644
index 0000000000000000000000000000000000000000..816174ef2b9f1c5db855aaa4380ca5e053465272
--- /dev/null
+++ b/deps/vomp/configs/sim/zag_and_falling_oranges.json
@@ -0,0 +1,98 @@
+{
+ "simulation": {
+ "dt": 0.01,
+ "gravity": [0.0, -9.8, 0.0],
+ "max_frames": 500,
+ "auto_start": false,
+ "contact": {
+ "friction_enable": true,
+ "d_hat": 0.005
+ }
+ },
+ "engine": {
+ "type": "cuda"
+ },
+ "contact_model": {
+ "friction": 0.5,
+ "contact_resistance": 2.0
+ },
+ "ground": {
+ "enable": true,
+ "height": 0.0
+ },
+ "objects": [
+ {
+ "name": "zag_middle_base",
+ "type": "msh",
+ "msh_path": "assets/zag/zag_middle_inst_base.msh",
+ "normalize_visual_mesh": false,
+ "scale": 1.0,
+ "translation": [0.0, 0.01, 0.0],
+ "rotation": [270.0, 0.0, 0.0],
+ "apply_boundary_conditions": true,
+ "boundary_fix_percentage": 0.15,
+ "material": {
+ "file": "assets/zag/materials.npz"
+ }
+ },
+
+ {
+ "name": "lemon_2",
+ "type": "msh",
+ "msh_path": "assets/orange/orange_02_inst_base.msh",
+ "normalize_visual_mesh": false,
+ "scale": 2.0,
+ "translation": [0.15, 1.0, 0.0],
+ "rotation": [0.0, 0.0, 0.0],
+ "material": {
+ "file": "assets/orange/materials.npz"
+ }
+ },
+ {
+ "name": "lemon_3",
+ "type": "msh",
+ "msh_path": "assets/orange/orange_02_inst_base.msh",
+ "normalize_visual_mesh": false,
+ "scale": 2.0,
+ "translation": [-0.15, 1.0, 0.0],
+ "rotation": [0.0, 0.0, 0.0],
+ "material": {
+ "file": "assets/orange/materials.npz"
+ }
+ },
+ {
+ "name": "lemon_4",
+ "type": "msh",
+ "msh_path": "assets/orange/orange_02_inst_base.msh",
+ "normalize_visual_mesh": false,
+ "scale": 2.0,
+ "translation": [0.0, 1.0, 0.15],
+ "rotation": [0.0, 0.0, 0.0],
+ "material": {
+ "file": "assets/orange/materials.npz"
+ }
+ },
+ {
+ "name": "lemon_5",
+ "type": "msh",
+ "msh_path": "assets/orange/orange_02_inst_base.msh",
+ "normalize_visual_mesh": false,
+ "scale": 2.0,
+ "translation": [0.0, 1.0, -0.15],
+ "rotation": [0.0, 0.0, 0.0],
+ "material": {
+ "file": "assets/orange/materials.npz"
+ }
+ }
+ ],
+ "gui": {
+ "enable": true
+ },
+ "output": {
+ "directory": "./outputs/simulation_output/zag_and_falling_oranges",
+ "save_meshes": true
+ },
+ "logging": {
+ "level": "warn"
+ }
+}
\ No newline at end of file
diff --git a/deps/vomp/dataset_toolkits/abo/ABO500.py b/deps/vomp/dataset_toolkits/abo/ABO500.py
new file mode 100644
index 0000000000000000000000000000000000000000..4fdc931a7d8c651a5b322b72deacb4d1ef2c3bae
--- /dev/null
+++ b/deps/vomp/dataset_toolkits/abo/ABO500.py
@@ -0,0 +1,204 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import json
+import argparse
+import pandas as pd
+from concurrent.futures import ThreadPoolExecutor
+from tqdm import tqdm
+import hashlib
+
+
+def add_args(parser: argparse.ArgumentParser):
+ parser.add_argument(
+ "--abo_500_dir",
+ type=str,
+ default="/home/rdagli/code/datasets/abo_500",
+ help="Path to the ABO 500 dataset directory",
+ )
+ parser.add_argument(
+ "--abo_3d_dir",
+ type=str,
+ default="/home/rdagli/code/datasets/abo-3dmodels/3dmodels",
+ help="Path to the ABO 3D models directory",
+ )
+ parser.add_argument(
+ "--split",
+ type=str,
+ default="all",
+ choices=["train", "val", "test", "all"],
+ help="Which split to process",
+ )
+ parser.add_argument(
+ "--limit",
+ type=int,
+ default=None,
+ help="Limit to first N objects for testing",
+ )
+
+
+def get_file_hash(file_path):
+ """Get SHA256 hash of a file."""
+ hasher = hashlib.sha256()
+ with open(file_path, "rb") as f:
+ for chunk in iter(lambda: f.read(4096), b""):
+ hasher.update(chunk)
+ return hasher.hexdigest()
+
+
+def get_metadata(abo_500_dir, abo_3d_dir, split="all", limit=None, **kwargs):
+ """Get metadata for ABO 500 dataset."""
+ splits_path = os.path.join(abo_500_dir, "splits.json")
+
+ if not os.path.exists(splits_path):
+ raise FileNotFoundError(f"Splits file not found at {splits_path}")
+
+ with open(splits_path, "r") as f:
+ splits_data = json.load(f)
+
+ if split == "all":
+ object_ids = splits_data["train"] + splits_data["val"] + splits_data["test"]
+ else:
+ object_ids = splits_data[split]
+
+ # Apply limit if specified
+ if limit is not None:
+ object_ids = object_ids[:limit]
+
+ print(f"Processing {len(object_ids)} objects from {split} split")
+
+ # Create metadata records
+ metadata_records = []
+ missing_files = []
+
+ for object_id in tqdm(object_ids, desc="Building metadata"):
+ # Extract base ID (remove suffix after underscore if present)
+ base_id = object_id.split("_")[0]
+
+ # Search for GLB file - try multiple patterns and locations
+ glb_path = None
+
+ # Pattern 1: Try with base_id in the directory based on first character
+ first_char = base_id[0]
+ candidate_path = os.path.join(
+ abo_3d_dir, "original", first_char, f"{base_id}.glb"
+ )
+ if os.path.exists(candidate_path):
+ glb_path = candidate_path
+ else:
+ # Pattern 2: Try with full object_id (without underscore splitting)
+ first_char_full = object_id[0]
+ candidate_path = os.path.join(
+ abo_3d_dir, "original", first_char_full, f"{object_id}.glb"
+ )
+ if os.path.exists(candidate_path):
+ glb_path = candidate_path
+ else:
+ # Pattern 3: Search in all directories for the base_id
+ for dir_name in os.listdir(os.path.join(abo_3d_dir, "original")):
+ dir_path = os.path.join(abo_3d_dir, "original", dir_name)
+ if os.path.isdir(dir_path):
+ candidate_path = os.path.join(dir_path, f"{base_id}.glb")
+ if os.path.exists(candidate_path):
+ glb_path = candidate_path
+ break
+ # Also try the full object_id
+ candidate_path = os.path.join(dir_path, f"{object_id}.glb")
+ if os.path.exists(candidate_path):
+ glb_path = candidate_path
+ break
+
+ if glb_path and os.path.exists(glb_path):
+ # Get file hash
+ try:
+ sha256 = get_file_hash(glb_path)
+ metadata_records.append(
+ {
+ "object_id": object_id,
+ "sha256": sha256,
+ "local_path": glb_path,
+ "file_type": "glb",
+ "split": split,
+ "dataset": "ABO500",
+ }
+ )
+ except Exception as e:
+ print(f"Error processing {object_id}: {e}")
+ missing_files.append(object_id)
+ else:
+ print(
+ f"Warning: GLB file not found for {object_id} (tried base_id: {base_id})"
+ )
+ missing_files.append(object_id)
+
+ if missing_files:
+ print(f"Warning: {len(missing_files)} objects have missing GLB files")
+
+ metadata = pd.DataFrame(metadata_records)
+ return metadata
+
+
+def download(metadata, output_dir, **kwargs):
+ """For ABO 500, files are already downloaded, so just return local paths."""
+ download_records = []
+
+ for _, row in metadata.iterrows():
+ download_records.append(
+ {"sha256": row["sha256"], "local_path": row["local_path"]}
+ )
+
+ return pd.DataFrame(download_records)
+
+
+def foreach_instance(
+ metadata, output_dir, func, max_workers=None, desc="Processing objects"
+) -> pd.DataFrame:
+ """Process each instance in the metadata."""
+ import os
+ from concurrent.futures import ThreadPoolExecutor
+ from tqdm import tqdm
+
+ # Convert to list of records
+ metadata_records = metadata.to_dict("records")
+
+ # Processing objects
+ records = []
+ max_workers = max_workers or os.cpu_count()
+
+ try:
+ with (
+ ThreadPoolExecutor(max_workers=max_workers) as executor,
+ tqdm(total=len(metadata_records), desc=desc) as pbar,
+ ):
+
+ def worker(metadatum):
+ try:
+ local_path = metadatum["local_path"]
+ sha256 = metadatum["sha256"]
+ record = func(local_path, sha256)
+ if record is not None:
+ records.append(record)
+ pbar.update()
+ except Exception as e:
+ print(f"Error processing object {sha256}: {e}")
+ pbar.update()
+
+ executor.map(worker, metadata_records)
+ executor.shutdown(wait=True)
+ except Exception as e:
+ print(f"Error happened during processing: {e}")
+
+ return pd.DataFrame.from_records(records)
diff --git a/deps/vomp/dataset_toolkits/abo/build_metadata.py b/deps/vomp/dataset_toolkits/abo/build_metadata.py
new file mode 100644
index 0000000000000000000000000000000000000000..5a216667b39b82fc7c02c006776b62dd97bf88a8
--- /dev/null
+++ b/deps/vomp/dataset_toolkits/abo/build_metadata.py
@@ -0,0 +1,108 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import sys
+import argparse
+import pandas as pd
+from easydict import EasyDict as edict
+
+# Add current directory to path to import dataset modules
+sys.path.insert(0, os.path.dirname(__file__))
+
+# Import the local ABO500 module directly
+import ABO500 as dataset_utils
+
+
+def main():
+ parser = argparse.ArgumentParser(description="Build metadata for ABO 500 dataset")
+ parser.add_argument(
+ "--output_dir",
+ type=str,
+ required=True,
+ help="Directory to save the metadata and processed files",
+ )
+
+ # Add dataset-specific arguments
+ dataset_utils.add_args(parser)
+
+ args = parser.parse_args()
+ opt = edict(vars(args))
+
+ # Create output directory
+ os.makedirs(opt.output_dir, exist_ok=True)
+
+ # Get metadata
+ print("Building metadata for ABO 500 dataset...")
+ metadata = dataset_utils.get_metadata(**opt)
+
+ # Add default columns for tracking processing status
+ metadata["rendered"] = False
+ metadata["voxelized"] = False
+ metadata["feature_dinov2_vitl14_reg"] = False
+
+ # Check for existing processed files and update flags
+ for idx, row in metadata.iterrows():
+ sha256 = row["sha256"]
+
+ # Check if voxel file exists
+ voxel_path = os.path.join(opt.output_dir, "voxels", f"{sha256}.ply")
+ if os.path.exists(voxel_path):
+ metadata.at[idx, "voxelized"] = True
+
+ # Check if render file exists (transforms.json)
+ render_path = os.path.join(opt.output_dir, "renders", sha256, "transforms.json")
+ if os.path.exists(render_path):
+ metadata.at[idx, "rendered"] = True
+
+ # Check if feature file exists
+ feature_path = os.path.join(
+ opt.output_dir, "features", "dinov2_vitl14_reg", f"{sha256}.npz"
+ )
+ if os.path.exists(feature_path):
+ metadata.at[idx, "feature_dinov2_vitl14_reg"] = True
+
+ # Save metadata
+ metadata_path = os.path.join(opt.output_dir, "metadata.csv")
+ metadata.to_csv(metadata_path, index=False)
+
+ print(f"Metadata saved to {metadata_path}")
+ print(f"Total objects: {len(metadata)}")
+ print(f"Objects by split:")
+ if "split" in metadata.columns:
+ print(metadata["split"].value_counts())
+
+ # Also save a summary file
+ summary = {
+ "total_objects": len(metadata),
+ "dataset": "ABO500",
+ "splits": (
+ metadata["split"].value_counts().to_dict()
+ if "split" in metadata.columns
+ else {}
+ ),
+ "output_dir": opt.output_dir,
+ }
+
+ import json
+
+ with open(os.path.join(opt.output_dir, "dataset_summary.json"), "w") as f:
+ json.dump(summary, f, indent=2)
+
+ print("Dataset summary saved to dataset_summary.json")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/deps/vomp/dataset_toolkits/abo/extract_feature.py b/deps/vomp/dataset_toolkits/abo/extract_feature.py
new file mode 100644
index 0000000000000000000000000000000000000000..1174c4019e64296fcfd76241d30bf5284d249cec
--- /dev/null
+++ b/deps/vomp/dataset_toolkits/abo/extract_feature.py
@@ -0,0 +1,381 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import copy
+import sys
+import json
+import argparse
+import torch
+import torch.nn.functional as F
+import numpy as np
+import pandas as pd
+import utils3d
+from tqdm import tqdm
+from easydict import EasyDict as edict
+from torchvision import transforms
+from PIL import Image
+
+# Add current directory to path to import dataset modules
+sys.path.insert(0, os.path.dirname(__file__))
+import ABO500 as dataset_utils
+
+torch.set_grad_enabled(False)
+
+
+def get_data(frames, sha256, output_dir):
+ """
+ Load and preprocess rendered images for feature extraction.
+
+ Args:
+ frames (list): List of frame data from transforms.json
+ sha256 (str): SHA256 hash of the object
+ output_dir (str): Output directory containing renders
+
+ Returns:
+ list: List of processed image data
+ """
+ valid_data = []
+
+ for view in frames:
+ image_path = os.path.join(output_dir, "renders", sha256, view["file_path"])
+ try:
+ # Check if file exists before trying to open it
+ if not os.path.exists(image_path):
+ print(f"Warning: Image file {image_path} not found, skipping")
+ continue
+
+ image = Image.open(image_path)
+ except Exception as e:
+ print(f"Error loading image {image_path}: {e}")
+ continue
+
+ try:
+ # Resize and normalize image
+ image = image.resize((518, 518), Image.Resampling.LANCZOS)
+ image = np.array(image).astype(np.float32) / 255
+ image = image[:, :, :3] * image[:, :, 3:] # Apply alpha channel
+ image = torch.from_numpy(image).permute(2, 0, 1).float()
+
+ # Extract camera parameters
+ c2w = torch.tensor(view["transform_matrix"])
+ c2w[:3, 1:3] *= -1
+ extrinsics = torch.inverse(c2w)
+ fov = view["camera_angle_x"]
+ intrinsics = utils3d.torch.intrinsics_from_fov_xy(
+ torch.tensor(fov), torch.tensor(fov)
+ )
+
+ valid_data.append(
+ {"image": image, "extrinsics": extrinsics, "intrinsics": intrinsics}
+ )
+ except Exception as e:
+ print(f"Error processing image {image_path}: {e}")
+ continue
+
+ if len(valid_data) == 0:
+ print(f"Warning: No valid images found for {sha256}")
+ else:
+ print(f"Loaded {len(valid_data)}/{len(frames)} valid images for {sha256}")
+
+ return valid_data
+
+
+def extract_features(
+ file_path,
+ sha256,
+ output_dir=None,
+ model=None,
+ transform=None,
+ batch_size=16,
+ feature_name="dinov2_vitl14_reg",
+):
+ """
+ Extract features for a single object.
+
+ Args:
+ file_path (str): Path to the GLB file (not used directly, but needed for interface)
+ sha256 (str): SHA256 hash of the object
+ output_dir (str): Output directory
+ model: Pre-loaded feature extraction model
+ transform: Image transformation pipeline
+ batch_size (int): Batch size for processing
+ feature_name (str): Name of the feature extraction method
+
+ Returns:
+ dict: Result dictionary with processing info
+ """
+ try:
+ # Load transforms.json
+ transforms_path = os.path.join(output_dir, "renders", sha256, "transforms.json")
+ if not os.path.exists(transforms_path):
+ print(f"transforms.json not found for {sha256}")
+ return {"sha256": sha256, f"feature_{feature_name}": False}
+
+ with open(transforms_path, "r") as f:
+ metadata_json = json.load(f)
+
+ frames = metadata_json["frames"]
+ data = get_data(frames, sha256, output_dir)
+
+ if len(data) == 0:
+ print(f"Skipping {sha256}: no valid image data")
+ return {"sha256": sha256, f"feature_{feature_name}": False}
+
+ # Apply transform to images
+ for datum in data:
+ datum["image"] = transform(datum["image"])
+
+ # Load voxel positions
+ voxel_path = os.path.join(output_dir, "voxels", f"{sha256}.ply")
+ if not os.path.exists(voxel_path):
+ print(f"Voxel file not found for {sha256}")
+ return {"sha256": sha256, f"feature_{feature_name}": False}
+
+ positions = utils3d.io.read_ply(voxel_path)[0]
+ positions = torch.from_numpy(positions).float().cuda()
+ indices = ((positions + 0.5) * 64).long()
+ # Clamp indices to valid range [0, 63] to handle floating point precision issues
+ indices = torch.clamp(indices, 0, 63)
+
+ n_views = len(data)
+ n_patch = 518 // 14
+ pack = {
+ "indices": indices.cpu().numpy().astype(np.uint8),
+ }
+
+ patchtokens_lst = []
+ uv_lst = []
+
+ # Process in batches
+ for i in range(0, n_views, batch_size):
+ batch_data = data[i : i + batch_size]
+ bs = len(batch_data)
+ batch_images = torch.stack([d["image"] for d in batch_data]).cuda()
+ batch_extrinsics = torch.stack([d["extrinsics"] for d in batch_data]).cuda()
+ batch_intrinsics = torch.stack([d["intrinsics"] for d in batch_data]).cuda()
+
+ # Extract features using the model
+ features = model(batch_images, is_training=True)
+
+ # Project 3D positions to 2D
+ uv = (
+ utils3d.torch.project_cv(positions, batch_extrinsics, batch_intrinsics)[
+ 0
+ ]
+ * 2
+ - 1
+ )
+
+ # Extract patch tokens
+ patchtokens = (
+ features["x_prenorm"][:, model.num_register_tokens + 1 :]
+ .permute(0, 2, 1)
+ .reshape(bs, 1024, n_patch, n_patch)
+ )
+ patchtokens_lst.append(patchtokens)
+ uv_lst.append(uv)
+
+ patchtokens = torch.cat(patchtokens_lst, dim=0)
+ uv = torch.cat(uv_lst, dim=0)
+
+ # Sample features at voxel positions
+ pack["patchtokens"] = (
+ F.grid_sample(
+ patchtokens,
+ uv.unsqueeze(1),
+ mode="bilinear",
+ align_corners=False,
+ )
+ .squeeze(2)
+ .permute(0, 2, 1)
+ .cpu()
+ .numpy()
+ )
+ pack["patchtokens"] = np.mean(pack["patchtokens"], axis=0).astype(np.float16)
+
+ # Save features
+ save_path = os.path.join(output_dir, "features", feature_name, f"{sha256}.npz")
+ os.makedirs(os.path.dirname(save_path), exist_ok=True)
+ np.savez_compressed(save_path, **pack)
+
+ return {"sha256": sha256, f"feature_{feature_name}": True}
+
+ except Exception as e:
+ print(f"Error processing {sha256}: {e}")
+ import traceback
+
+ traceback.print_exc()
+ return {"sha256": sha256, f"feature_{feature_name}": False}
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(description="Extract features for ABO 500 dataset")
+ parser.add_argument(
+ "--output_dir",
+ type=str,
+ required=True,
+ help="Directory containing metadata and where to save features",
+ )
+ parser.add_argument(
+ "--model",
+ type=str,
+ default="dinov2_vitl14_reg",
+ help="Feature extraction model",
+ )
+ parser.add_argument(
+ "--instances",
+ type=str,
+ default=None,
+ help="Specific instances to process (comma-separated or file path)",
+ )
+ parser.add_argument("--batch_size", type=int, default=16)
+ parser.add_argument("--rank", type=int, default=0)
+ parser.add_argument("--world_size", type=int, default=1)
+ parser.add_argument(
+ "--force",
+ action="store_true",
+ help="Force feature extraction even if already processed",
+ )
+ parser.add_argument(
+ "--limit", type=int, default=None, help="Process only the first N objects"
+ )
+
+ args = parser.parse_args()
+ opt = edict(vars(args))
+
+ feature_name = opt.model
+
+ # Create features directory
+ os.makedirs(os.path.join(opt.output_dir, "features", feature_name), exist_ok=True)
+
+ # Load model
+ print(f"Loading model: {opt.model}")
+ dinov2_model = torch.hub.load("facebookresearch/dinov2", opt.model)
+ dinov2_model.eval().cuda()
+ transform = transforms.Compose(
+ [
+ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+ ]
+ )
+
+ # Load metadata
+ metadata_path = os.path.join(opt.output_dir, "metadata.csv")
+ if not os.path.exists(metadata_path):
+ raise ValueError(f"metadata.csv not found at {metadata_path}")
+
+ metadata = pd.read_csv(metadata_path)
+
+ # Filter instances if specified
+ if opt.instances is not None:
+ if os.path.exists(opt.instances):
+ with open(opt.instances, "r") as f:
+ instances = f.read().splitlines()
+ else:
+ instances = opt.instances.split(",")
+ metadata = metadata[metadata["sha256"].isin(instances)]
+ else:
+ # Only process objects that have been rendered and voxelized
+ if "rendered" in metadata.columns:
+ metadata = metadata[metadata["rendered"] == True]
+ if "voxelized" in metadata.columns:
+ metadata = metadata[metadata["voxelized"] == True]
+
+ # Only process objects that haven't had features extracted yet
+ if f"feature_{feature_name}" in metadata.columns and not opt.force:
+ metadata = metadata[metadata[f"feature_{feature_name}"] == False]
+
+ # Apply distributed processing
+ start = len(metadata) * opt.rank // opt.world_size
+ end = len(metadata) * (opt.rank + 1) // opt.world_size
+ metadata = metadata[start:end]
+
+ # Apply limit if specified
+ if opt.limit is not None:
+ metadata = metadata.head(opt.limit)
+
+ print(f"Processing {len(metadata)} objects...")
+
+ # Track already processed objects
+ records = []
+ sha256s = list(metadata["sha256"].values)
+
+ # Filter out objects that are already processed
+ if not opt.force:
+ for sha256 in copy.copy(sha256s):
+ feature_path = os.path.join(
+ opt.output_dir, "features", feature_name, f"{sha256}.npz"
+ )
+ if os.path.exists(feature_path):
+ records.append({"sha256": sha256, f"feature_{feature_name}": True})
+ sha256s.remove(sha256)
+
+ # Filter out objects that don't have required prerequisite files
+ initial_count = len(sha256s)
+ filtered_sha256s = []
+
+ for sha256 in sha256s:
+ # Check for voxel file
+ voxel_path = os.path.join(opt.output_dir, "voxels", f"{sha256}.ply")
+ if not os.path.exists(voxel_path):
+ print(f"Skipping {sha256}: voxel file not found")
+ continue
+
+ # Check for transforms.json
+ transforms_path = os.path.join(
+ opt.output_dir, "renders", sha256, "transforms.json"
+ )
+ if not os.path.exists(transforms_path):
+ print(f"Skipping {sha256}: transforms.json not found")
+ continue
+
+ filtered_sha256s.append(sha256)
+
+ sha256s = filtered_sha256s
+ print(
+ f"Filtered from {initial_count} to {len(sha256s)} objects with required files"
+ )
+
+ # Extract features for remaining objects
+ if len(sha256s) > 0:
+ for sha256 in tqdm(sha256s, desc="Extracting features"):
+ # Get the file path (not used directly but needed for interface consistency)
+ file_path = metadata[metadata["sha256"] == sha256]["local_path"].iloc[0]
+
+ result = extract_features(
+ file_path=file_path,
+ sha256=sha256,
+ output_dir=opt.output_dir,
+ model=dinov2_model,
+ transform=transform,
+ batch_size=opt.batch_size,
+ feature_name=feature_name,
+ )
+
+ if result is not None:
+ records.append(result)
+
+ # Save results
+ if len(records) > 0:
+ results_df = pd.DataFrame.from_records(records)
+ results_df.to_csv(
+ os.path.join(opt.output_dir, f"feature_{feature_name}_{opt.rank}.csv"),
+ index=False,
+ )
+ print(
+ f"Feature extraction complete. Results saved to feature_{feature_name}_{opt.rank}.csv"
+ )
+ else:
+ print("No objects processed.")
diff --git a/deps/vomp/dataset_toolkits/abo/render.py b/deps/vomp/dataset_toolkits/abo/render.py
new file mode 100644
index 0000000000000000000000000000000000000000..c3f935e3bbd1325983a8ba8ce68195225991d4d0
--- /dev/null
+++ b/deps/vomp/dataset_toolkits/abo/render.py
@@ -0,0 +1,241 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import json
+import copy
+import sys
+import argparse
+import pandas as pd
+from easydict import EasyDict as edict
+from functools import partial
+from subprocess import DEVNULL, call
+import numpy as np
+
+# Add current directory to path to import dataset modules
+sys.path.insert(0, os.path.dirname(__file__))
+import ABO500 as dataset_utils
+
+# Import from the existing render.py utils
+sys.path.append(os.path.dirname(os.path.dirname(__file__)))
+from utils import sphere_hammersley_sequence
+
+BLENDER_LINK = (
+ "https://download.blender.org/release/Blender3.0/blender-3.0.1-linux-x64.tar.xz"
+)
+BLENDER_INSTALLATION_PATH = "/tmp"
+BLENDER_PATH = f"{BLENDER_INSTALLATION_PATH}/blender-3.0.1-linux-x64/blender"
+
+
+def _install_blender():
+ """Install Blender if not already installed."""
+ if not os.path.exists(BLENDER_PATH):
+ os.system("sudo apt-get update")
+ os.system(
+ "sudo apt-get install -y libxrender1 libxi6 libxkbcommon-x11-0 libsm6"
+ )
+ os.system(f"wget {BLENDER_LINK} -P {BLENDER_INSTALLATION_PATH}")
+ os.system(
+ f"tar -xvf {BLENDER_INSTALLATION_PATH}/blender-3.0.1-linux-x64.tar.xz -C {BLENDER_INSTALLATION_PATH}"
+ )
+
+
+def _render_glb(file_path, sha256, output_dir, num_views):
+ """
+ Render a GLB file from multiple viewpoints.
+
+ Args:
+ file_path (str): Path to the GLB file
+ sha256 (str): SHA256 hash of the file
+ output_dir (str): Directory to save renders
+ num_views (int): Number of viewpoints to render
+
+ Returns:
+ dict: Result dictionary with rendering info
+ """
+ # Convert to absolute path to avoid issues with relative paths
+ output_dir = os.path.abspath(output_dir)
+ output_folder = os.path.join(output_dir, "renders", sha256)
+
+ # Build camera parameters {yaw, pitch, radius, fov}
+ yaws = []
+ pitchs = []
+ offset = (np.random.rand(), np.random.rand())
+ for i in range(num_views):
+ y, p = sphere_hammersley_sequence(i, num_views, offset)
+ yaws.append(y)
+ pitchs.append(p)
+
+ radius = [2] * num_views
+ fov = [40 / 180 * np.pi] * num_views
+ views = [
+ {"yaw": y, "pitch": p, "radius": r, "fov": f}
+ for y, p, r, f in zip(yaws, pitchs, radius, fov)
+ ]
+
+ # Construct Blender command
+ blender_script_path = os.path.join(
+ os.path.dirname(os.path.dirname(__file__)), "blender_script", "render.py"
+ )
+
+ args = [
+ BLENDER_PATH,
+ "-b",
+ "-P",
+ blender_script_path,
+ "--",
+ "--views",
+ json.dumps(views),
+ "--object",
+ os.path.expanduser(file_path),
+ "--resolution",
+ "512",
+ "--output_folder",
+ output_folder,
+ "--engine",
+ "CYCLES",
+ "--save_mesh",
+ ]
+
+ try:
+ # Execute Blender rendering
+ result = call(args, stdout=DEVNULL, stderr=DEVNULL)
+
+ # Check if rendering was successful
+ if result == 0 and os.path.exists(
+ os.path.join(output_folder, "transforms.json")
+ ):
+ return {"sha256": sha256, "rendered": True}
+ else:
+ print(f"Rendering failed for {sha256}")
+ return {"sha256": sha256, "rendered": False}
+
+ except Exception as e:
+ print(f"Error rendering {file_path}: {e}")
+ return {"sha256": sha256, "rendered": False}
+
+
+def _render(file_path, sha256, output_dir=None, num_views=150):
+ """Wrapper function for rendering."""
+ return _render_glb(file_path, sha256, output_dir, num_views)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(description="Render ABO 500 dataset")
+ parser.add_argument(
+ "--output_dir",
+ type=str,
+ required=True,
+ help="Directory containing metadata and where to save renders",
+ )
+ parser.add_argument(
+ "--instances",
+ type=str,
+ default=None,
+ help="Specific instances to process (comma-separated or file path)",
+ )
+ parser.add_argument(
+ "--num_views", type=int, default=150, help="Number of views to render"
+ )
+ parser.add_argument(
+ "--force", action="store_true", help="Force rendering even if already processed"
+ )
+ parser.add_argument("--rank", type=int, default=0)
+ parser.add_argument("--world_size", type=int, default=1)
+ parser.add_argument("--max_workers", type=int, default=8)
+ parser.add_argument(
+ "--limit", type=int, default=None, help="Process only the first N objects"
+ )
+
+ args = parser.parse_args()
+ opt = edict(vars(args))
+
+ # Create renders directory
+ os.makedirs(os.path.join(opt.output_dir, "renders"), exist_ok=True)
+
+ # Install Blender if needed
+ print("Checking Blender installation...", flush=True)
+ _install_blender()
+
+ # Load metadata
+ metadata_path = os.path.join(opt.output_dir, "metadata.csv")
+ if not os.path.exists(metadata_path):
+ raise ValueError(f"metadata.csv not found at {metadata_path}")
+
+ metadata = pd.read_csv(metadata_path)
+
+ # Filter instances if specified
+ if opt.instances is not None:
+ if os.path.exists(opt.instances):
+ with open(opt.instances, "r") as f:
+ instances = f.read().splitlines()
+ else:
+ instances = opt.instances.split(",")
+ metadata = metadata[metadata["sha256"].isin(instances)]
+ else:
+ # Only process objects that have valid local paths
+ metadata = metadata[metadata["local_path"].notna()]
+
+ # Only process objects that haven't been rendered yet
+ if "rendered" in metadata.columns and not opt.force:
+ metadata = metadata[metadata["rendered"] == False]
+
+ # Apply distributed processing
+ start = len(metadata) * opt.rank // opt.world_size
+ end = len(metadata) * (opt.rank + 1) // opt.world_size
+ metadata = metadata[start:end]
+
+ # Apply limit if specified
+ if opt.limit is not None:
+ metadata = metadata.head(opt.limit)
+
+ print(f"Processing {len(metadata)} objects...")
+
+ # Track already processed objects
+ records = []
+
+ # Filter out objects that are already processed
+ if not opt.force:
+ for sha256 in copy.copy(metadata["sha256"].values):
+ transforms_path = os.path.join(
+ opt.output_dir, "renders", sha256, "transforms.json"
+ )
+ if os.path.exists(transforms_path):
+ records.append({"sha256": sha256, "rendered": True})
+ metadata = metadata[metadata["sha256"] != sha256]
+
+ # Process remaining objects
+ if len(metadata) > 0:
+ func = partial(_render, output_dir=opt.output_dir, num_views=opt.num_views)
+ rendered = dataset_utils.foreach_instance(
+ metadata,
+ opt.output_dir,
+ func,
+ max_workers=opt.max_workers,
+ desc="Rendering objects",
+ )
+
+ # Combine results
+ if len(records) > 0:
+ rendered = pd.concat([rendered, pd.DataFrame.from_records(records)])
+
+ # Save results
+ rendered.to_csv(
+ os.path.join(opt.output_dir, f"rendered_{opt.rank}.csv"), index=False
+ )
+
+ print(f"Rendering complete. Results saved to rendered_{opt.rank}.csv")
+ else:
+ print("No objects to process.")
diff --git a/deps/vomp/dataset_toolkits/abo/voxelize.py b/deps/vomp/dataset_toolkits/abo/voxelize.py
new file mode 100644
index 0000000000000000000000000000000000000000..d3ae49d1ca32440cd7cf78916d97d657c0f62c8e
--- /dev/null
+++ b/deps/vomp/dataset_toolkits/abo/voxelize.py
@@ -0,0 +1,306 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import copy
+import sys
+import argparse
+import pandas as pd
+from easydict import EasyDict as edict
+from functools import partial
+import numpy as np
+import open3d as o3d
+import utils3d
+import trimesh
+import tempfile
+import shutil
+
+# Add current directory to path to import dataset modules
+sys.path.insert(0, os.path.dirname(__file__))
+
+import ABO500 as dataset_utils
+
+
+def voxelize_mesh(
+ vertices, faces, voxel_size=1 / 64, center_scale=None, max_voxels=None
+):
+ """
+ Voxelize a mesh represented by vertices and faces using volumetric voxelization.
+
+ Args:
+ vertices (numpy.ndarray): Array of vertices
+ faces (numpy.ndarray): Array of faces
+ voxel_size (float): Size of each voxel
+ center_scale (tuple): Optional center and scale for normalization
+ max_voxels (int): Maximum number of voxels to return (will subsample if exceeded)
+
+ Returns:
+ tuple: (voxel_centers, voxel_grid) - center coordinates of voxels and Trimesh voxel grid
+ """
+ # Create a Trimesh mesh
+ mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
+
+ # Normalize the mesh to [-0.5, 0.5] range
+ vertices = mesh.vertices.copy()
+
+ if center_scale is None:
+ vertices_min = np.min(vertices, axis=0)
+ vertices_max = np.max(vertices, axis=0)
+ center = (vertices_min + vertices_max) / 2
+ scale = np.max(vertices_max - vertices_min)
+ else:
+ center, scale = center_scale
+
+ vertices = (vertices - center) / scale
+ vertices = np.clip(vertices, -0.5 + 1e-6, 0.5 - 1e-6)
+
+ # Update mesh with normalized vertices
+ mesh.vertices = vertices
+
+ # Create volumetric voxel grid using Trimesh
+ voxel_grid = mesh.voxelized(pitch=voxel_size).fill()
+
+ # Get voxel centers from the filled voxel grid
+ voxel_centers = voxel_grid.points
+
+ # Subsample if we have too many voxels
+ if max_voxels is not None and len(voxel_centers) > max_voxels:
+ print(f"Subsampling voxels: {len(voxel_centers):,} -> {max_voxels:,}")
+ # Use random sampling to maintain spatial distribution
+ np.random.seed(42) # For reproducibility
+ indices = np.random.choice(len(voxel_centers), max_voxels, replace=False)
+ voxel_centers = voxel_centers[indices]
+
+ return voxel_centers, voxel_grid
+
+
+def load_glb_mesh(glb_path):
+ """
+ Load a GLB file and extract mesh data.
+
+ Args:
+ glb_path (str): Path to the GLB file
+
+ Returns:
+ tuple: (vertices, faces) - mesh vertices and faces
+ """
+ try:
+ # Load the GLB file using trimesh
+ mesh = trimesh.load(glb_path)
+
+ # Handle different mesh types
+ if isinstance(mesh, trimesh.Scene):
+ # If it's a scene, combine all meshes
+ combined_mesh = trimesh.util.concatenate(
+ [
+ geometry
+ for geometry in mesh.geometry.values()
+ if isinstance(geometry, trimesh.Trimesh)
+ ]
+ )
+ if combined_mesh is None:
+ raise ValueError("No valid meshes found in GLB file")
+ mesh = combined_mesh
+ elif not isinstance(mesh, trimesh.Trimesh):
+ raise ValueError("GLB file does not contain a valid mesh")
+
+ # Ensure the mesh has faces
+ if len(mesh.faces) == 0:
+ raise ValueError("Mesh has no faces")
+
+ return mesh.vertices, mesh.faces
+
+ except Exception as e:
+ print(f"Error loading GLB file {glb_path}: {e}")
+ return None, None
+
+
+def voxelize_glb(glb_path, sha256, output_dir, max_voxels=None):
+ """
+ Voxelize a GLB file and save the result.
+
+ Args:
+ glb_path (str): Path to the GLB file
+ sha256 (str): SHA256 hash of the file
+ output_dir (str): Directory to save the voxelized data
+ max_voxels (int): Maximum number of voxels to generate
+
+ Returns:
+ dict: Result dictionary with processing info
+ """
+ try:
+ # Load the GLB mesh
+ vertices, faces = load_glb_mesh(glb_path)
+
+ if vertices is None or faces is None:
+ print(f"Failed to load mesh from {glb_path}")
+ return {"sha256": sha256, "voxelized": False, "num_voxels": 0}
+
+ print(f"Loaded mesh with {len(vertices)} vertices and {len(faces)} faces")
+
+ # Voxelize the mesh
+ voxel_centers, voxel_grid = voxelize_mesh(
+ vertices, faces, max_voxels=max_voxels
+ )
+
+ if len(voxel_centers) == 0:
+ print(f"No voxels generated for {sha256}")
+ return {"sha256": sha256, "voxelized": False, "num_voxels": 0}
+
+ # Save voxel centers as PLY file
+ ply_output_path = os.path.join(output_dir, "voxels", f"{sha256}.ply")
+ save_ply(ply_output_path, voxel_centers)
+
+ print(f"Voxelized {sha256}: {len(voxel_centers)} voxels")
+
+ return {"sha256": sha256, "voxelized": True, "num_voxels": len(voxel_centers)}
+
+ except Exception as e:
+ print(f"Error voxelizing {glb_path}: {e}")
+ import traceback
+
+ traceback.print_exc()
+ return {"sha256": sha256, "voxelized": False, "num_voxels": 0}
+
+
+def save_ply(filename, points):
+ """
+ Save points as a PLY file.
+
+ Args:
+ filename (str): Output filename
+ points (numpy.ndarray): Array of 3D points
+ """
+ os.makedirs(os.path.dirname(filename), exist_ok=True)
+ pcd = o3d.geometry.PointCloud()
+ pcd.points = o3d.utility.Vector3dVector(points)
+ o3d.io.write_point_cloud(filename, pcd)
+
+
+def _voxelize(file_path, sha256, output_dir=None, max_voxels=None):
+ """Wrapper function for voxelization."""
+ return voxelize_glb(file_path, sha256, output_dir, max_voxels=max_voxels)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(description="Voxelize ABO 500 dataset")
+ parser.add_argument(
+ "--output_dir",
+ type=str,
+ required=True,
+ help="Directory containing metadata and where to save voxelized data",
+ )
+ parser.add_argument(
+ "--instances",
+ type=str,
+ default=None,
+ help="Specific instances to process (comma-separated or file path)",
+ )
+ parser.add_argument(
+ "--force",
+ action="store_true",
+ help="Force voxelization even if already processed",
+ )
+ parser.add_argument("--rank", type=int, default=0)
+ parser.add_argument("--world_size", type=int, default=1)
+ parser.add_argument("--max_workers", type=int, default=None)
+ parser.add_argument(
+ "--limit", type=int, default=None, help="Process only the first N objects"
+ )
+ parser.add_argument(
+ "--max_voxels",
+ type=int,
+ default=70000,
+ help="Maximum number of voxels per asset",
+ )
+
+ args = parser.parse_args()
+ opt = edict(vars(args))
+
+ # Create voxels directory
+ os.makedirs(os.path.join(opt.output_dir, "voxels"), exist_ok=True)
+
+ # Load metadata
+ metadata_path = os.path.join(opt.output_dir, "metadata.csv")
+ if not os.path.exists(metadata_path):
+ raise ValueError(f"metadata.csv not found at {metadata_path}")
+
+ metadata = pd.read_csv(metadata_path)
+
+ # Filter instances if specified
+ if opt.instances is not None:
+ if os.path.exists(opt.instances):
+ with open(opt.instances, "r") as f:
+ instances = f.read().splitlines()
+ else:
+ instances = opt.instances.split(",")
+ metadata = metadata[metadata["sha256"].isin(instances)]
+ else:
+ # Only process objects that haven't been voxelized yet
+ if "voxelized" in metadata.columns and not opt.force:
+ metadata = metadata[metadata["voxelized"] == False]
+
+ # Apply distributed processing
+ start = len(metadata) * opt.rank // opt.world_size
+ end = len(metadata) * (opt.rank + 1) // opt.world_size
+ metadata = metadata[start:end]
+
+ # Apply limit if specified
+ if opt.limit is not None:
+ metadata = metadata.head(opt.limit)
+
+ print(f"Processing {len(metadata)} objects with max_voxels={opt.max_voxels:,}...")
+
+ # Track already processed objects
+ records = []
+
+ # Filter out objects that are already processed
+ if not opt.force:
+ for sha256 in copy.copy(metadata["sha256"].values):
+ ply_path = os.path.join(opt.output_dir, "voxels", f"{sha256}.ply")
+ if os.path.exists(ply_path):
+ try:
+ pts = utils3d.io.read_ply(ply_path)[0]
+ records.append(
+ {"sha256": sha256, "voxelized": True, "num_voxels": len(pts)}
+ )
+ metadata = metadata[metadata["sha256"] != sha256]
+ except:
+ # If file is corrupted, re-process it
+ pass
+
+ # Process remaining objects
+ if len(metadata) > 0:
+ func = partial(_voxelize, output_dir=opt.output_dir, max_voxels=opt.max_voxels)
+ voxelized = dataset_utils.foreach_instance(
+ metadata,
+ opt.output_dir,
+ func,
+ max_workers=opt.max_workers,
+ desc="Voxelizing",
+ )
+
+ # Combine results
+ if len(records) > 0:
+ voxelized = pd.concat([voxelized, pd.DataFrame.from_records(records)])
+
+ # Save results
+ voxelized.to_csv(
+ os.path.join(opt.output_dir, f"voxelized_{opt.rank}.csv"), index=False
+ )
+
+ print(f"Voxelization complete. Results saved to voxelized_{opt.rank}.csv")
+ else:
+ print("No objects to process.")
diff --git a/deps/vomp/dataset_toolkits/blender_script/render.py b/deps/vomp/dataset_toolkits/blender_script/render.py
new file mode 100644
index 0000000000000000000000000000000000000000..2146d5c4eb87ed6c4056a7cd4e1428b38ed8e257
--- /dev/null
+++ b/deps/vomp/dataset_toolkits/blender_script/render.py
@@ -0,0 +1,695 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse, sys, os, math, re, glob
+from typing import *
+import bpy
+from mathutils import Vector, Matrix
+import numpy as np
+import json
+import glob
+
+"""=============== BLENDER ==============="""
+
+IMPORT_FUNCTIONS: Dict[str, Callable] = {
+ "obj": bpy.ops.import_scene.obj,
+ "glb": bpy.ops.import_scene.gltf,
+ "gltf": bpy.ops.import_scene.gltf,
+ "usd": bpy.ops.import_scene.usd,
+ "fbx": bpy.ops.import_scene.fbx,
+ "stl": bpy.ops.import_mesh.stl,
+ "usda": bpy.ops.import_scene.usda,
+ "dae": bpy.ops.wm.collada_import,
+ "ply": bpy.ops.import_mesh.ply,
+ "abc": bpy.ops.wm.alembic_import,
+ "blend": bpy.ops.wm.append,
+}
+
+EXT = {
+ "PNG": "png",
+ "JPEG": "jpg",
+ "OPEN_EXR": "exr",
+ "TIFF": "tiff",
+ "BMP": "bmp",
+ "HDR": "hdr",
+ "TARGA": "tga",
+}
+
+
+def setup_gpu_devices():
+ """Setup GPU devices for Cycles rendering using the same approach as viz_fields.py."""
+ try:
+ cycles_prefs = bpy.context.preferences.addons["cycles"].preferences
+ except KeyError:
+ print("[ERROR] Cycles addon not found or not enabled")
+ return False
+
+ # Get available device types and try them in order of preference
+ available_types = cycles_prefs.get_device_types(bpy.context)
+ available_type_names = [dt[0] for dt in available_types]
+
+ # Try device types in order of preference (same as viz_fields.py)
+ preferred_types = ["OPTIX", "CUDA", "HIP", "ONEAPI", "OPENCL", "METAL"]
+ available_preferred = [t for t in preferred_types if t in available_type_names]
+
+ best_type = None
+ for device_type in available_preferred:
+ try:
+ cycles_prefs.compute_device_type = device_type
+ cycles_prefs.get_devices()
+
+ gpu_devices = [
+ dev for dev in cycles_prefs.devices if dev.type == device_type
+ ]
+
+ if gpu_devices:
+ best_type = device_type
+ print(f"[INFO] โ Using compute device type: {device_type}")
+ break
+
+ except Exception as e:
+ print(f"[WARNING] โ Could not set {device_type}: {e}")
+ continue
+
+ if not best_type:
+ print("[ERROR] โ No GPU compute device type available")
+ return False
+
+ # Enable all GPU devices for the selected type
+ gpu_devices_enabled = 0
+ for device in cycles_prefs.devices:
+ if device.type == best_type:
+ device.use = True
+ gpu_devices_enabled += 1
+ print(f"[INFO] โ Enabled GPU device: {device.name}")
+
+ print(
+ f"[INFO] GPU setup complete: {gpu_devices_enabled} device(s) enabled with {best_type}"
+ )
+ return gpu_devices_enabled > 0
+
+
+def init_render(
+ engine="CYCLES", resolution=512, geo_mode=False, use_gpu=True, gpu_device="OPTIX"
+):
+ bpy.context.scene.render.engine = engine
+ bpy.context.scene.render.resolution_x = resolution
+ bpy.context.scene.render.resolution_y = resolution
+ bpy.context.scene.render.resolution_percentage = 100
+ bpy.context.scene.render.image_settings.file_format = "PNG"
+ bpy.context.scene.render.image_settings.color_mode = "RGBA"
+ bpy.context.scene.render.film_transparent = True
+
+ # Enhanced GPU setup using the same approach as viz_fields.py
+ if use_gpu:
+ gpu_success = setup_gpu_devices()
+ if gpu_success:
+ bpy.context.scene.cycles.device = "GPU"
+ print("[INFO] โ
GPU rendering enabled")
+ else:
+ bpy.context.scene.cycles.device = "CPU"
+ print("[WARNING] โ GPU setup failed, using CPU rendering")
+ else:
+ bpy.context.scene.cycles.device = "CPU"
+ print("[INFO] CPU rendering requested")
+
+ bpy.context.scene.cycles.samples = 128 if not geo_mode else 1
+ bpy.context.scene.cycles.filter_type = "BOX"
+ bpy.context.scene.cycles.filter_width = 1
+ bpy.context.scene.cycles.diffuse_bounces = 1
+ bpy.context.scene.cycles.glossy_bounces = 1
+ bpy.context.scene.cycles.transparent_max_bounces = 3 if not geo_mode else 0
+ bpy.context.scene.cycles.transmission_bounces = 3 if not geo_mode else 1
+ bpy.context.scene.cycles.use_denoising = True
+
+
+def init_nodes(save_depth=False, save_normal=False, save_albedo=False, save_mist=False):
+ if not any([save_depth, save_normal, save_albedo, save_mist]):
+ return {}, {}
+ outputs = {}
+ spec_nodes = {}
+
+ bpy.context.scene.use_nodes = True
+ bpy.context.scene.view_layers["View Layer"].use_pass_z = save_depth
+ bpy.context.scene.view_layers["View Layer"].use_pass_normal = save_normal
+ bpy.context.scene.view_layers["View Layer"].use_pass_diffuse_color = save_albedo
+ bpy.context.scene.view_layers["View Layer"].use_pass_mist = save_mist
+
+ nodes = bpy.context.scene.node_tree.nodes
+ links = bpy.context.scene.node_tree.links
+ for n in nodes:
+ nodes.remove(n)
+
+ render_layers = nodes.new("CompositorNodeRLayers")
+
+ if save_depth:
+ depth_file_output = nodes.new("CompositorNodeOutputFile")
+ depth_file_output.base_path = ""
+ depth_file_output.file_slots[0].use_node_format = True
+ depth_file_output.format.file_format = "PNG"
+ depth_file_output.format.color_depth = "16"
+ depth_file_output.format.color_mode = "BW"
+ # Remap to 0-1
+ map = nodes.new(type="CompositorNodeMapRange")
+ map.inputs[1].default_value = 0 # (min value you will be getting)
+ map.inputs[2].default_value = 10 # (max value you will be getting)
+ map.inputs[3].default_value = 0 # (min value you will map to)
+ map.inputs[4].default_value = 1 # (max value you will map to)
+
+ links.new(render_layers.outputs["Depth"], map.inputs[0])
+ links.new(map.outputs[0], depth_file_output.inputs[0])
+
+ outputs["depth"] = depth_file_output
+ spec_nodes["depth_map"] = map
+
+ if save_normal:
+ normal_file_output = nodes.new("CompositorNodeOutputFile")
+ normal_file_output.base_path = ""
+ normal_file_output.file_slots[0].use_node_format = True
+ normal_file_output.format.file_format = "OPEN_EXR"
+ normal_file_output.format.color_mode = "RGB"
+ normal_file_output.format.color_depth = "16"
+
+ links.new(render_layers.outputs["Normal"], normal_file_output.inputs[0])
+
+ outputs["normal"] = normal_file_output
+
+ if save_albedo:
+ albedo_file_output = nodes.new("CompositorNodeOutputFile")
+ albedo_file_output.base_path = ""
+ albedo_file_output.file_slots[0].use_node_format = True
+ albedo_file_output.format.file_format = "PNG"
+ albedo_file_output.format.color_mode = "RGBA"
+ albedo_file_output.format.color_depth = "8"
+
+ alpha_albedo = nodes.new("CompositorNodeSetAlpha")
+
+ links.new(render_layers.outputs["DiffCol"], alpha_albedo.inputs["Image"])
+ links.new(render_layers.outputs["Alpha"], alpha_albedo.inputs["Alpha"])
+ links.new(alpha_albedo.outputs["Image"], albedo_file_output.inputs[0])
+
+ outputs["albedo"] = albedo_file_output
+
+ if save_mist:
+ bpy.data.worlds["World"].mist_settings.start = 0
+ bpy.data.worlds["World"].mist_settings.depth = 10
+
+ mist_file_output = nodes.new("CompositorNodeOutputFile")
+ mist_file_output.base_path = ""
+ mist_file_output.file_slots[0].use_node_format = True
+ mist_file_output.format.file_format = "PNG"
+ mist_file_output.format.color_mode = "BW"
+ mist_file_output.format.color_depth = "16"
+
+ links.new(render_layers.outputs["Mist"], mist_file_output.inputs[0])
+
+ outputs["mist"] = mist_file_output
+
+ return outputs, spec_nodes
+
+
+def init_scene() -> None:
+ """Resets the scene to a clean state.
+
+ Returns:
+ None
+ """
+ # delete everything
+ for obj in bpy.data.objects:
+ bpy.data.objects.remove(obj, do_unlink=True)
+
+ # delete all the materials
+ for material in bpy.data.materials:
+ bpy.data.materials.remove(material, do_unlink=True)
+
+ # delete all the textures
+ for texture in bpy.data.textures:
+ bpy.data.textures.remove(texture, do_unlink=True)
+
+ # delete all the images
+ for image in bpy.data.images:
+ bpy.data.images.remove(image, do_unlink=True)
+
+
+def init_camera():
+ cam = bpy.data.objects.new("Camera", bpy.data.cameras.new("Camera"))
+ bpy.context.collection.objects.link(cam)
+ bpy.context.scene.camera = cam
+ cam.data.sensor_height = cam.data.sensor_width = 32
+ cam_constraint = cam.constraints.new(type="TRACK_TO")
+ cam_constraint.track_axis = "TRACK_NEGATIVE_Z"
+ cam_constraint.up_axis = "UP_Y"
+ cam_empty = bpy.data.objects.new("Empty", None)
+ cam_empty.location = (0, 0, 0)
+ bpy.context.scene.collection.objects.link(cam_empty)
+ cam_constraint.target = cam_empty
+ return cam
+
+
+def init_lighting():
+ # Clear existing lights
+ bpy.ops.object.select_all(action="DESELECT")
+ bpy.ops.object.select_by_type(type="LIGHT")
+ bpy.ops.object.delete()
+
+ # Create key light
+ default_light = bpy.data.objects.new(
+ "Default_Light", bpy.data.lights.new("Default_Light", type="POINT")
+ )
+ bpy.context.collection.objects.link(default_light)
+ default_light.data.energy = 1000
+ default_light.location = (4, 1, 6)
+ default_light.rotation_euler = (0, 0, 0)
+
+ # create top light
+ top_light = bpy.data.objects.new(
+ "Top_Light", bpy.data.lights.new("Top_Light", type="AREA")
+ )
+ bpy.context.collection.objects.link(top_light)
+ top_light.data.energy = 10000
+ top_light.location = (0, 0, 10)
+ top_light.scale = (100, 100, 100)
+
+ # create bottom light
+ bottom_light = bpy.data.objects.new(
+ "Bottom_Light", bpy.data.lights.new("Bottom_Light", type="AREA")
+ )
+ bpy.context.collection.objects.link(bottom_light)
+ bottom_light.data.energy = 1000
+ bottom_light.location = (0, 0, -10)
+ bottom_light.rotation_euler = (0, 0, 0)
+
+ return {
+ "default_light": default_light,
+ "top_light": top_light,
+ "bottom_light": bottom_light,
+ }
+
+
+def load_object(object_path: str) -> None:
+ """Loads a model with a supported file extension into the scene.
+
+ Args:
+ object_path (str): Path to the model file.
+
+ Raises:
+ ValueError: If the file extension is not supported.
+
+ Returns:
+ None
+ """
+ file_extension = object_path.split(".")[-1].lower()
+ if file_extension is None:
+ raise ValueError(f"Unsupported file type: {object_path}")
+
+ if file_extension == "usdz":
+ # install usdz io package
+ dirname = os.path.dirname(os.path.realpath(__file__))
+ usdz_package = os.path.join(dirname, "io_scene_usdz.zip")
+ bpy.ops.preferences.addon_install(filepath=usdz_package)
+ # enable it
+ addon_name = "io_scene_usdz"
+ bpy.ops.preferences.addon_enable(module=addon_name)
+ # import the usdz
+ from io_scene_usdz.import_usdz import import_usdz
+
+ import_usdz(context, filepath=object_path, materials=True, animations=True)
+ return None
+
+ # load from existing import functions
+ import_function = IMPORT_FUNCTIONS[file_extension]
+
+ print(f"Loading object from {object_path}")
+ if file_extension == "blend":
+ import_function(directory=object_path, link=False)
+ elif file_extension in {"glb", "gltf"}:
+ import_function(
+ filepath=object_path, merge_vertices=True, import_shading="NORMALS"
+ )
+ else:
+ import_function(filepath=object_path)
+
+
+def delete_invisible_objects() -> None:
+ """Deletes all invisible objects in the scene.
+
+ Returns:
+ None
+ """
+ # bpy.ops.object.mode_set(mode="OBJECT")
+ bpy.ops.object.select_all(action="DESELECT")
+ for obj in bpy.context.scene.objects:
+ if obj.hide_viewport or obj.hide_render:
+ obj.hide_viewport = False
+ obj.hide_render = False
+ obj.hide_select = False
+ obj.select_set(True)
+ bpy.ops.object.delete()
+
+ # Delete invisible collections
+ invisible_collections = [col for col in bpy.data.collections if col.hide_viewport]
+ for col in invisible_collections:
+ bpy.data.collections.remove(col)
+
+
+def split_mesh_normal():
+ bpy.ops.object.select_all(action="DESELECT")
+ objs = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"]
+ bpy.context.view_layer.objects.active = objs[0]
+ for obj in objs:
+ obj.select_set(True)
+ bpy.ops.object.mode_set(mode="EDIT")
+ bpy.ops.mesh.select_all(action="SELECT")
+ bpy.ops.mesh.split_normals()
+ bpy.ops.object.mode_set(mode="OBJECT")
+ bpy.ops.object.select_all(action="DESELECT")
+
+
+def delete_custom_normals():
+ for this_obj in bpy.data.objects:
+ if this_obj.type == "MESH":
+ bpy.context.view_layer.objects.active = this_obj
+ bpy.ops.mesh.customdata_custom_splitnormals_clear()
+
+
+def override_material():
+ new_mat = bpy.data.materials.new(name="Override0123456789")
+ new_mat.use_nodes = True
+ new_mat.node_tree.nodes.clear()
+ bsdf = new_mat.node_tree.nodes.new("ShaderNodeBsdfDiffuse")
+ bsdf.inputs[0].default_value = (0.5, 0.5, 0.5, 1)
+ bsdf.inputs[1].default_value = 1
+ output = new_mat.node_tree.nodes.new("ShaderNodeOutputMaterial")
+ new_mat.node_tree.links.new(bsdf.outputs["BSDF"], output.inputs["Surface"])
+ bpy.context.scene.view_layers["View Layer"].material_override = new_mat
+
+
+def unhide_all_objects() -> None:
+ """Unhides all objects in the scene.
+
+ Returns:
+ None
+ """
+ for obj in bpy.context.scene.objects:
+ obj.hide_set(False)
+
+
+def convert_to_meshes() -> None:
+ """Converts all objects in the scene to meshes.
+
+ Returns:
+ None
+ """
+ bpy.ops.object.select_all(action="DESELECT")
+ bpy.context.view_layer.objects.active = [
+ obj for obj in bpy.context.scene.objects if obj.type == "MESH"
+ ][0]
+ for obj in bpy.context.scene.objects:
+ obj.select_set(True)
+ bpy.ops.object.convert(target="MESH")
+
+
+def triangulate_meshes() -> None:
+ """Triangulates all meshes in the scene.
+
+ Returns:
+ None
+ """
+ bpy.ops.object.select_all(action="DESELECT")
+ objs = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"]
+ bpy.context.view_layer.objects.active = objs[0]
+ for obj in objs:
+ obj.select_set(True)
+ bpy.ops.object.mode_set(mode="EDIT")
+ bpy.ops.mesh.reveal()
+ bpy.ops.mesh.select_all(action="SELECT")
+ bpy.ops.mesh.quads_convert_to_tris(quad_method="BEAUTY", ngon_method="BEAUTY")
+ bpy.ops.object.mode_set(mode="OBJECT")
+ bpy.ops.object.select_all(action="DESELECT")
+
+
+def scene_bbox() -> Tuple[Vector, Vector]:
+ """Returns the bounding box of the scene.
+
+ Taken from Shap-E rendering script
+ (https://github.com/openai/shap-e/blob/main/shap_e/rendering/blender/blender_script.py#L68-L82)
+
+ Returns:
+ Tuple[Vector, Vector]: The minimum and maximum coordinates of the bounding box.
+ """
+ bbox_min = (math.inf,) * 3
+ bbox_max = (-math.inf,) * 3
+ found = False
+ scene_meshes = [
+ obj
+ for obj in bpy.context.scene.objects.values()
+ if isinstance(obj.data, bpy.types.Mesh)
+ ]
+ for obj in scene_meshes:
+ found = True
+ for coord in obj.bound_box:
+ coord = Vector(coord)
+ coord = obj.matrix_world @ coord
+ bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
+ bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))
+ if not found:
+ raise RuntimeError("no objects in scene to compute bounding box for")
+ return Vector(bbox_min), Vector(bbox_max)
+
+
+def normalize_scene() -> Tuple[float, Vector]:
+ """Normalizes the scene by scaling and translating it to fit in a unit cube centered
+ at the origin.
+
+ Mostly taken from the Point-E / Shap-E rendering script
+ (https://github.com/openai/point-e/blob/main/point_e/evals/scripts/blender_script.py#L97-L112),
+ but fix for multiple root objects: (see bug report here:
+ https://github.com/openai/shap-e/pull/60).
+
+ Returns:
+ Tuple[float, Vector]: The scale factor and the offset applied to the scene.
+ """
+ scene_root_objects = [
+ obj for obj in bpy.context.scene.objects.values() if not obj.parent
+ ]
+ if len(scene_root_objects) > 1:
+ # create an empty object to be used as a parent for all root objects
+ scene = bpy.data.objects.new("ParentEmpty", None)
+ bpy.context.scene.collection.objects.link(scene)
+
+ # parent all root objects to the empty object
+ for obj in scene_root_objects:
+ obj.parent = scene
+ else:
+ scene = scene_root_objects[0]
+
+ bbox_min, bbox_max = scene_bbox()
+ scale = 1 / max(bbox_max - bbox_min)
+ scene.scale = scene.scale * scale
+
+ # Apply scale to matrix_world.
+ bpy.context.view_layer.update()
+ bbox_min, bbox_max = scene_bbox()
+ offset = -(bbox_min + bbox_max) / 2
+ scene.matrix_world.translation += offset
+ bpy.ops.object.select_all(action="DESELECT")
+
+ return scale, offset
+
+
+def get_transform_matrix(obj: bpy.types.Object) -> list:
+ pos, rt, _ = obj.matrix_world.decompose()
+ rt = rt.to_matrix()
+ matrix = []
+ for ii in range(3):
+ a = []
+ for jj in range(3):
+ a.append(rt[ii][jj])
+ a.append(pos[ii])
+ matrix.append(a)
+ matrix.append([0, 0, 0, 1])
+ return matrix
+
+
+def main(arg):
+ os.makedirs(arg.output_folder, exist_ok=True)
+
+ # Initialize context
+ init_render(
+ engine=arg.engine,
+ resolution=arg.resolution,
+ geo_mode=arg.geo_mode,
+ use_gpu=getattr(arg, "use_gpu", True),
+ gpu_device=getattr(arg, "gpu_device", "OPTIX"),
+ )
+ outputs, spec_nodes = init_nodes(
+ save_depth=arg.save_depth,
+ save_normal=arg.save_normal,
+ save_albedo=arg.save_albedo,
+ save_mist=arg.save_mist,
+ )
+ if arg.object.endswith(".blend"):
+ delete_invisible_objects()
+ else:
+ init_scene()
+ load_object(arg.object)
+ if arg.split_normal:
+ split_mesh_normal()
+ # delete_custom_normals()
+ print("[INFO] Scene initialized.")
+
+ # normalize scene
+ scale, offset = normalize_scene()
+ print("[INFO] Scene normalized.")
+
+ # Initialize camera and lighting
+ cam = init_camera()
+ init_lighting()
+ print("[INFO] Camera and lighting initialized.")
+
+ # Override material
+ if arg.geo_mode:
+ override_material()
+
+ # Create a list of views
+ to_export = {
+ "aabb": [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+ "scale": scale,
+ "offset": [offset.x, offset.y, offset.z],
+ "frames": [],
+ }
+ views = json.loads(arg.views)
+ for i, view in enumerate(views):
+ cam.location = (
+ view["radius"] * np.cos(view["yaw"]) * np.cos(view["pitch"]),
+ view["radius"] * np.sin(view["yaw"]) * np.cos(view["pitch"]),
+ view["radius"] * np.sin(view["pitch"]),
+ )
+ cam.data.lens = 16 / np.tan(view["fov"] / 2)
+
+ if arg.save_depth:
+ spec_nodes["depth_map"].inputs[1].default_value = view[
+ "radius"
+ ] - 0.5 * np.sqrt(3)
+ spec_nodes["depth_map"].inputs[2].default_value = view[
+ "radius"
+ ] + 0.5 * np.sqrt(3)
+
+ bpy.context.scene.render.filepath = os.path.join(
+ arg.output_folder, f"{i:03d}.png"
+ )
+ for name, output in outputs.items():
+ output.file_slots[0].path = os.path.join(
+ arg.output_folder, f"{i:03d}_{name}"
+ )
+
+ # Render the scene
+ bpy.ops.render.render(write_still=True)
+ bpy.context.view_layer.update()
+ for name, output in outputs.items():
+ ext = EXT[output.format.file_format]
+ path = glob.glob(f"{output.file_slots[0].path}*.{ext}")[0]
+ os.rename(path, f"{output.file_slots[0].path}.{ext}")
+
+ # Save camera parameters
+ metadata = {
+ "file_path": f"{i:03d}.png",
+ "camera_angle_x": view["fov"],
+ "transform_matrix": get_transform_matrix(cam),
+ }
+ if arg.save_depth:
+ metadata["depth"] = {
+ "min": view["radius"] - 0.5 * np.sqrt(3),
+ "max": view["radius"] + 0.5 * np.sqrt(3),
+ }
+ to_export["frames"].append(metadata)
+
+ # Save the camera parameters
+ with open(os.path.join(arg.output_folder, "transforms.json"), "w") as f:
+ json.dump(to_export, f, indent=4)
+
+ if arg.save_mesh:
+ # triangulate meshes
+ unhide_all_objects()
+ convert_to_meshes()
+ triangulate_meshes()
+ print("[INFO] Meshes triangulated.")
+
+ # export ply mesh
+ bpy.ops.export_mesh.ply(filepath=os.path.join(arg.output_folder, "mesh.ply"))
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(
+ description="Renders given obj file by rotation a camera around it."
+ )
+ parser.add_argument(
+ "--views",
+ type=str,
+ help="JSON string of views. Contains a list of {yaw, pitch, radius, fov} object.",
+ )
+ parser.add_argument(
+ "--object", type=str, help="Path to the 3D model file to be rendered."
+ )
+ parser.add_argument(
+ "--output_folder",
+ type=str,
+ default="/tmp",
+ help="The path the output will be dumped to.",
+ )
+ parser.add_argument(
+ "--resolution", type=int, default=512, help="Resolution of the images."
+ )
+ parser.add_argument(
+ "--engine",
+ type=str,
+ default="CYCLES",
+ help="Blender internal engine for rendering. E.g. CYCLES, BLENDER_EEVEE, ...",
+ )
+ parser.add_argument(
+ "--geo_mode", action="store_true", help="Geometry mode for rendering."
+ )
+ parser.add_argument(
+ "--save_depth", action="store_true", help="Save the depth maps."
+ )
+ parser.add_argument(
+ "--save_normal", action="store_true", help="Save the normal maps."
+ )
+ parser.add_argument(
+ "--save_albedo", action="store_true", help="Save the albedo maps."
+ )
+ parser.add_argument(
+ "--save_mist", action="store_true", help="Save the mist distance maps."
+ )
+ parser.add_argument(
+ "--split_normal", action="store_true", help="Split the normals of the mesh."
+ )
+ parser.add_argument(
+ "--save_mesh", action="store_true", help="Save the mesh as a .ply file."
+ )
+ parser.add_argument(
+ "--use_gpu", action="store_true", help="Use GPU acceleration for rendering."
+ )
+ parser.add_argument(
+ "--gpu_device",
+ type=str,
+ default="OPTIX",
+ choices=["OPTIX", "CUDA", "OPENCL"],
+ help="GPU device type for rendering (OPTIX, CUDA, or OPENCL).",
+ )
+ argv = sys.argv[sys.argv.index("--") + 1 :]
+ args = parser.parse_args(argv)
+
+ main(args)
diff --git a/deps/vomp/dataset_toolkits/build_metadata.py b/deps/vomp/dataset_toolkits/build_metadata.py
new file mode 100644
index 0000000000000000000000000000000000000000..37ad95fe928c31e5a8a82a1b2d4613d2809c1d88
--- /dev/null
+++ b/deps/vomp/dataset_toolkits/build_metadata.py
@@ -0,0 +1,551 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import shutil
+import sys
+import time
+import importlib
+import argparse
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+from easydict import EasyDict as edict
+from concurrent.futures import ThreadPoolExecutor
+import utils3d
+
+
+def get_first_directory(path):
+ with os.scandir(path) as it:
+ for entry in it:
+ if entry.is_dir():
+ return entry.name
+ return None
+
+
+def need_process(key):
+ return key in opt.field or opt.field == ["all"]
+
+
+if __name__ == "__main__":
+ dataset_utils = importlib.import_module(f"dataset_toolkits.datasets.{sys.argv[1]}")
+
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--output_dir", type=str, required=True, help="Directory to save the metadata"
+ )
+ parser.add_argument(
+ "--field",
+ type=str,
+ default="all",
+ help="Fields to process, separated by commas",
+ )
+ parser.add_argument(
+ "--from_file",
+ action="store_true",
+ help="Build metadata from file instead of from records of processings."
+ + "Useful when some processing fail to generate records but file already exists.",
+ )
+ parser.add_argument(
+ "--force_update_class_split",
+ action="store_true",
+ help="Force updating class and split information even if metadata file exists",
+ )
+ parser.add_argument(
+ "--skip_class_split_on_error",
+ action="store_true",
+ help="Skip updating class and split if an error occurs, instead of failing",
+ )
+ dataset_utils.add_args(parser)
+ opt = parser.parse_args(sys.argv[2:])
+ opt = edict(vars(opt))
+
+ os.makedirs(opt.output_dir, exist_ok=True)
+ os.makedirs(os.path.join(opt.output_dir, "merged_records"), exist_ok=True)
+
+ opt.field = opt.field.split(",")
+
+ timestamp = str(int(time.time()))
+
+ # Check if metadata file exists
+ metadata_exists = os.path.exists(os.path.join(opt.output_dir, "metadata.csv"))
+
+ # Load or create metadata
+ if metadata_exists:
+ print("Loading previous metadata...")
+ metadata = pd.read_csv(os.path.join(opt.output_dir, "metadata.csv"))
+
+ # Check if class and split information needs to be updated
+ requires_class_update = (
+ "class" not in metadata.columns or opt.force_update_class_split
+ )
+ requires_split_update = (
+ "split" not in metadata.columns or opt.force_update_class_split
+ )
+
+ if requires_class_update or requires_split_update:
+ # Generate fresh metadata with class and split information
+ print("Updating class and split information...")
+ try:
+ fresh_metadata = dataset_utils.get_metadata(**opt)
+
+ # Set index on sha256 for both DataFrames
+ metadata.set_index("sha256", inplace=True)
+ fresh_metadata.set_index("sha256", inplace=True)
+
+ # Update class information if needed
+ if requires_class_update and "class" in fresh_metadata.columns:
+ if "class" not in metadata.columns:
+ metadata["class"] = "unknown"
+ metadata.update(fresh_metadata[["class"]])
+
+ # Update split information if needed
+ if requires_split_update and "split" in fresh_metadata.columns:
+ if "split" not in metadata.columns:
+ metadata["split"] = "train" # Default value
+ metadata.update(fresh_metadata[["split"]])
+ except Exception as e:
+ if opt.skip_class_split_on_error:
+ print(f"Warning: Error updating class and split information: {e}")
+ print("Continuing with existing metadata...")
+ if "class" not in metadata.columns:
+ metadata["class"] = "unknown"
+ if "split" not in metadata.columns:
+ metadata["split"] = "train"
+ metadata.set_index("sha256", inplace=True)
+ else:
+ raise e
+ else:
+ metadata.set_index("sha256", inplace=True)
+ else:
+ # Create new metadata with all required information
+ print("Creating new metadata...")
+ try:
+ metadata = dataset_utils.get_metadata(**opt)
+ metadata.set_index("sha256", inplace=True)
+ except Exception as e:
+ if opt.skip_class_split_on_error:
+ print(
+ f"Warning: Error creating metadata with class and split information: {e}"
+ )
+ print("Creating basic metadata without class and split information...")
+ metadata = dataset_utils.get_metadata(skip_split=True, **opt)
+ metadata.set_index("sha256", inplace=True)
+ if "class" not in metadata.columns:
+ metadata["class"] = "unknown"
+ if "split" not in metadata.columns:
+ metadata["split"] = "train"
+ else:
+ raise e
+
+ # merge downloaded
+ df_files = [
+ f
+ for f in os.listdir(opt.output_dir)
+ if f.startswith("downloaded_") and f.endswith(".csv")
+ ]
+ df_parts = []
+ for f in df_files:
+ try:
+ df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f)))
+ except:
+ pass
+ if len(df_parts) > 0:
+ df = pd.concat(df_parts)
+ df.set_index("sha256", inplace=True)
+ if "local_path" in metadata.columns:
+ metadata.update(df, overwrite=True)
+ else:
+ metadata = metadata.join(df, on="sha256", how="left")
+ for f in df_files:
+ shutil.move(
+ os.path.join(opt.output_dir, f),
+ os.path.join(opt.output_dir, "merged_records", f"{timestamp}_{f}"),
+ )
+
+ # detect models
+ image_models = []
+ if os.path.exists(os.path.join(opt.output_dir, "features")):
+ image_models = os.listdir(os.path.join(opt.output_dir, "features"))
+ latent_models = []
+ if os.path.exists(os.path.join(opt.output_dir, "latents")):
+ latent_models = os.listdir(os.path.join(opt.output_dir, "latents"))
+ ss_latent_models = []
+ if os.path.exists(os.path.join(opt.output_dir, "ss_latents")):
+ ss_latent_models = os.listdir(os.path.join(opt.output_dir, "ss_latents"))
+ print(f"Image models: {image_models}")
+ print(f"Latent models: {latent_models}")
+ print(f"Sparse Structure latent models: {ss_latent_models}")
+
+ if "rendered" not in metadata.columns:
+ metadata["rendered"] = [False] * len(metadata)
+ if "voxelized" not in metadata.columns:
+ metadata["voxelized"] = [False] * len(metadata)
+ if "num_voxels" not in metadata.columns:
+ metadata["num_voxels"] = [0] * len(metadata)
+ if "cond_rendered" not in metadata.columns:
+ metadata["cond_rendered"] = [False] * len(metadata)
+ for model in image_models:
+ if f"feature_{model}" not in metadata.columns:
+ metadata[f"feature_{model}"] = [False] * len(metadata)
+ for model in latent_models:
+ if f"latent_{model}" not in metadata.columns:
+ metadata[f"latent_{model}"] = [False] * len(metadata)
+ for model in ss_latent_models:
+ if f"ss_latent_{model}" not in metadata.columns:
+ metadata[f"ss_latent_{model}"] = [False] * len(metadata)
+
+ # merge rendered
+ df_files = [
+ f
+ for f in os.listdir(opt.output_dir)
+ if f.startswith("rendered_") and f.endswith(".csv")
+ ]
+ df_parts = []
+ for f in df_files:
+ try:
+ df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f)))
+ except:
+ pass
+ if len(df_parts) > 0:
+ df = pd.concat(df_parts)
+ df.set_index("sha256", inplace=True)
+ metadata.update(df, overwrite=True)
+ for f in df_files:
+ shutil.move(
+ os.path.join(opt.output_dir, f),
+ os.path.join(opt.output_dir, "merged_records", f"{timestamp}_{f}"),
+ )
+
+ # merge voxelized
+ df_files = [
+ f
+ for f in os.listdir(opt.output_dir)
+ if f.startswith("voxelized_") and f.endswith(".csv")
+ ]
+ df_parts = []
+ for f in df_files:
+ try:
+ df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f)))
+ except:
+ pass
+ if len(df_parts) > 0:
+ df = pd.concat(df_parts)
+ df.set_index("sha256", inplace=True)
+ metadata.update(df, overwrite=True)
+ for f in df_files:
+ shutil.move(
+ os.path.join(opt.output_dir, f),
+ os.path.join(opt.output_dir, "merged_records", f"{timestamp}_{f}"),
+ )
+
+ # merge cond_rendered
+ df_files = [
+ f
+ for f in os.listdir(opt.output_dir)
+ if f.startswith("cond_rendered_") and f.endswith(".csv")
+ ]
+ df_parts = []
+ for f in df_files:
+ try:
+ df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f)))
+ except:
+ pass
+ if len(df_parts) > 0:
+ df = pd.concat(df_parts)
+ df.set_index("sha256", inplace=True)
+ metadata.update(df, overwrite=True)
+ for f in df_files:
+ shutil.move(
+ os.path.join(opt.output_dir, f),
+ os.path.join(opt.output_dir, "merged_records", f"{timestamp}_{f}"),
+ )
+
+ # merge features
+ for model in image_models:
+ df_files = [
+ f
+ for f in os.listdir(opt.output_dir)
+ if f.startswith(f"feature_{model}_") and f.endswith(".csv")
+ ]
+ df_parts = []
+ for f in df_files:
+ try:
+ df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f)))
+ except:
+ pass
+ if len(df_parts) > 0:
+ df = pd.concat(df_parts)
+ df.set_index("sha256", inplace=True)
+ metadata.update(df, overwrite=True)
+ for f in df_files:
+ shutil.move(
+ os.path.join(opt.output_dir, f),
+ os.path.join(opt.output_dir, "merged_records", f"{timestamp}_{f}"),
+ )
+
+ # merge latents
+ for model in latent_models:
+ df_files = [
+ f
+ for f in os.listdir(opt.output_dir)
+ if f.startswith(f"latent_{model}_") and f.endswith(".csv")
+ ]
+ df_parts = []
+ for f in df_files:
+ try:
+ df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f)))
+ except:
+ pass
+ if len(df_parts) > 0:
+ df = pd.concat(df_parts)
+ df.set_index("sha256", inplace=True)
+ metadata.update(df, overwrite=True)
+ for f in df_files:
+ shutil.move(
+ os.path.join(opt.output_dir, f),
+ os.path.join(opt.output_dir, "merged_records", f"{timestamp}_{f}"),
+ )
+
+ # merge sparse structure latents
+ for model in ss_latent_models:
+ df_files = [
+ f
+ for f in os.listdir(opt.output_dir)
+ if f.startswith(f"ss_latent_{model}_") and f.endswith(".csv")
+ ]
+ df_parts = []
+ for f in df_files:
+ try:
+ df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f)))
+ except:
+ pass
+ if len(df_parts) > 0:
+ df = pd.concat(df_parts)
+ df.set_index("sha256", inplace=True)
+ metadata.update(df, overwrite=True)
+ for f in df_files:
+ shutil.move(
+ os.path.join(opt.output_dir, f),
+ os.path.join(opt.output_dir, "merged_records", f"{timestamp}_{f}"),
+ )
+
+ # build metadata from files
+ if opt.from_file:
+ with (
+ ThreadPoolExecutor(max_workers=os.cpu_count()) as executor,
+ tqdm(total=len(metadata), desc="Building metadata") as pbar,
+ ):
+
+ def worker(sha256):
+ try:
+ if (
+ need_process("rendered")
+ and metadata.loc[sha256, "rendered"] == False
+ and os.path.exists(
+ os.path.join(
+ opt.output_dir, "renders", sha256, "transforms.json"
+ )
+ )
+ ):
+ metadata.loc[sha256, "rendered"] = True
+ if (
+ need_process("voxelized")
+ and metadata.loc[sha256, "rendered"] == True
+ and metadata.loc[sha256, "voxelized"] == False
+ and os.path.exists(
+ os.path.join(opt.output_dir, "voxels", f"{sha256}.ply")
+ )
+ ):
+ try:
+ pts = utils3d.io.read_ply(
+ os.path.join(opt.output_dir, "voxels", f"{sha256}.ply")
+ )[0]
+ metadata.loc[sha256, "voxelized"] = True
+ metadata.loc[sha256, "num_voxels"] = len(pts)
+ except Exception as e:
+ pass
+ if (
+ need_process("cond_rendered")
+ and metadata.loc[sha256, "cond_rendered"] == False
+ and os.path.exists(
+ os.path.join(
+ opt.output_dir,
+ "renders_cond",
+ sha256,
+ "transforms.json",
+ )
+ )
+ ):
+ metadata.loc[sha256, "cond_rendered"] = True
+ for model in image_models:
+ if (
+ need_process(f"feature_{model}")
+ and metadata.loc[sha256, f"feature_{model}"] == False
+ and metadata.loc[sha256, "rendered"] == True
+ and metadata.loc[sha256, "voxelized"] == True
+ and os.path.exists(
+ os.path.join(
+ opt.output_dir, "features", model, f"{sha256}.npz"
+ )
+ )
+ ):
+ metadata.loc[sha256, f"feature_{model}"] = True
+ for model in latent_models:
+ if (
+ need_process(f"latent_{model}")
+ and metadata.loc[sha256, f"latent_{model}"] == False
+ and metadata.loc[sha256, "rendered"] == True
+ and metadata.loc[sha256, "voxelized"] == True
+ and os.path.exists(
+ os.path.join(
+ opt.output_dir, "latents", model, f"{sha256}.npz"
+ )
+ )
+ ):
+ metadata.loc[sha256, f"latent_{model}"] = True
+ for model in ss_latent_models:
+ if (
+ need_process(f"ss_latent_{model}")
+ and metadata.loc[sha256, f"ss_latent_{model}"] == False
+ and metadata.loc[sha256, "voxelized"] == True
+ and os.path.exists(
+ os.path.join(
+ opt.output_dir, "ss_latents", model, f"{sha256}.npz"
+ )
+ )
+ ):
+ metadata.loc[sha256, f"ss_latent_{model}"] = True
+ pbar.update()
+ except Exception as e:
+ print(f"Error processing {sha256}: {e}")
+ pbar.update()
+
+ executor.map(worker, metadata.index)
+ executor.shutdown(wait=True)
+
+ # Save dataset splits if we have split information
+ if "split" in metadata.columns:
+ os.makedirs(os.path.join(opt.output_dir, "splits"), exist_ok=True)
+ # Reset index to include sha256 in the exported files
+ metadata_export = metadata.reset_index()
+ for split in ["train", "val", "test"]:
+ split_df = metadata_export[metadata_export["split"] == split]
+ if not split_df.empty:
+ split_df.to_csv(
+ os.path.join(opt.output_dir, "splits", f"{split}.csv"), index=False
+ )
+
+ # statistics
+ metadata.to_csv(os.path.join(opt.output_dir, "metadata.csv"))
+ num_downloaded = (
+ metadata["local_path"].count() if "local_path" in metadata.columns else 0
+ )
+
+ # If from_file is True, update metadata to reflect actual files on disk before writing statistics
+ if opt.from_file:
+ print("Updating metadata to reflect actual files on disk...")
+ for model in image_models:
+ for sha256 in metadata.index:
+ actual_exists = os.path.exists(
+ os.path.join(opt.output_dir, "features", model, f"{sha256}.npz")
+ )
+ metadata.loc[sha256, f"feature_{model}"] = actual_exists
+
+ for model in latent_models:
+ for sha256 in metadata.index:
+ actual_exists = os.path.exists(
+ os.path.join(opt.output_dir, "latents", model, f"{sha256}.npz")
+ )
+ metadata.loc[sha256, f"latent_{model}"] = actual_exists
+
+ for model in ss_latent_models:
+ for sha256 in metadata.index:
+ actual_exists = os.path.exists(
+ os.path.join(opt.output_dir, "ss_latents", model, f"{sha256}.npz")
+ )
+ metadata.loc[sha256, f"ss_latent_{model}"] = actual_exists
+
+ # Save updated metadata
+ metadata.to_csv(os.path.join(opt.output_dir, "metadata.csv"))
+
+ with open(os.path.join(opt.output_dir, "statistics.txt"), "w") as f:
+ f.write("Statistics:\n")
+ f.write(f" - Number of assets: {len(metadata)}\n")
+ f.write(f" - Number of assets downloaded: {num_downloaded}\n")
+ f.write(f' - Number of assets rendered: {metadata["rendered"].sum()}\n')
+ f.write(f' - Number of assets voxelized: {metadata["voxelized"].sum()}\n')
+ if len(image_models) != 0:
+ f.write(f" - Number of assets with image features extracted:\n")
+ for model in image_models:
+ # Always use metadata counts since they're now accurate when from_file=True
+ f.write(f' - {model}: {metadata[f"feature_{model}"].sum()}\n')
+ if len(latent_models) != 0:
+ f.write(f" - Number of assets with latents extracted:\n")
+ for model in latent_models:
+ f.write(f' - {model}: {metadata[f"latent_{model}"].sum()}\n')
+ if len(ss_latent_models) != 0:
+ f.write(f" - Number of assets with sparse structure latents extracted:\n")
+ for model in ss_latent_models:
+ f.write(f' - {model}: {metadata[f"ss_latent_{model}"].sum()}\n')
+
+ # Only report captions if the column exists (it may not for Gaussian splats)
+ if "captions" in metadata.columns:
+ f.write(
+ f' - Number of assets with captions: {metadata["captions"].count()}\n'
+ )
+ else:
+ f.write(
+ f" - Number of assets with captions: N/A (no caption data available)\n"
+ )
+
+ f.write(
+ f' - Number of assets with image conditions: {metadata["cond_rendered"].sum()}\n'
+ )
+
+ # Add class distribution statistics
+ if "class" in metadata.columns:
+ f.write("\nClass distribution:\n")
+ class_counts = metadata["class"].value_counts()
+ for class_name, count in class_counts.items():
+ f.write(f" - {class_name}: {count} ({count/len(metadata)*100:.1f}%)\n")
+
+ # Add split statistics if split column exists
+ if "split" in metadata.columns:
+ f.write("\nDataset splits:\n")
+ split_counts = metadata["split"].value_counts()
+ for split_name, count in split_counts.items():
+ f.write(f" - {split_name}: {count} ({count/len(metadata)*100:.1f}%)\n")
+
+ # Add class distribution per split if both columns exist
+ if "class" in metadata.columns:
+ f.write("\nClass distribution per split:\n")
+ # Reset index to allow cross-tabulation
+ metadata_reset = metadata.reset_index()
+ # For each split, show class distribution
+ for split_name in ["train", "val", "test"]:
+ if split_name in split_counts:
+ f.write(f" {split_name.upper()}:\n")
+ split_data = metadata_reset[
+ metadata_reset["split"] == split_name
+ ]
+ class_in_split = split_data["class"].value_counts()
+ for class_name, count in class_in_split.items():
+ f.write(
+ f" - {class_name}: {count} ({count/len(split_data)*100:.1f}%)\n"
+ )
+
+ with open(os.path.join(opt.output_dir, "statistics.txt"), "r") as f:
+ print(f.read())
diff --git a/deps/vomp/dataset_toolkits/datasets/ABO.py b/deps/vomp/dataset_toolkits/datasets/ABO.py
new file mode 100644
index 0000000000000000000000000000000000000000..407893efc5007aa570b9b3dd6b09266aab9805d6
--- /dev/null
+++ b/deps/vomp/dataset_toolkits/datasets/ABO.py
@@ -0,0 +1,132 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import re
+import argparse
+import tarfile
+from concurrent.futures import ThreadPoolExecutor
+from tqdm import tqdm
+import pandas as pd
+from utils import get_file_hash
+
+
+def add_args(parser: argparse.ArgumentParser):
+ pass
+
+
+def get_metadata(**kwargs):
+ metadata = pd.read_csv("hf://datasets/JeffreyXiang/TRELLIS-500K/ABO.csv")
+ return metadata
+
+
+def download(metadata, output_dir, **kwargs):
+ os.makedirs(os.path.join(output_dir, "raw"), exist_ok=True)
+
+ if not os.path.exists(os.path.join(output_dir, "raw", "abo-3dmodels.tar")):
+ try:
+ os.makedirs(os.path.join(output_dir, "raw"), exist_ok=True)
+ os.system(
+ f"wget -O {output_dir}/raw/abo-3dmodels.tar https://amazon-berkeley-objects.s3.amazonaws.com/archives/abo-3dmodels.tar"
+ )
+ except:
+ print("\033[93m")
+ print(
+ "Error downloading ABO dataset. Please check your internet connection and try again."
+ )
+ print(
+ "Or, you can manually download the abo-3dmodels.tar file and place it in the {output_dir}/raw directory"
+ )
+ print(
+ "Visit https://amazon-berkeley-objects.s3.amazonaws.com/index.html for more information"
+ )
+ print("\033[0m")
+ raise FileNotFoundError("Error downloading ABO dataset")
+
+ downloaded = {}
+ metadata = metadata.set_index("file_identifier")
+ with tarfile.open(os.path.join(output_dir, "raw", "abo-3dmodels.tar")) as tar:
+ with (
+ ThreadPoolExecutor(max_workers=1) as executor,
+ tqdm(total=len(metadata), desc="Extracting") as pbar,
+ ):
+
+ def worker(instance: str) -> str:
+ try:
+ tar.extract(
+ f"3dmodels/original/{instance}",
+ path=os.path.join(output_dir, "raw"),
+ )
+ sha256 = get_file_hash(
+ os.path.join(output_dir, "raw/3dmodels/original", instance)
+ )
+ pbar.update()
+ return sha256
+ except Exception as e:
+ pbar.update()
+ print(f"Error extracting for {instance}: {e}")
+ return None
+
+ sha256s = executor.map(worker, metadata.index)
+ executor.shutdown(wait=True)
+
+ for k, sha256 in zip(metadata.index, sha256s):
+ if sha256 is not None:
+ if sha256 == metadata.loc[k, "sha256"]:
+ downloaded[sha256] = os.path.join("raw/3dmodels/original", k)
+ else:
+ print(f"Error downloading {k}: sha256s do not match")
+
+ return pd.DataFrame(downloaded.items(), columns=["sha256", "local_path"])
+
+
+def foreach_instance(
+ metadata, output_dir, func, max_workers=None, desc="Processing objects"
+) -> pd.DataFrame:
+ import os
+ from concurrent.futures import ThreadPoolExecutor
+ from tqdm import tqdm
+
+ # load metadata
+ metadata = metadata.to_dict("records")
+
+ # processing objects
+ records = []
+ max_workers = max_workers or os.cpu_count()
+ try:
+ with (
+ ThreadPoolExecutor(max_workers=max_workers) as executor,
+ tqdm(total=len(metadata), desc=desc) as pbar,
+ ):
+
+ def worker(metadatum):
+ try:
+ local_path = metadatum["local_path"]
+ sha256 = metadatum["sha256"]
+ file = os.path.join(output_dir, local_path)
+ record = func(file, sha256)
+ if record is not None:
+ records.append(record)
+ pbar.update()
+ except Exception as e:
+ print(f"Error processing object {sha256}: {e}")
+ pbar.update()
+
+ executor.map(worker, metadata)
+ executor.shutdown(wait=True)
+ except:
+ print("Error happened during processing.")
+
+ return pd.DataFrame.from_records(records)
diff --git a/deps/vomp/dataset_toolkits/datasets/__init__.py b/deps/vomp/dataset_toolkits/datasets/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..5fcd7fab2c8148a7bd3eb6a8aa2105e1a8f0e6a5
--- /dev/null
+++ b/deps/vomp/dataset_toolkits/datasets/__init__.py
@@ -0,0 +1,16 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Dataset modules for TRELLIS preprocessing
diff --git a/deps/vomp/dataset_toolkits/datasets/allmats.py b/deps/vomp/dataset_toolkits/datasets/allmats.py
new file mode 100644
index 0000000000000000000000000000000000000000..f4267535eb526c798a64cb7a3ffb6d1a34854bc5
--- /dev/null
+++ b/deps/vomp/dataset_toolkits/datasets/allmats.py
@@ -0,0 +1,510 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import os
+import json
+import pandas as pd
+import numpy as np
+import hashlib
+import random
+from glob import glob
+from sklearn.model_selection import train_test_split
+from typing import Dict, List, Optional, Any, Union
+
+
+from dataset_toolkits.material_objects.vlm_annotations.utils.utils import (
+ SIMREADY_PROPS_DIR,
+ COMMERCIAL_BASE_DIR,
+ RESIDENTIAL_BASE_DIR,
+ VEGETATION_BASE_DIR,
+ SIMREADY_ASSET_CLASS_MAPPING,
+ SIMREADY_ASSET_INFO_PATH,
+)
+from dataset_toolkits.material_objects.vlm_annotations.data_subsets import (
+ simready,
+ commercial,
+ vegetation,
+ residential,
+ common,
+)
+
+
+def set_seeds(seed=42):
+ random.seed(seed)
+ np.random.seed(seed)
+
+
+def add_args(parser):
+ parser.add_argument(
+ "--simready_dir",
+ type=str,
+ default=SIMREADY_PROPS_DIR,
+ help="Path to the SimReady props directory",
+ )
+ parser.add_argument(
+ "--commercial_dir",
+ type=str,
+ default=COMMERCIAL_BASE_DIR,
+ help="Path to the Commercial models directory",
+ )
+ parser.add_argument(
+ "--residential_dir",
+ type=str,
+ default=RESIDENTIAL_BASE_DIR,
+ help="Path to the Residential models directory",
+ )
+ parser.add_argument(
+ "--vegetation_dir",
+ type=str,
+ default=VEGETATION_BASE_DIR,
+ help="Path to the Vegetation models directory",
+ )
+ parser.add_argument(
+ "--asset_info_path",
+ type=str,
+ default=SIMREADY_ASSET_INFO_PATH,
+ help="Path to the SimReady asset_info.json file",
+ )
+ parser.add_argument(
+ "--seed",
+ type=int,
+ default=42,
+ help="Random seed for reproducibility",
+ )
+ parser.add_argument(
+ "--default_class",
+ type=str,
+ default="unknown",
+ help="Default class label to use when class information is not available",
+ )
+ parser.add_argument(
+ "--include_datasets",
+ type=str,
+ default="simready,commercial,residential,vegetation",
+ help="Comma-separated list of datasets to include",
+ )
+
+
+def split_dataset(metadata, seed=42):
+
+ np.random.seed(seed)
+ random.seed(seed)
+
+ metadata_copy = metadata.copy()
+
+ metadata_copy["split"] = "train"
+
+ classes = metadata_copy["class"].unique()
+
+ large_classes = []
+ small_classes = []
+
+ for cls in classes:
+
+ cls_indices = metadata_copy[metadata_copy["class"] == cls].index.tolist()
+ if len(cls_indices) >= 10:
+ large_classes.append(cls)
+ else:
+ small_classes.append(cls)
+
+ for cls in large_classes:
+ cls_indices = metadata_copy[metadata_copy["class"] == cls].index.tolist()
+ random.shuffle(cls_indices)
+
+ n_samples = len(cls_indices)
+ n_train = int(0.8 * n_samples)
+ n_val = int(0.1 * n_samples)
+
+ train_indices = cls_indices[:n_train]
+ val_indices = cls_indices[n_train : n_train + n_val]
+ test_indices = cls_indices[n_train + n_val :]
+
+ metadata_copy.loc[train_indices, "split"] = "train"
+ metadata_copy.loc[val_indices, "split"] = "val"
+ metadata_copy.loc[test_indices, "split"] = "test"
+
+ total_samples = len(metadata_copy)
+ goal_train = int(0.8 * total_samples)
+ goal_val = int(0.1 * total_samples)
+ goal_test = total_samples - goal_train - goal_val
+
+ current_train = (metadata_copy["split"] == "train").sum()
+ current_val = (metadata_copy["split"] == "val").sum()
+ current_test = (metadata_copy["split"] == "test").sum()
+
+ small_indices = []
+ for cls in small_classes:
+ cls_indices = metadata_copy[metadata_copy["class"] == cls].index.tolist()
+ small_indices.extend(cls_indices)
+
+ random.shuffle(small_indices)
+
+ need_train = max(0, goal_train - current_train)
+ need_val = max(0, goal_val - current_val)
+ need_test = max(0, goal_test - current_test)
+
+ idx = 0
+ while idx < len(small_indices):
+ if need_train > 0:
+ metadata_copy.loc[small_indices[idx], "split"] = "train"
+ need_train -= 1
+ idx += 1
+ elif need_val > 0:
+ metadata_copy.loc[small_indices[idx], "split"] = "val"
+ need_val -= 1
+ idx += 1
+ elif need_test > 0:
+ metadata_copy.loc[small_indices[idx], "split"] = "test"
+ need_test -= 1
+ idx += 1
+ else:
+
+ metadata_copy.loc[small_indices[idx:], "split"] = "train"
+ break
+
+ train_count = (metadata_copy["split"] == "train").sum()
+ val_count = (metadata_copy["split"] == "val").sum()
+ test_count = (metadata_copy["split"] == "test").sum()
+
+ print(
+ f"Dataset split: Train: {train_count} ({train_count/len(metadata_copy)*100:.1f}%), "
+ f"Val: {val_count} ({val_count/len(metadata_copy)*100:.1f}%), "
+ f"Test: {test_count} ({test_count/len(metadata_copy)*100:.1f}%)"
+ )
+
+ if small_classes:
+ print("\nSmall class distribution across splits:")
+ for cls in small_classes:
+ cls_data = metadata_copy[metadata_copy["class"] == cls]
+ cls_train = (cls_data["split"] == "train").sum()
+ cls_val = (cls_data["split"] == "val").sum()
+ cls_test = (cls_data["split"] == "test").sum()
+ cls_total = len(cls_data)
+ print(
+ f" - {cls} (total {cls_total}): Train: {cls_train}, Val: {cls_val}, Test: {cls_test}"
+ )
+
+ return metadata_copy
+
+
+def get_simready_metadata(simready_dir, asset_info_path, default_class="unknown"):
+
+ asset_class_mapping = SIMREADY_ASSET_CLASS_MAPPING
+
+ if not asset_class_mapping and asset_info_path and os.path.exists(asset_info_path):
+ try:
+ with open(asset_info_path, "r") as f:
+ asset_info = json.load(f)
+
+ asset_class_mapping = {}
+ for asset in asset_info:
+ simple_name = asset.get("Simple Name")
+ if simple_name and "Labels" in asset and "Class" in asset["Labels"]:
+ asset_class_mapping[simple_name] = asset["Labels"]["Class"]
+
+ print(f"Loaded class information for {len(asset_class_mapping)} assets")
+ except Exception as e:
+ print(f"Error loading asset info: {e}")
+
+ prop_dirs = []
+ if os.path.exists(simready_dir):
+ prop_dirs = [
+ d
+ for d in os.listdir(simready_dir)
+ if os.path.isdir(os.path.join(simready_dir, d))
+ ]
+
+ metadata = []
+
+ for prop_name in prop_dirs:
+ prop_dir = os.path.join(simready_dir, prop_name)
+
+ usd_files = glob(os.path.join(prop_dir, "*.usd"))
+ if not usd_files:
+ continue
+
+ inst_base_files = [f for f in usd_files if "_inst_base.usd" in f]
+ base_files = [f for f in usd_files if "_base.usd" in f]
+
+ if inst_base_files:
+ usd_file = inst_base_files[0]
+ elif base_files:
+ usd_file = base_files[0]
+ else:
+ usd_file = usd_files[0]
+
+ sha256 = hashlib.sha256(prop_name.encode()).hexdigest()
+
+ prop_class = asset_class_mapping.get(prop_name, default_class)
+
+ metadata.append(
+ {
+ "sha256": sha256,
+ "local_path": usd_file,
+ "original_name": prop_name,
+ "aesthetic_score": 1.0,
+ "rendered": False,
+ "class": prop_class,
+ "dataset": "simready",
+ }
+ )
+
+ return metadata
+
+
+def get_commercial_metadata(commercial_dir, default_class="commercial"):
+ metadata = []
+
+ if not os.path.exists(commercial_dir):
+ print(f"Commercial directory not found: {commercial_dir}")
+ return metadata
+
+ for root, _, files in os.walk(commercial_dir):
+ for file in files:
+ if file.endswith(".usd") and not os.path.basename(root).startswith("."):
+ usd_file = os.path.join(root, file)
+
+ object_name = os.path.basename(os.path.dirname(usd_file))
+
+ sha256 = hashlib.sha256(f"{object_name}_{file}".encode()).hexdigest()
+
+ try:
+ material_info = common.extract_materials_from_usd(
+ usd_file, "commercial"
+ )
+ category = material_info.get("category", default_class)
+ except Exception:
+ category = default_class
+
+ metadata.append(
+ {
+ "sha256": sha256,
+ "local_path": usd_file,
+ "original_name": f"{object_name}/{file}",
+ "aesthetic_score": 1.0,
+ "rendered": False,
+ "class": category,
+ "dataset": "commercial",
+ }
+ )
+
+ return metadata
+
+
+def get_residential_metadata(residential_dir, default_class="residential"):
+ metadata = []
+
+ if not os.path.exists(residential_dir):
+ print(f"Residential directory not found: {residential_dir}")
+ return metadata
+
+ for root, _, files in os.walk(residential_dir):
+ for file in files:
+ if file.endswith(".usd") and not os.path.basename(root).startswith("."):
+ usd_file = os.path.join(root, file)
+
+ object_name = os.path.basename(os.path.dirname(usd_file))
+
+ sha256 = hashlib.sha256(f"{object_name}_{file}".encode()).hexdigest()
+
+ try:
+ material_info = common.extract_materials_from_usd(
+ usd_file, "residential"
+ )
+ category = material_info.get("category", default_class)
+ except Exception:
+ category = default_class
+
+ metadata.append(
+ {
+ "sha256": sha256,
+ "local_path": usd_file,
+ "original_name": f"{object_name}/{file}",
+ "aesthetic_score": 1.0,
+ "rendered": False,
+ "class": category,
+ "dataset": "residential",
+ }
+ )
+
+ return metadata
+
+
+def get_vegetation_metadata(vegetation_dir, default_class="vegetation"):
+ metadata = []
+
+ if not os.path.exists(vegetation_dir):
+ print(f"Vegetation directory not found: {vegetation_dir}")
+ return metadata
+
+ for root, _, files in os.walk(vegetation_dir):
+ for file in files:
+ if file.endswith(".usd") and not os.path.basename(root).startswith("."):
+ usd_file = os.path.join(root, file)
+
+ object_name = os.path.basename(os.path.dirname(usd_file))
+
+ sha256 = hashlib.sha256(f"{object_name}_{file}".encode()).hexdigest()
+
+ try:
+ material_info = common.extract_materials_from_usd(
+ usd_file, "vegetation"
+ )
+ category = material_info.get("category", default_class)
+ except Exception:
+ category = default_class
+
+ metadata.append(
+ {
+ "sha256": sha256,
+ "local_path": usd_file,
+ "original_name": f"{object_name}/{file}",
+ "aesthetic_score": 1.0,
+ "rendered": False,
+ "class": category,
+ "dataset": "vegetation",
+ }
+ )
+
+ return metadata
+
+
+def get_metadata(
+ simready_dir=None,
+ commercial_dir=None,
+ residential_dir=None,
+ vegetation_dir=None,
+ output_dir=None,
+ asset_info_path=None,
+ include_datasets="simready,commercial,residential,vegetation",
+ seed=42,
+ default_class="unknown",
+ skip_split=False,
+ **kwargs,
+):
+
+ set_seeds(seed)
+
+ if simready_dir is None:
+ simready_dir = SIMREADY_PROPS_DIR
+ if commercial_dir is None:
+ commercial_dir = COMMERCIAL_BASE_DIR
+ if residential_dir is None:
+ residential_dir = RESIDENTIAL_BASE_DIR
+ if vegetation_dir is None:
+ vegetation_dir = VEGETATION_BASE_DIR
+ if asset_info_path is None:
+ asset_info_path = SIMREADY_ASSET_INFO_PATH
+
+ datasets = [d.strip() for d in include_datasets.split(",")]
+
+ metadata = []
+
+ if "simready" in datasets:
+ print(f"Processing SimReady dataset from {simready_dir}")
+ simready_metadata = get_simready_metadata(
+ simready_dir, asset_info_path, default_class
+ )
+ metadata.extend(simready_metadata)
+ print(f"Added {len(simready_metadata)} items from SimReady dataset")
+
+ if "commercial" in datasets:
+ print(f"Processing Commercial dataset from {commercial_dir}")
+ commercial_metadata = get_commercial_metadata(commercial_dir)
+ metadata.extend(commercial_metadata)
+ print(f"Added {len(commercial_metadata)} items from Commercial dataset")
+
+ if "residential" in datasets:
+ print(f"Processing Residential dataset from {residential_dir}")
+ residential_metadata = get_residential_metadata(residential_dir)
+ metadata.extend(residential_metadata)
+ print(f"Added {len(residential_metadata)} items from Residential dataset")
+
+ if "vegetation" in datasets:
+ print(f"Processing Vegetation dataset from {vegetation_dir}")
+ vegetation_metadata = get_vegetation_metadata(vegetation_dir)
+ metadata.extend(vegetation_metadata)
+ print(f"Added {len(vegetation_metadata)} items from Vegetation dataset")
+
+ df = pd.DataFrame(metadata)
+
+ if df.empty:
+ print("Warning: No metadata collected from any dataset")
+ return df
+
+ class_counts = df["class"].value_counts()
+ print("\nClass distribution in combined dataset:")
+ for class_name, count in class_counts.items():
+ print(f" - {class_name}: {count} ({count/len(df)*100:.1f}%)")
+
+ dataset_counts = df["dataset"].value_counts()
+ print("\nDataset distribution:")
+ for dataset_name, count in dataset_counts.items():
+ print(f" - {dataset_name}: {count} ({count/len(df)*100:.1f}%)")
+
+ if not skip_split:
+ df = split_dataset(df, seed=seed)
+ else:
+ print("Skipping dataset splitting as requested")
+ df["split"] = "train"
+
+ if output_dir:
+ os.makedirs(output_dir, exist_ok=True)
+
+ df.to_csv(os.path.join(output_dir, "metadata.csv"), index=False)
+
+ splits_dir = os.path.join(output_dir, "splits")
+ os.makedirs(splits_dir, exist_ok=True)
+
+ for split in ["train", "val", "test"]:
+ split_df = df[df["split"] == split]
+ if not split_df.empty:
+ split_df.to_csv(os.path.join(splits_dir, f"{split}.csv"), index=False)
+
+ class_stats = df.groupby(["class", "split"]).size().unstack(fill_value=0)
+ class_stats.to_csv(os.path.join(output_dir, "class_distribution.csv"))
+
+ dataset_stats = df.groupby(["dataset", "split"]).size().unstack(fill_value=0)
+ dataset_stats.to_csv(os.path.join(output_dir, "dataset_distribution.csv"))
+
+ return df
+
+
+def foreach_instance(metadata, output_dir, func, max_workers=8, desc="Processing"):
+ from concurrent.futures import ThreadPoolExecutor
+ from tqdm import tqdm
+ import pandas as pd
+
+ results = []
+
+ with ThreadPoolExecutor(max_workers=max_workers) as executor:
+ futures = []
+ for _, row in metadata.iterrows():
+ sha256 = row["sha256"]
+ local_path = row["local_path"]
+ dataset = row.get("dataset", "unknown")
+
+ futures.append(executor.submit(func, local_path, sha256, dataset))
+
+ for future in tqdm(futures, desc=desc, total=len(futures)):
+ try:
+ result = future.result()
+ if result is not None:
+ results.append(result)
+ except Exception as e:
+ print(f"Error in worker: {e}")
+
+ return pd.DataFrame.from_records(results)
diff --git a/deps/vomp/dataset_toolkits/datasets/simready.py b/deps/vomp/dataset_toolkits/datasets/simready.py
new file mode 100644
index 0000000000000000000000000000000000000000..27e2745c46626ed567d024cd32a0fd59bdcb1c9d
--- /dev/null
+++ b/deps/vomp/dataset_toolkits/datasets/simready.py
@@ -0,0 +1,297 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import os
+import json
+import pandas as pd
+import numpy as np
+import hashlib
+import random
+from glob import glob
+from sklearn.model_selection import train_test_split
+
+
+def set_seeds(seed=42):
+ random.seed(seed)
+ np.random.seed(seed)
+
+
+def add_args(parser):
+ parser.add_argument(
+ "--simready_dir",
+ type=str,
+ default="datasets/raw/simready/common_assets/props",
+ help="Path to the SimReady props directory",
+ )
+ parser.add_argument(
+ "--asset_info_path",
+ type=str,
+ default="datasets/raw/simready/asset_info.json",
+ help="Path to the SimReady asset_info.json file",
+ )
+ parser.add_argument(
+ "--seed",
+ type=int,
+ default=42,
+ help="Random seed for reproducibility",
+ )
+ parser.add_argument(
+ "--default_class",
+ type=str,
+ default="unknown",
+ help="Default class label to use when class information is not available",
+ )
+
+
+def get_asset_class_mapping(asset_info_path):
+ if not os.path.exists(asset_info_path):
+ print(f"Warning: Asset info file not found at {asset_info_path}")
+ return {}
+
+ try:
+ with open(asset_info_path, "r") as f:
+ asset_info = json.load(f)
+
+ asset_class_mapping = {}
+ for asset in asset_info:
+ simple_name = asset.get("Simple Name")
+ if simple_name and "Labels" in asset and "Class" in asset["Labels"]:
+ asset_class_mapping[simple_name] = asset["Labels"]["Class"]
+
+ print(f"Loaded class information for {len(asset_class_mapping)} assets")
+ return asset_class_mapping
+ except Exception as e:
+ print(f"Error loading asset info: {e}")
+ return {}
+
+
+def split_dataset(metadata, seed=42):
+
+ np.random.seed(seed)
+ random.seed(seed)
+
+ metadata_copy = metadata.copy()
+
+ metadata_copy["split"] = "train"
+
+ classes = metadata_copy["class"].unique()
+
+ large_classes = []
+ small_classes = []
+
+ for cls in classes:
+
+ cls_indices = metadata_copy[metadata_copy["class"] == cls].index.tolist()
+ if len(cls_indices) >= 10:
+ large_classes.append(cls)
+ else:
+ small_classes.append(cls)
+
+ for cls in large_classes:
+ cls_indices = metadata_copy[metadata_copy["class"] == cls].index.tolist()
+ random.shuffle(cls_indices)
+
+ n_samples = len(cls_indices)
+ n_train = int(0.8 * n_samples)
+ n_val = int(0.1 * n_samples)
+
+ train_indices = cls_indices[:n_train]
+ val_indices = cls_indices[n_train : n_train + n_val]
+ test_indices = cls_indices[n_train + n_val :]
+
+ metadata_copy.loc[train_indices, "split"] = "train"
+ metadata_copy.loc[val_indices, "split"] = "val"
+ metadata_copy.loc[test_indices, "split"] = "test"
+
+ total_samples = len(metadata_copy)
+ goal_train = int(0.8 * total_samples)
+ goal_val = int(0.1 * total_samples)
+ goal_test = total_samples - goal_train - goal_val
+
+ current_train = (metadata_copy["split"] == "train").sum()
+ current_val = (metadata_copy["split"] == "val").sum()
+ current_test = (metadata_copy["split"] == "test").sum()
+
+ small_indices = []
+ for cls in small_classes:
+ cls_indices = metadata_copy[metadata_copy["class"] == cls].index.tolist()
+ small_indices.extend(cls_indices)
+
+ random.shuffle(small_indices)
+
+ need_train = max(0, goal_train - current_train)
+ need_val = max(0, goal_val - current_val)
+ need_test = max(0, goal_test - current_test)
+
+ idx = 0
+ while idx < len(small_indices):
+ if need_train > 0:
+ metadata_copy.loc[small_indices[idx], "split"] = "train"
+ need_train -= 1
+ idx += 1
+ elif need_val > 0:
+ metadata_copy.loc[small_indices[idx], "split"] = "val"
+ need_val -= 1
+ idx += 1
+ elif need_test > 0:
+ metadata_copy.loc[small_indices[idx], "split"] = "test"
+ need_test -= 1
+ idx += 1
+ else:
+
+ metadata_copy.loc[small_indices[idx:], "split"] = "train"
+ break
+
+ train_count = (metadata_copy["split"] == "train").sum()
+ val_count = (metadata_copy["split"] == "val").sum()
+ test_count = (metadata_copy["split"] == "test").sum()
+
+ print(
+ f"Dataset split: Train: {train_count} ({train_count/len(metadata_copy)*100:.1f}%), "
+ f"Val: {val_count} ({val_count/len(metadata_copy)*100:.1f}%), "
+ f"Test: {test_count} ({test_count/len(metadata_copy)*100:.1f}%)"
+ )
+
+ if small_classes:
+ print("\nSmall class distribution across splits:")
+ for cls in small_classes:
+ cls_data = metadata_copy[metadata_copy["class"] == cls]
+ cls_train = (cls_data["split"] == "train").sum()
+ cls_val = (cls_data["split"] == "val").sum()
+ cls_test = (cls_data["split"] == "test").sum()
+ cls_total = len(cls_data)
+ print(
+ f" - {cls} (total {cls_total}): Train: {cls_train}, Val: {cls_val}, Test: {cls_test}"
+ )
+
+ return metadata_copy
+
+
+def get_metadata(
+ simready_dir=None,
+ output_dir=None,
+ asset_info_path=None,
+ seed=42,
+ default_class="unknown",
+ skip_split=False,
+ **kwargs,
+):
+
+ set_seeds(seed)
+
+ if simready_dir is None:
+ simready_dir = "datasets/raw/simready/common_assets/props"
+
+ if asset_info_path is None:
+ asset_info_path = "datasets/raw/simready/asset_info.json"
+
+ asset_class_mapping = get_asset_class_mapping(asset_info_path)
+
+ prop_dirs = [
+ d
+ for d in os.listdir(simready_dir)
+ if os.path.isdir(os.path.join(simready_dir, d))
+ ]
+
+ metadata = []
+
+ for prop_name in prop_dirs:
+ prop_dir = os.path.join(simready_dir, prop_name)
+
+ usd_files = glob(os.path.join(prop_dir, "*.usd"))
+ if not usd_files:
+ continue
+
+ inst_base_files = [f for f in usd_files if "_inst_base.usd" in f]
+ base_files = [f for f in usd_files if "_base.usd" in f]
+
+ if inst_base_files:
+ usd_file = inst_base_files[0]
+ elif base_files:
+ usd_file = base_files[0]
+ else:
+ usd_file = usd_files[0]
+
+ sha256 = hashlib.sha256(prop_name.encode()).hexdigest()
+
+ prop_class = asset_class_mapping.get(prop_name, default_class)
+
+ metadata.append(
+ {
+ "sha256": sha256,
+ "local_path": usd_file,
+ "original_name": prop_name,
+ "aesthetic_score": 1.0,
+ "rendered": False,
+ "class": prop_class,
+ }
+ )
+
+ df = pd.DataFrame(metadata)
+
+ class_counts = df["class"].value_counts()
+ print("\nClass distribution in dataset:")
+ for class_name, count in class_counts.items():
+ print(f" - {class_name}: {count} ({count/len(df)*100:.1f}%)")
+
+ if not skip_split:
+ df = split_dataset(df, seed=seed)
+ else:
+ print("Skipping dataset splitting as requested")
+ df["split"] = "train"
+
+ if output_dir:
+ os.makedirs(output_dir, exist_ok=True)
+
+ df.to_csv(os.path.join(output_dir, "metadata.csv"), index=False)
+
+ splits_dir = os.path.join(output_dir, "splits")
+ os.makedirs(splits_dir, exist_ok=True)
+
+ for split in ["train", "val", "test"]:
+ split_df = df[df["split"] == split]
+ if not split_df.empty:
+ split_df.to_csv(os.path.join(splits_dir, f"{split}.csv"), index=False)
+
+ class_stats = df.groupby(["class", "split"]).size().unstack(fill_value=0)
+ class_stats.to_csv(os.path.join(output_dir, "class_distribution.csv"))
+
+ return df
+
+
+def foreach_instance(metadata, output_dir, func, max_workers=8, desc="Processing"):
+ from concurrent.futures import ThreadPoolExecutor
+ from tqdm import tqdm
+ import pandas as pd
+
+ results = []
+
+ with ThreadPoolExecutor(max_workers=max_workers) as executor:
+ futures = []
+ for _, row in metadata.iterrows():
+ sha256 = row["sha256"]
+ local_path = row["local_path"]
+ futures.append(executor.submit(func, local_path, sha256))
+
+ for future in tqdm(futures, desc=desc, total=len(futures)):
+ try:
+ result = future.result()
+ if result is not None:
+ results.append(result)
+ except Exception as e:
+ print(f"Error in worker: {e}")
+
+ return pd.DataFrame.from_records(results)
diff --git a/deps/vomp/dataset_toolkits/extract_feature.py b/deps/vomp/dataset_toolkits/extract_feature.py
new file mode 100644
index 0000000000000000000000000000000000000000..32521998411455e25671486f70d11013587c5217
--- /dev/null
+++ b/deps/vomp/dataset_toolkits/extract_feature.py
@@ -0,0 +1,273 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import copy
+import sys
+import json
+import importlib
+import argparse
+import torch
+import torch.nn.functional as F
+import numpy as np
+import pandas as pd
+import utils3d
+from tqdm import tqdm
+from easydict import EasyDict as edict
+from torchvision import transforms
+from PIL import Image
+
+torch.set_grad_enabled(False)
+
+
+def get_data(frames, sha256):
+ valid_data = []
+
+ for view in frames:
+ image_path = os.path.join(opt.output_dir, "renders", sha256, view["file_path"])
+ try:
+ # Check if file exists before trying to open it
+ if not os.path.exists(image_path):
+ print(f"Warning: Image file {image_path} not found, skipping")
+ continue
+
+ image = Image.open(image_path)
+ except Exception as e:
+ print(f"Error loading image {image_path}: {e}")
+ continue
+
+ try:
+ image = image.resize((518, 518), Image.Resampling.LANCZOS)
+ image = np.array(image).astype(np.float32) / 255
+ image = image[:, :, :3] * image[:, :, 3:]
+ image = torch.from_numpy(image).permute(2, 0, 1).float()
+
+ c2w = torch.tensor(view["transform_matrix"])
+ c2w[:3, 1:3] *= -1
+ extrinsics = torch.inverse(c2w)
+ fov = view["camera_angle_x"]
+ intrinsics = utils3d.torch.intrinsics_from_fov_xy(
+ torch.tensor(fov), torch.tensor(fov)
+ )
+
+ valid_data.append(
+ {"image": image, "extrinsics": extrinsics, "intrinsics": intrinsics}
+ )
+ except Exception as e:
+ print(f"Error processing image {image_path}: {e}")
+ continue
+
+ if len(valid_data) == 0:
+ print(f"Warning: No valid images found for {sha256}")
+ else:
+ print(f"Loaded {len(valid_data)}/{len(frames)} valid images for {sha256}")
+
+ return valid_data
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--output_dir", type=str, required=True, help="Directory to save the metadata"
+ )
+ parser.add_argument(
+ "--filter_low_aesthetic_score",
+ type=float,
+ default=None,
+ help="Filter objects with aesthetic score lower than this value",
+ )
+ parser.add_argument(
+ "--model",
+ type=str,
+ default="dinov2_vitl14_reg",
+ help="Feature extraction model",
+ )
+ parser.add_argument(
+ "--instances", type=str, default=None, help="Instances to process"
+ )
+ parser.add_argument("--batch_size", type=int, default=16)
+ parser.add_argument("--rank", type=int, default=0)
+ parser.add_argument("--world_size", type=int, default=1)
+ parser.add_argument(
+ "--force",
+ action="store_true",
+ help="Force feature extraction even if feature files already exist",
+ )
+ opt = parser.parse_args()
+ opt = edict(vars(opt))
+
+ feature_name = opt.model
+ os.makedirs(os.path.join(opt.output_dir, "features", feature_name), exist_ok=True)
+
+ # load model
+ dinov2_model = torch.hub.load("facebookresearch/dinov2", opt.model)
+ dinov2_model.eval().cuda()
+ transform = transforms.Compose(
+ [
+ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+ ]
+ )
+ n_patch = 518 // 14
+
+ # get file list
+ if os.path.exists(os.path.join(opt.output_dir, "metadata.csv")):
+ metadata = pd.read_csv(os.path.join(opt.output_dir, "metadata.csv"))
+ else:
+ raise ValueError("metadata.csv not found")
+ if opt.instances is not None:
+ with open(opt.instances, "r") as f:
+ instances = f.read().splitlines()
+ metadata = metadata[metadata["sha256"].isin(instances)]
+ else:
+ if opt.filter_low_aesthetic_score is not None:
+ metadata = metadata[
+ metadata["aesthetic_score"] >= opt.filter_low_aesthetic_score
+ ]
+ if f"feature_{feature_name}" in metadata.columns and not opt.force:
+ metadata = metadata[metadata[f"feature_{feature_name}"] == False]
+ metadata = metadata[metadata["voxelized"] == True]
+ metadata = metadata[metadata["rendered"] == True]
+
+ start = len(metadata) * opt.rank // opt.world_size
+ end = len(metadata) * (opt.rank + 1) // opt.world_size
+ metadata = metadata[start:end]
+ records = []
+
+ # filter out objects that are already processed
+ sha256s = list(metadata["sha256"].values)
+ if not opt.force:
+ for sha256 in copy.copy(sha256s):
+ if os.path.exists(
+ os.path.join(opt.output_dir, "features", feature_name, f"{sha256}.npz")
+ ):
+ records.append({"sha256": sha256, f"feature_{feature_name}": True})
+ sha256s.remove(sha256)
+ else:
+ print(
+ f"Force mode enabled. Processing all {len(sha256s)} objects regardless of existing features."
+ )
+
+ # filter out objects that don't have voxel files
+ initial_count = len(sha256s)
+ sha256s_with_voxels = []
+ for sha256 in sha256s:
+ voxel_path = os.path.join(opt.output_dir, "voxels", f"{sha256}.ply")
+ if os.path.exists(voxel_path):
+ sha256s_with_voxels.append(sha256)
+ else:
+ print(f"Skipping {sha256}: voxel file not found at {voxel_path}")
+
+ sha256s = sha256s_with_voxels
+ print(f"Filtered from {initial_count} to {len(sha256s)} objects with voxel files")
+
+ # extract features
+ for sha256 in tqdm(sha256s, desc="Extracting features"):
+ try:
+ # Load data
+ with open(
+ os.path.join(opt.output_dir, "renders", sha256, "transforms.json"),
+ "r",
+ ) as f:
+ metadata_json = json.load(f)
+ frames = metadata_json["frames"]
+ data = get_data(frames, sha256)
+
+ if len(data) == 0:
+ print(f"Skipping {sha256}: no valid image data")
+ continue
+
+ # Apply transform to images
+ for datum in data:
+ datum["image"] = transform(datum["image"])
+
+ # Load positions
+ positions = utils3d.io.read_ply(
+ os.path.join(opt.output_dir, "voxels", f"{sha256}.ply")
+ )[0]
+ positions = torch.from_numpy(positions).float().cuda()
+ indices = ((positions + 0.5) * 64).long()
+ # Clamp indices to valid range [0, 63] to handle floating point precision issues
+ indices = torch.clamp(indices, 0, 63)
+
+ n_views = len(data)
+ N = positions.shape[0]
+ pack = {
+ "indices": indices.cpu().numpy().astype(np.uint8),
+ }
+
+ patchtokens_lst = []
+ uv_lst = []
+
+ # Process in batches
+ for i in range(0, n_views, opt.batch_size):
+ batch_data = data[i : i + opt.batch_size]
+ bs = len(batch_data)
+ batch_images = torch.stack([d["image"] for d in batch_data]).cuda()
+ batch_extrinsics = torch.stack(
+ [d["extrinsics"] for d in batch_data]
+ ).cuda()
+ batch_intrinsics = torch.stack(
+ [d["intrinsics"] for d in batch_data]
+ ).cuda()
+ features = dinov2_model(batch_images, is_training=True)
+ uv = (
+ utils3d.torch.project_cv(
+ positions, batch_extrinsics, batch_intrinsics
+ )[0]
+ * 2
+ - 1
+ )
+ patchtokens = (
+ features["x_prenorm"][:, dinov2_model.num_register_tokens + 1 :]
+ .permute(0, 2, 1)
+ .reshape(bs, 1024, n_patch, n_patch)
+ )
+ patchtokens_lst.append(patchtokens)
+ uv_lst.append(uv)
+
+ patchtokens = torch.cat(patchtokens_lst, dim=0)
+ uv = torch.cat(uv_lst, dim=0)
+
+ # Save features
+ pack["patchtokens"] = (
+ F.grid_sample(
+ patchtokens,
+ uv.unsqueeze(1),
+ mode="bilinear",
+ align_corners=False,
+ )
+ .squeeze(2)
+ .permute(0, 2, 1)
+ .cpu()
+ .numpy()
+ )
+ pack["patchtokens"] = np.mean(pack["patchtokens"], axis=0).astype(
+ np.float16
+ )
+ save_path = os.path.join(
+ opt.output_dir, "features", feature_name, f"{sha256}.npz"
+ )
+ np.savez_compressed(save_path, **pack)
+ records.append({"sha256": sha256, f"feature_{feature_name}": True})
+
+ except Exception as e:
+ print(f"Error processing {sha256}: {e}")
+ continue
+
+ records = pd.DataFrame.from_records(records)
+ records.to_csv(
+ os.path.join(opt.output_dir, f"feature_{feature_name}_{opt.rank}.csv"),
+ index=False,
+ )
diff --git a/deps/vomp/dataset_toolkits/latent_space/analyze_data_distribution.py b/deps/vomp/dataset_toolkits/latent_space/analyze_data_distribution.py
new file mode 100644
index 0000000000000000000000000000000000000000..0938d36e4ecad34d7067767c5becfabff32715dc
--- /dev/null
+++ b/deps/vomp/dataset_toolkits/latent_space/analyze_data_distribution.py
@@ -0,0 +1,111 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+from pathlib import Path
+
+df = pd.read_csv("datasets/latent_space/materials.csv")
+
+print("Data shape:", df.shape)
+print("\nColumn names:", df.columns.tolist())
+
+
+for col in ["youngs_modulus", "poisson_ratio", "density"]:
+ print(f"\n{col}:")
+ print(f" Min: {df[col].min():.2e}")
+ print(f" Max: {df[col].max():.2e}")
+ print(f" Mean: {df[col].mean():.2e}")
+ print(f" Median: {df[col].median():.2e}")
+ print(f" Std: {df[col].std():.2e}")
+
+ Q1 = df[col].quantile(0.25)
+ Q3 = df[col].quantile(0.75)
+ IQR = Q3 - Q1
+ lower_bound = Q1 - 1.5 * IQR
+ upper_bound = Q3 + 1.5 * IQR
+
+ outliers = df[(df[col] < lower_bound) | (df[col] > upper_bound)]
+ print(
+ f" Outliers (IQR method): {len(outliers)} ({len(outliers)/len(df)*100:.1f}%)"
+ )
+
+ if col in ["youngs_modulus", "density"]:
+ log_vals = np.log10(df[col])
+ print(f" Log10 range: [{log_vals.min():.2f}, {log_vals.max():.2f}]")
+ print(
+ f" Log10 span: {log_vals.max() - log_vals.min():.2f} orders of magnitude"
+ )
+
+
+print("\n\nMaterials with extreme Young's modulus (< 1e7 Pa):")
+low_E = df[df["youngs_modulus"] < 1e7]
+if len(low_E) > 0:
+ material_counts = low_E["material_name"].value_counts().head(10)
+ for mat, count in material_counts.items():
+ print(f" {mat}: {count} samples")
+
+print("\n\nMaterials with extreme density (< 100 kg/mยณ):")
+low_rho = df[df["density"] < 100]
+if len(low_rho) > 0:
+ material_counts = low_rho["material_name"].value_counts().head(10)
+ for mat, count in material_counts.items():
+ print(f" {mat}: {count} samples")
+
+
+print("\n\nPercentile analysis:")
+percentiles = [1, 5, 10, 25, 50, 75, 90, 95, 99]
+for col in ["youngs_modulus", "poisson_ratio", "density"]:
+ print(f"\n{col} percentiles:")
+ for p in percentiles:
+ val = df[col].quantile(p / 100)
+ print(f" {p}%: {val:.2e}")
+
+
+print("\n\nCreating filtered dataset...")
+
+filtered_df = df[
+ (df["youngs_modulus"] >= 1e5)
+ & (df["youngs_modulus"] <= 1e12)
+ & (df["density"] >= 100)
+ & (df["density"] <= 20000)
+ & (df["poisson_ratio"] >= 0.0)
+ & (df["poisson_ratio"] <= 0.49)
+]
+
+print(f"Original size: {len(df)}")
+print(f"Filtered size: {len(filtered_df)}")
+print(
+ f"Removed: {len(df) - len(filtered_df)} ({(len(df) - len(filtered_df))/len(df)*100:.1f}%)"
+)
+
+
+print("\nRanges in filtered dataset (only Poisson ratio filtering):")
+for col in ["youngs_modulus", "poisson_ratio", "density"]:
+ print(f"\n{col}:")
+ print(f" Min: {filtered_df[col].min():.2e}")
+ print(f" Max: {filtered_df[col].max():.2e}")
+ print(f" Range span: {filtered_df[col].max() - filtered_df[col].min():.2e}")
+
+ if col in ["youngs_modulus", "density"]:
+ log_min = np.log10(filtered_df[col].min())
+ log_max = np.log10(filtered_df[col].max())
+ print(f" Log10 range: [{log_min:.2f}, {log_max:.2f}]")
+ print(f" Orders of magnitude: {log_max - log_min:.2f}")
+
+filtered_df.to_csv("datasets/latent_space/materials_filtered.csv", index=False)
+print("\nSaved filtered dataset to materials_filtered.csv")
diff --git a/deps/vomp/dataset_toolkits/latent_space/make_csv.py b/deps/vomp/dataset_toolkits/latent_space/make_csv.py
new file mode 100644
index 0000000000000000000000000000000000000000..62c438132071b1ced400256e964f60d8410feb6c
--- /dev/null
+++ b/deps/vomp/dataset_toolkits/latent_space/make_csv.py
@@ -0,0 +1,411 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import argparse
+import json
+import csv
+import random
+from pathlib import Path
+from typing import Tuple, Set, List
+import math
+from dataset_toolkits.material_objects.vlm_annotations.utils.utils import (
+ parse_numerical_range_str,
+)
+
+
+def parse_args() -> Path:
+ parser = argparse.ArgumentParser(
+ description="Generate a materials.csv file from material_ranges.csv in the provided directory."
+ )
+ parser.add_argument(
+ "directory",
+ type=str,
+ help="Path to the directory containing material_ranges.csv",
+ )
+ args = parser.parse_args()
+ directory = Path(args.directory).expanduser().resolve()
+
+ if not directory.is_dir():
+ parser.error(f"Provided path '{directory}' is not a directory.")
+
+ return directory
+
+
+def read_dataset(json_path: Path):
+ try:
+ with json_path.open("r", encoding="utf-8") as f:
+ return json.load(f)
+ except FileNotFoundError:
+ raise FileNotFoundError(
+ f"Could not find '{json_path}'. Ensure the directory contains the file."
+ )
+
+
+def extract_unique_rows(
+ dataset, unique_triplets: Set[Tuple[float, float, float]] | None = None
+) -> tuple[list, Set[Tuple[float, float, float]]]:
+ if unique_triplets is None:
+ unique_triplets = set()
+
+ rows: List[Tuple[str, float, float, float]] = []
+
+ for obj in dataset:
+ segments = obj.get("segments", {})
+ for seg_key, seg_data in segments.items():
+ try:
+ youngs = float(seg_data["youngs_modulus"])
+ poisson = float(seg_data["poissons_ratio"])
+ density = float(seg_data["density"])
+ except (KeyError, ValueError, TypeError):
+ continue
+
+ if youngs <= 0 or youngs > 1e13:
+ print(
+ f"WARNING: Skipping material with invalid Young's modulus: {youngs}"
+ )
+ continue
+
+ if poisson < -1.0 or poisson > 0.5:
+ print(
+ f"WARNING: Skipping material with invalid Poisson's ratio: {poisson}"
+ )
+ continue
+
+ if density <= 0 or density > 50000:
+ print(f"WARNING: Skipping material with invalid density: {density}")
+ continue
+
+ triplet = (youngs, poisson, density)
+ if triplet in unique_triplets:
+ continue
+
+ unique_triplets.add(triplet)
+ material_name = seg_data.get("name", seg_key)
+ rows.append((material_name, youngs, poisson, density))
+
+ return rows, unique_triplets
+
+
+def sample_ranges(
+ csv_path: Path,
+ unique_triplets: Set[Tuple[float, float, float]],
+ min_samples_per_material: int = 100,
+ max_samples_per_material: int = 2500,
+ target_total_samples: int = 100_000,
+) -> list:
+ if not csv_path.exists():
+
+ return []
+
+ parsed_rows: list[dict] = []
+ dynamic_indices: list[int] = []
+
+ with csv_path.open("r", encoding="utf-8") as f:
+ lines = f.readlines()
+
+ header = lines[0].strip().split(",")
+
+ for idx, line in enumerate(lines[1:], 0):
+
+ parts = []
+ current = ""
+ in_brackets = False
+ for char in line.strip() + ",":
+ if char == "," and not in_brackets:
+ parts.append(current)
+ current = ""
+ else:
+ if char == "[":
+ in_brackets = True
+ elif char == "]":
+ in_brackets = False
+ current += char
+
+ if len(parts) < 4:
+ print(f"WARNING: Line {idx+1} has incorrect format: {line.strip()}")
+ continue
+
+ material_name = parts[0].strip().strip('"')
+
+ y_range_str = parts[1].strip().strip('"')
+ p_range_str = parts[2].strip().strip('"')
+ d_range_str = parts[3].strip().strip('"')
+
+ try:
+ y_low, y_high = parse_numerical_range_str(y_range_str)
+ p_low, p_high = parse_numerical_range_str(p_range_str)
+ d_low, d_high = parse_numerical_range_str(d_range_str)
+ except ValueError as e:
+ print(
+ f"WARNING: Error parsing ranges for {material_name} on line {idx+1}: {e} - Skipping material."
+ )
+ continue
+
+ y_low *= 1e9
+ y_high *= 1e9
+
+ y_low = max(1e6, min(y_low, 1e13))
+ y_high = max(y_low, min(y_high, 1e13))
+
+ p_low = max(-0.999, min(p_low, 0.499))
+ p_high = max(p_low, min(p_high, 0.499))
+
+ d_low = max(10.0, min(d_low, 50000.0))
+ d_high = max(d_low, min(d_high, 50000.0))
+
+ y_has_range = abs(y_high - y_low) > 1e-6
+ p_has_range = abs(p_high - p_low) > 1e-6
+ d_has_range = abs(d_high - d_low) > 1e-6
+
+ has_range = y_has_range or p_has_range or d_has_range
+
+ y_width = max(y_high - y_low, 1.0) if y_has_range else 1.0
+ p_width = max(p_high - p_low, 0.001) if p_has_range else 0.001
+ d_width = max(d_high - d_low, 1.0) if d_has_range else 1.0
+
+ y_width_norm = y_width / 1e9
+
+ volume = y_width_norm * p_width * d_width
+
+ if has_range:
+ dynamic_indices.append(idx)
+
+ parsed_rows.append(
+ {
+ "material_name": material_name,
+ "y_low": y_low,
+ "y_high": y_high,
+ "p_low": p_low,
+ "p_high": p_high,
+ "d_low": d_low,
+ "d_high": d_high,
+ "has_range": has_range,
+ "y_has_range": y_has_range,
+ "p_has_range": p_has_range,
+ "d_has_range": d_has_range,
+ "volume": volume,
+ }
+ )
+
+ if not parsed_rows:
+ return []
+
+ print(
+ f"Found {len(dynamic_indices)} materials with ranges out of {len(parsed_rows)} total"
+ )
+
+ fixed_count = len(parsed_rows) - len(dynamic_indices)
+ print(f"Number of materials with fixed values: {fixed_count}")
+
+ if dynamic_indices:
+ print("\nExample materials with ranges:")
+ for i in range(min(5, len(dynamic_indices))):
+ idx = dynamic_indices[i]
+ info = parsed_rows[idx]
+ ranges_info = []
+ if info["y_has_range"]:
+ ranges_info.append(
+ f"Young's: {info['y_low']/1e9:.3f}-{info['y_high']/1e9:.3f} GPa"
+ )
+ if info["p_has_range"]:
+ ranges_info.append(
+ f"Poisson's: {info['p_low']:.3f}-{info['p_high']:.3f}"
+ )
+ if info["d_has_range"]:
+ ranges_info.append(
+ f"Density: {info['d_low']:.1f}-{info['d_high']:.1f} kg/mยณ"
+ )
+
+ print(
+ f" {info['material_name']}: {', '.join(ranges_info)} (volume: {info['volume']:.4f})"
+ )
+
+ total_volume = sum(parsed_rows[idx]["volume"] for idx in dynamic_indices)
+ print(f"\nTotal parameter space volume: {total_volume:.4f}")
+
+ volume_scale_factor = 13.0
+
+ samples_per_material = {}
+
+ for idx in dynamic_indices:
+ volume_ratio = parsed_rows[idx]["volume"] / total_volume
+ proportional_samples = max(
+ math.ceil(target_total_samples * volume_ratio * volume_scale_factor),
+ min_samples_per_material,
+ )
+
+ samples_per_material[idx] = min(proportional_samples, max_samples_per_material)
+
+ fixed_total = 0
+
+ dynamic_total = sum(samples_per_material.values())
+ total_planned = dynamic_total + fixed_total
+
+ print(f"\nSampling strategy (scaled by {volume_scale_factor}x):")
+ print(f" Minimum samples per material with ranges: {min_samples_per_material}")
+ print(f" Maximum samples per material: {max_samples_per_material}")
+ print(f" Planned total samples: {total_planned}")
+
+ sorted_materials = sorted(
+ [
+ (
+ idx,
+ parsed_rows[idx]["material_name"],
+ parsed_rows[idx]["volume"],
+ samples_per_material.get(idx, 1) if idx in dynamic_indices else 1,
+ )
+ for idx in range(len(parsed_rows))
+ ],
+ key=lambda x: x[2],
+ reverse=True,
+ )
+
+ print("\nTop 15 highest volume materials:")
+ for idx, name, volume, samples in sorted_materials[:15]:
+ if idx in dynamic_indices:
+ volume_percent = volume / total_volume * 100
+ print(
+ f" {name}: volume {volume:.4f} ({volume_percent:.2f}%), {samples} samples"
+ )
+ else:
+ print(f" {name}: fixed values, 1 sample")
+
+ rows: list[Tuple[str, float, float, float]] = []
+
+ def _add_triplet(material: str, y: float, p: float, d: float):
+
+ if y <= 0 or y > 1e13:
+ return False
+ if p < -1.0 or p > 0.5:
+ return False
+ if d <= 0 or d > 50000:
+ return False
+
+ triplet = (y, p, d)
+ if triplet in unique_triplets:
+ return False
+ unique_triplets.add(triplet)
+ rows.append((material, y, p, d))
+ return True
+
+ total_generated = 0
+ duplicate_avoidance_failures = 0
+
+ for idx, info in enumerate(parsed_rows):
+ if not info["has_range"]:
+ name = info["material_name"]
+ y_val = info["y_low"]
+ p_val = info["p_low"]
+ d_val = info["d_low"]
+
+ if _add_triplet(name, y_val, p_val, d_val):
+ total_generated += 1
+
+ print(f"Added {total_generated} materials with fixed values")
+
+ for idx in dynamic_indices:
+ info = parsed_rows[idx]
+ name = info["material_name"]
+ y_low, y_high = info["y_low"], info["y_high"]
+ p_low, p_high = info["p_low"], info["p_high"]
+ d_low, d_high = info["d_low"], info["d_high"]
+
+ required = samples_per_material.get(idx, 0)
+
+ report_progress = required > 100
+
+ attempts = 0
+ generated = 0
+
+ max_attempts = required * 50
+
+ if info["volume"] > 10.0:
+ max_attempts *= 2
+
+ if report_progress:
+ print(
+ f"Generating {required} samples for {name} (volume: {info['volume']:.4f})"
+ )
+
+ while generated < required and attempts < max_attempts:
+ attempts += 1
+
+ y_val = random.uniform(y_low, y_high) if info["y_has_range"] else y_low
+ p_val = random.uniform(p_low, p_high) if info["p_has_range"] else p_low
+ d_val = random.uniform(d_low, d_high) if info["d_has_range"] else d_low
+
+ y_val = round(y_val, 10)
+ p_val = round(p_val, 10)
+ d_val = round(d_val, 10)
+
+ if _add_triplet(name, y_val, p_val, d_val):
+ generated += 1
+ total_generated += 1
+ else:
+
+ duplicate_avoidance_failures += 1
+
+ if report_progress and generated > 0 and generated % 100 == 0:
+ print(f" Generated {generated}/{required} samples for {name}")
+
+ if required > 0 and report_progress:
+
+ success_rate = (generated / attempts) * 100 if attempts > 0 else 0
+ print(
+ f"Material {name}: Generated {generated}/{required} samples after {attempts} attempts (success rate: {success_rate:.1f}%)"
+ )
+
+ print(f"Successfully generated {len(rows)} unique material property combinations")
+ print(
+ f"Duplicate avoidance prevented {duplicate_avoidance_failures} potential duplicates"
+ )
+ return rows
+
+
+def write_csv(rows: list, csv_path: Path):
+ csv_path.parent.mkdir(parents=True, exist_ok=True)
+ with csv_path.open("w", newline="", encoding="utf-8") as csvfile:
+ writer = csv.writer(csvfile)
+ writer.writerow(["material_name", "youngs_modulus", "poisson_ratio", "density"])
+ for row in rows:
+ writer.writerow(row)
+
+
+def main():
+ directory = parse_args()
+ csv_path = directory / "materials.csv"
+
+ print("Generating materials data from ranges only (skipping JSON file)...")
+
+ unique_triplets = set()
+
+ ranges_csv_path = directory / "material_ranges.csv"
+ if not ranges_csv_path.exists():
+ print(f"ERROR: material_ranges.csv not found at {ranges_csv_path}")
+ return
+
+ sampled_rows = sample_ranges(ranges_csv_path, unique_triplets)
+
+ write_csv(sampled_rows, csv_path)
+
+ print(
+ f"materials.csv generated with {len(sampled_rows)} unique rows at '{csv_path}'."
+ )
+ print("All data generated from material_ranges.csv with validation applied.")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/deps/vomp/dataset_toolkits/material_objects/render_usd.py b/deps/vomp/dataset_toolkits/material_objects/render_usd.py
new file mode 100644
index 0000000000000000000000000000000000000000..0d4dbd57caf5f4c1700e3a9dc262b6e8e4611131
--- /dev/null
+++ b/deps/vomp/dataset_toolkits/material_objects/render_usd.py
@@ -0,0 +1,1176 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Clean USD rendering pipeline.
+
+This script extracts meshes and textures directly from USD files,
+similar to how Omniverse exports meshes. It does NOT search the filesystem
+for textures - all texture paths come from the USD shaders themselves.
+
+For vegetation datasets that use MDL materials, it parses the MDL files
+to extract texture references.
+"""
+
+import os
+import sys
+import json
+import re
+import argparse
+import tempfile
+import shutil
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+from collections import defaultdict
+from subprocess import call, DEVNULL
+
+import numpy as np
+import pandas as pd
+from pxr import Usd, UsdGeom, UsdShade, Sdf, Gf
+
+sys.path.append(os.path.dirname(os.path.dirname(__file__)))
+from utils import sphere_hammersley_sequence
+
+BLENDER_LINK = (
+ "https://download.blender.org/release/Blender3.0/blender-3.0.1-linux-x64.tar.xz"
+)
+BLENDER_INSTALLATION_PATH = "/tmp"
+BLENDER_PATH = f"{BLENDER_INSTALLATION_PATH}/blender-3.0.1-linux-x64/blender"
+
+
+class USDMaterialExtractor:
+ """
+ Extracts materials and textures directly from USD files.
+
+ This class reads shader inputs from USD prims and resolves texture paths
+ relative to the USD file. For MDL materials (used in vegetation), it
+ parses the MDL files to extract texture references.
+ """
+
+ # MDL texture patterns (for vegetation)
+ MDL_TEXTURE_PATTERNS = [
+ r'diffuse_texture:\s*texture_2d\("([^"]+)"',
+ r'normalmap_texture:\s*texture_2d\("([^"]+)"',
+ r'reflectionroughness_texture:\s*texture_2d\("([^"]+)"',
+ r'metallic_texture:\s*texture_2d\("([^"]+)"',
+ r'ORM_texture:\s*texture_2d\("([^"]+)"',
+ ]
+
+ def __init__(self, usd_path: str, verbose: bool = False):
+ self.usd_path = Path(usd_path).resolve()
+ self.usd_dir = self.usd_path.parent
+ self.verbose = verbose
+ self.stage = None
+
+ # Extracted data
+ self.materials = {} # material_path -> {input_name: texture_path}
+ self.meshes = {} # mesh_path -> {name, material, vertices, faces, uvs}
+ self.mesh_materials = {} # mesh_path -> material_path
+
+ def _log(self, msg: str):
+ if self.verbose:
+ print(msg)
+
+ def _resolve_texture_path(self, texture_path: str) -> Optional[Path]:
+ texture_path = texture_path.strip("@")
+
+ # Handle UDIM textures - replace Feed-forward, fine-grained, physically based volumetric material properties from Splats, Meshes, NeRFs, and more.
+
+
+
+
+ 
+