Vendored VoMP with local debug edits; single commit (history rewritten to drop binary blobs for HF Hub).

.gitattributes

@@ -0,0 +1,37 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.ply filter=lfs diff=lfs merge=lfs -text
+*.whl filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+venv/
+weights/

README.md

@@ -0,0 +1,15 @@
+---
+title: VoMP
+emoji: 🚀
+colorFrom: green
+colorTo: green
+sdk: gradio
+python_version: 3.12
+sdk_version: 6.2.0
+app_file: app.py
+pinned: true
+license: apache-2.0
+short_description: Volumetric physics materials for interactive worlds
+suggested_hardware: a100-large
+suggested_storage: medium
+---

app.py

@@ -0,0 +1,561 @@
+import glob
+import os
+import shutil
+import tempfile
+from typing import Dict, List, Optional, Tuple
+
+import gradio as gr
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import matplotlib.colors as mcolors
+from matplotlib.colorbar import ColorbarBase
+import numpy as np
+import spaces
+import torch
+from huggingface_hub import snapshot_download
+
+from vomp.inference import Vomp
+from vomp.inference.utils import save_materials
+
+NUM_VIEWS = 150
+PROPERTY_NAMES = ["youngs_modulus", "poissons_ratio", "density"]
+PROPERTY_DISPLAY_NAMES = {
+    "youngs_modulus": "Young's Modulus",
+    "poissons_ratio": "Poisson's Ratio",
+    "density": "Density",
+}
+
+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"
+
+EXAMPLES_DIR = "examples"
+
+model_id = "nvidia/PhysicalAI-Simulation-VoMP-Model"
+base_path = snapshot_download(repo_id=model_id, local_dir="weights")
+print(os.listdir(base_path))
+
+def _install_blender():
+    if not os.path.exists(BLENDER_PATH):
+        print("Installing Blender...")
+        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}"
+        )
+        print("Blender installed successfully!")
+
+
+def _is_gaussian_splat(file_path: str) -> bool:
+    if not file_path.lower().endswith(".ply"):
+        return False
+
+    try:
+        with open(file_path, "rb") as f:
+            header = b""
+            while True:
+                line = f.readline()
+                header += line
+                if b"end_header" in line:
+                    break
+                if len(header) > 10000:
+                    break
+
+            header_str = header.decode("utf-8", errors="ignore").lower()
+            gaussian_indicators = ["f_dc", "opacity", "scale_0", "rot_0"]
+            return any(indicator in header_str for indicator in gaussian_indicators)
+    except Exception:
+        return False
+
+
+def _setup_examples():
+    """Ensure examples directory exists."""
+    os.makedirs(EXAMPLES_DIR, exist_ok=True)
+
+
+_setup_examples()
+
+
+print("Loading VoMP model...")
+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",
+)
+print("VoMP model loaded successfully!")
+
+
+def _get_render_images(output_dir: str) -> List[str]:
+    renders_dir = os.path.join(output_dir, "renders")
+    if not os.path.exists(renders_dir):
+        return []
+    image_paths = sorted(glob.glob(os.path.join(renders_dir, "*.png")))
+    return image_paths
+
+
+def _create_colorbar(
+    data: np.ndarray, property_name: str, output_path: str, colormap: str = "viridis"
+) -> str:
+    fig, ax = plt.subplots(figsize=(6, 0.8))
+    fig.subplots_adjust(bottom=0.5)
+    ax.remove()
+
+    cmap = plt.cm.get_cmap(colormap)
+    norm = mcolors.Normalize(vmin=np.min(data), vmax=np.max(data))
+
+    cbar_ax = fig.add_axes([0.1, 0.4, 0.8, 0.35])
+    cb = ColorbarBase(cbar_ax, cmap=cmap, norm=norm, orientation="horizontal")
+    cb.ax.set_xlabel(
+        f"{PROPERTY_DISPLAY_NAMES.get(property_name, property_name)}", fontsize=10
+    )
+
+    plt.savefig(
+        output_path, dpi=150, bbox_inches="tight", facecolor="white", transparent=False
+    )
+    plt.close()
+    return output_path
+
+
+def _render_point_cloud_views(
+    coords: np.ndarray,
+    values: np.ndarray,
+    output_dir: str,
+    property_name: str,
+    colormap: str = "viridis",
+) -> List[str]:
+    vmin, vmax = np.min(values), np.max(values)
+    if vmax - vmin > 1e-10:
+        normalized = (values - vmin) / (vmax - vmin)
+    else:
+        normalized = np.zeros_like(values)
+
+    cmap = plt.cm.get_cmap(colormap)
+    colors = cmap(normalized)
+
+    views = [
+        (30, 45, "view1"),
+        (30, 135, "view2"),
+        (80, 45, "view3"),
+    ]
+
+    image_paths = []
+
+    for elev, azim, view_name in views:
+        fig = plt.figure(figsize=(6, 6), facecolor="#1a1a1a")
+        ax = fig.add_subplot(111, projection="3d", facecolor="#1a1a1a")
+
+        ax.scatter(
+            coords[:, 0],
+            coords[:, 1],
+            coords[:, 2],
+            c=colors,
+            s=15,
+            alpha=0.9,
+        )
+
+        ax.view_init(elev=elev, azim=azim)
+        ax.set_xlim([-0.6, 0.6])
+        ax.set_ylim([-0.6, 0.6])
+        ax.set_zlim([-0.6, 0.6])
+        ax.set_axis_off()
+        ax.set_box_aspect([1, 1, 1])
+
+        output_path = os.path.join(output_dir, f"{property_name}_{view_name}.png")
+        plt.savefig(
+            output_path,
+            dpi=150,
+            bbox_inches="tight",
+            facecolor="#1a1a1a",
+            edgecolor="none",
+        )
+        plt.close()
+
+        image_paths.append(output_path)
+
+    return image_paths
+
+
+def _create_material_visualizations(
+    material_file: str, output_dir: str
+) -> Dict[str, Tuple[List[str], str]]:
+    result = {}
+    data = np.load(material_file, allow_pickle=True)
+
+    if "voxel_data" in data:
+        voxel_data = data["voxel_data"]
+        coords = np.column_stack([voxel_data["x"], voxel_data["y"], voxel_data["z"]])
+        properties = {
+            "youngs_modulus": voxel_data["youngs_modulus"],
+            "poissons_ratio": voxel_data["poissons_ratio"],
+            "density": voxel_data["density"],
+        }
+    else:
+        if "voxel_coords_world" in data:
+            coords = data["voxel_coords_world"]
+        elif "query_coords_world" in data:
+            coords = data["query_coords_world"]
+        elif "coords" in data:
+            coords = data["coords"]
+        else:
+            print(f"Warning: No coordinate data found in {material_file}")
+            return result
+
+        properties = {}
+        property_mapping = {
+            "youngs_modulus": ["youngs_modulus", "young_modulus"],
+            "poissons_ratio": ["poissons_ratio", "poisson_ratio"],
+            "density": ["density"],
+        }
+        for prop_name, possible_names in property_mapping.items():
+            for name in possible_names:
+                if name in data:
+                    properties[prop_name] = data[name]
+                    break
+
+    center = (np.min(coords, axis=0) + np.max(coords, axis=0)) / 2
+    max_range = np.max(np.max(coords, axis=0) - np.min(coords, axis=0))
+    if max_range > 1e-10:
+        coords_normalized = (coords - center) / max_range
+    else:
+        coords_normalized = coords - center
+
+    for prop_name, prop_data in properties.items():
+        if prop_data is not None:
+            view_paths = _render_point_cloud_views(
+                coords_normalized, prop_data, output_dir, prop_name
+            )
+            colorbar_path = os.path.join(output_dir, f"{prop_name}_colorbar.png")
+            _create_colorbar(prop_data, prop_name, colorbar_path)
+            result[prop_name] = (view_paths, colorbar_path)
+            print(f"Created visualization for {prop_name}: {len(view_paths)} views")
+
+    return result
+
+
+@spaces.GPU()
+@torch.no_grad()
+def process_3d_model(input_file):
+    empty_result = (
+        None,
+        [],
+        None,
+        [],
+        None,
+        None,
+        [],
+        None,
+        None,
+        [],
+        None,
+        None,
+    )
+    if input_file is None:
+        return empty_result
+    output_dir = tempfile.mkdtemp(prefix="vomp_")
+    material_file = os.path.join(output_dir, "materials.npz")
+    try:
+        if _is_gaussian_splat(input_file):
+            print(f"Processing as Gaussian splat: {input_file}")
+            results = model.get_splat_materials(
+                input_file,
+                voxel_method="kaolin",
+                query_points="voxel_centers",
+                output_dir=output_dir,
+            )
+        else:
+            print(f"Processing as mesh: {input_file}")
+            _install_blender()
+            results = model.get_mesh_materials(
+                input_file,
+                blender_path=BLENDER_PATH,
+                query_points="voxel_centers",
+                output_dir=output_dir,
+                return_original_scale=True,
+            )
+
+        save_materials(results, material_file)
+        print(f"Materials saved to: {material_file}")
+
+        all_images = _get_render_images(output_dir)
+        first_image = all_images[0] if all_images else None
+
+        visualizations = _create_material_visualizations(material_file, output_dir)
+
+        youngs_views = visualizations.get("youngs_modulus", ([], None))[0]
+        youngs_colorbar = visualizations.get("youngs_modulus", ([], None))[1]
+        youngs_first = youngs_views[0] if youngs_views else None
+
+        poissons_views = visualizations.get("poissons_ratio", ([], None))[0]
+        poissons_colorbar = visualizations.get("poissons_ratio", ([], None))[1]
+        poissons_first = poissons_views[0] if poissons_views else None
+
+        density_views = visualizations.get("density", ([], None))[0]
+        density_colorbar = visualizations.get("density", ([], None))[1]
+        density_first = density_views[0] if density_views else None
+
+        return (
+            first_image,
+            all_images,
+            youngs_first,
+            youngs_views,
+            youngs_colorbar,
+            poissons_first,
+            poissons_views,
+            poissons_colorbar,
+            density_first,
+            density_views,
+            density_colorbar,
+            material_file,
+        )
+
+    except Exception as e:
+        print(f"Error processing 3D model: {e}")
+        raise gr.Error(f"Failed to process 3D model: {str(e)}")
+
+
+def update_slider_image(slider_value: int, all_images: List[str]) -> Optional[str]:
+    if not all_images or slider_value < 0 or slider_value >= len(all_images):
+        return None
+    return all_images[slider_value]
+
+
+def update_property_view(slider_value: int, views: List[str]) -> Optional[str]:
+    if not views or slider_value < 0 or slider_value >= len(views):
+        return None
+    return views[slider_value]
+
+
+css = """
+.gradio-container {
+    font-family: 'IBM Plex Sans', sans-serif;
+}
+
+.title-container {
+    text-align: center;
+    padding: 20px 0;
+}
+
+.badge-container {
+    display: flex;
+    justify-content: center;
+    gap: 8px;
+    flex-wrap: wrap;
+    margin-bottom: 20px;
+}
+
+.badge-container a img {
+    height: 22px;
+}
+
+h1 {
+    text-align: center;
+    font-size: 2.5rem;
+    margin-bottom: 0.5rem;
+}
+
+.subtitle {
+    text-align: center;
+    color: #666;
+    font-size: 1.1rem;
+    margin-bottom: 1.5rem;
+}
+
+.input-column, .output-column {
+    min-height: 400px;
+}
+
+.output-column .row {
+    display: flex !important;
+    flex-wrap: nowrap !important;
+    gap: 16px;
+}
+
+.output-column .row > .column {
+    flex: 1 1 50% !important;
+    min-width: 0 !important;
+}
+"""
+
+title_md = """
+<div class="title-container">
+    <h1>VoMP: Predicting Volumetric Mechanical Properties</h1>
+    <p class="subtitle">Feed-forward, fine-grained, physically based volumetric material properties from Splats, Meshes, NeRFs, and more.</p>
+    <div class="badge-container">
+<a href="https://arxiv.org/abs/2510.22975"><img src='https://img.shields.io/badge/arXiv-VoMP-red' alt='Paper PDF'></a>
+<a href='https://research.nvidia.com/labs/sil/projects/vomp/'><img src='https://img.shields.io/badge/Project_Page-VoMP-green' alt='Project Page'></a>
+<a href='https://huggingface.co/nvidia/PhysicalAI-Simulation-VoMP-Model'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20-Models-yellow'></a>
+<a href='https://huggingface.co/datasets/nvidia/PhysicalAI-Robotics-PhysicalAssets-VoMP'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20-GVM%20Dataset-yellow'></a>
+    </div>
+</div>
+"""
+
+description_md = """
+Upload a Gaussian Splat (.ply) or Mesh (.obj, .glb, .stl, .gltf) to predict volumetric mechanical properties (Young's modulus, Poisson's ratio, density) for realistic physics simulation.
+"""
+
+with gr.Blocks(css=css, title="VoMP") as demo:
+    all_images_state = gr.State([])
+    youngs_views_state = gr.State([])
+    poissons_views_state = gr.State([])
+    density_views_state = gr.State([])
+
+    gr.HTML(title_md)
+    gr.Markdown(description_md)
+
+    with gr.Row():
+        # Input Column (50%)
+        with gr.Column(scale=1, elem_classes="input-column"):
+            gr.Markdown("### 📤 Input")
+            input_model = gr.Model3D(
+                label="Upload 3D Model",
+                clear_color=[0.1, 0.1, 0.1, 1.0],
+            )
+
+            submit_btn = gr.Button(
+                "🚀 Generate Materials", variant="primary", size="lg"
+            )
+
+            gr.Markdown("#### 🎬 Rendered Views")
+            rendered_image = gr.Image(label="Rendered View", height=250)
+
+            view_slider = gr.Slider(
+                minimum=0,
+                maximum=NUM_VIEWS - 1,
+                step=1,
+                value=0,
+                label="Browse All Views",
+                info=f"Slide to view all {NUM_VIEWS} rendered views",
+            )
+
+        # Output Column (50%)
+        with gr.Column(scale=1, elem_classes="output-column"):
+            gr.Markdown("### 📥 Output - Material Properties")
+
+            # Row 1: Young's Modulus and Poisson's Ratio
+            with gr.Row():
+                with gr.Column(scale=1, min_width=200):
+                    youngs_image = gr.Image(label="Young's Modulus", height=200)
+                    youngs_slider = gr.Slider(
+                        minimum=0,
+                        maximum=2,
+                        step=1,
+                        value=0,
+                        label="View",
+                        info="Switch between 3 views",
+                    )
+                    youngs_colorbar = gr.Image(height=50, show_label=False)
+
+                with gr.Column(scale=1, min_width=200):
+                    poissons_image = gr.Image(label="Poisson's Ratio", height=200)
+                    poissons_slider = gr.Slider(
+                        minimum=0,
+                        maximum=2,
+                        step=1,
+                        value=0,
+                        label="View",
+                        info="Switch between 3 views",
+                    )
+                    poissons_colorbar = gr.Image(height=50, show_label=False)
+
+            # Row 2: Density and Download
+            with gr.Row():
+                with gr.Column(scale=1, min_width=200):
+                    density_image = gr.Image(label="Density", height=200)
+                    density_slider = gr.Slider(
+                        minimum=0,
+                        maximum=2,
+                        step=1,
+                        value=0,
+                        label="View",
+                        info="Switch between 3 views",
+                    )
+                    density_colorbar = gr.Image(height=50, show_label=False)
+
+                with gr.Column(scale=1, min_width=200):
+                    gr.Markdown("#### 💾 Download")
+                    output_file = gr.File(
+                        label="Download Materials (.npz)",
+                        file_count="single",
+                    )
+
+    gr.Markdown("### 🎯 Examples")
+    gr.Examples(
+        examples=[
+            [os.path.join(EXAMPLES_DIR, "plant.ply")],
+            [os.path.join(EXAMPLES_DIR, "dog.ply")],
+            [os.path.join(EXAMPLES_DIR, "dozer.ply")],
+            [os.path.join(EXAMPLES_DIR, "fiscus.ply")],
+        ],
+        inputs=[input_model],
+        outputs=[
+            rendered_image,
+            all_images_state,
+            youngs_image,
+            youngs_views_state,
+            youngs_colorbar,
+            poissons_image,
+            poissons_views_state,
+            poissons_colorbar,
+            density_image,
+            density_views_state,
+            density_colorbar,
+            output_file,
+        ],
+        fn=process_3d_model,
+        cache_examples=False,
+    )
+
+    # Event handlers
+    submit_btn.click(
+        fn=process_3d_model,
+        inputs=[input_model],
+        outputs=[
+            rendered_image,
+            all_images_state,
+            youngs_image,
+            youngs_views_state,
+            youngs_colorbar,
+            poissons_image,
+            poissons_views_state,
+            poissons_colorbar,
+            density_image,
+            density_views_state,
+            density_colorbar,
+            output_file,
+        ],
+    )
+
+    view_slider.change(
+        fn=update_slider_image,
+        inputs=[view_slider, all_images_state],
+        outputs=[rendered_image],
+    )
+
+    youngs_slider.change(
+        fn=update_property_view,
+        inputs=[youngs_slider, youngs_views_state],
+        outputs=[youngs_image],
+    )
+
+    poissons_slider.change(
+        fn=update_property_view,
+        inputs=[poissons_slider, poissons_views_state],
+        outputs=[poissons_image],
+    )
+
+    density_slider.change(
+        fn=update_property_view,
+        inputs=[density_slider, density_views_state],
+        outputs=[density_image],
+    )
+
+if __name__ == "__main__":
+    demo.launch()

deps/vomp/.gitignore

@@ -0,0 +1,216 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#   Usually these files are written by a python script from a template
+#   before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+# Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+# uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+# poetry.lock
+# poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+# pdm.lock
+# pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+# pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# Redis
+*.rdb
+*.aof
+*.pid
+
+# RabbitMQ
+mnesia/
+rabbitmq/
+rabbitmq-data/
+
+# ActiveMQ
+activemq-data/
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#   JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#   be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#   and can be added to the global gitignore or merged into this file.  For a more nuclear
+#   option (not recommended) you can uncomment the following to ignore the entire idea folder.
+# .idea/
+
+# Abstra
+#   Abstra is an AI-powered process automation framework.
+#   Ignore directories containing user credentials, local state, and settings.
+#   Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#   Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#   that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#   and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#   you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/
+
+# Streamlit
+.streamlit/secrets.toml

deps/vomp/CONTRIBUTING.md

@@ -0,0 +1,51 @@
+# How to Contribute
+
+We'd love to receive your patches and contributions. Please keep your PRs as draft until such time that you would like us to review them.
+
+## Code Reviews
+
+All submissions, including submissions by project members, require review. We use GitHub pull requests for this purpose. Consult
+[GitHub Help](https://help.github.com/articles/about-pull-requests/) for more information on using pull requests.
+
+## Signing Your Work
+
+* We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
+
+  * Any contribution which contains commits that are not Signed-Off will not be accepted.
+
+* To sign off on a commit you simply use the `--signoff` (or `-s`) option when committing your changes:
+  ```bash
+  $ git commit -s -m "Add cool feature."
+  ```
+  This will append the following to your commit message:
+  ```
+  Signed-off-by: Your Name <your@email.com>
+  ```
+
+* Full text of the DCO:
+
+  ```
+    Developer Certificate of Origin
+    Version 1.1
+
+    Copyright (C) 2004, 2006 The Linux Foundation and its contributors.
+    1 Letterman Drive
+    Suite D4700
+    San Francisco, CA, 94129
+
+    Everyone is permitted to copy and distribute verbatim copies of this license document, but changing it is not allowed.
+  ```
+
+  ```
+    Developer's Certificate of Origin 1.1
+
+    By making a contribution to this project, I certify that:
+
+    (a) The contribution was created in whole or in part by me and I have the right to submit it under the open source license indicated in the file; or
+
+    (b) The contribution is based upon previous work that, to the best of my knowledge, is covered under an appropriate open source license and I have the right under that license to submit that work with modifications, whether created in whole or in part by me, under the same open source license (unless I am permitted to submit under a different license), as indicated in the file; or
+
+    (c) The contribution was provided directly to me by some other person who certified (a), (b) or (c) and I have not modified it.
+
+    (d) I understand and agree that this project and the contribution are public and that a record of the contribution (including all personal information I submit with it, including my sign-off) is maintained indefinitely and may be redistributed consistent with this project or the open source license(s) involved.
+  ```

deps/vomp/LICENSE

@@ -0,0 +1,201 @@
+                                Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

deps/vomp/README.md

@@ -0,0 +1,665 @@
+<div align="center">
+<h2>VoMP: Predicting Volumetric Mechanical Properties</h2>
+
+<a href="https://arxiv.org/abs/2510.22975"><img src='https://img.shields.io/badge/arXiv-VoMP-red' alt='Paper PDF'></a>
+<a href='https://research.nvidia.com/labs/sil/projects/vomp/'><img src='https://img.shields.io/badge/Project_Page-VoMP-green' alt='Project Page'></a>
+<a href='https://huggingface.co/nvidia/PhysicalAI-Simulation-VoMP-Model'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20-Models-yellow'></a>
+<a href='https://huggingface.co/datasets/nvidia/PhysicalAI-Robotics-PhysicalAssets-VoMP'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20-GVM%20Dataset-yellow'></a>
+</div>
+
+![](./images/teaser.png)
+
+This repository provides the implementation of **VoMP**. TL;DR: Feed-forward, fine-grained, physically based volumetric material properties from Splats, Meshes, NeRFs, etc. which can be used to produce realistic worlds. We recommend reading the [README_train.md](./README_train.md) if you need to fine-tune or train the model from scratch or know more details about the codebase.
+
+---
+
+## Contents
+
+- [🔧 Dependencies and Installation](#-dependencies-and-installation)
+  * [Setup a Virtual Environment (Recommended)](#setup-a-virtual-environment--recommended-)
+  * [Install a Mesh Renderer (Required for Mesh Processing Only)](#install-a-mesh-renderer--required-for-mesh-processing-only-)
+    + [Isaac Sim](#isaac-sim)
+    + [Blender](#blender)
+  * [Setup a Conda Environment (Alternative)](#setup-a-conda-environment--alternative-)
+  * [Trained Models](#trained-models)
+- [🌐 Quickstart: Web Demo](#-quickstart-web-demo)
+- [📥 Loading the Model](#-loading-the-model)
+  * [Using Inference Config (Recommended)](#using-inference-config--recommended-)
+  * [Using Direct File Paths](#using-direct-file-paths)
+  * [Using Directories (use for fine-tuning)](#using-directories--use-for-fine-tuning-)
+- [🎯 High-Level API](#-high-level-api)
+  * [Gaussian Splats](#gaussian-splats)
+  * [Meshes](#meshes)
+  * [USD Assets (including meshes)](#usd-assets--including-meshes-)
+    + [General USD Formats](#general-usd-formats)
+    + [SimReady Format USD](#simready-format-usd)
+- [🎨 Visualizing Material Results](#-visualizing-material-results)
+- [🔧 Low-Level API](#-low-level-api)
+  * [Gaussian Splats](#gaussian-splats-1)
+  * [Meshes](#meshes-1)
+  * [USD Assets](#usd-assets)
+- [🧩 Custom 3D Representations](#-custom-3d-representations)
+- [🧬 Material Upsampler](#-material-upsampler)
+- [💾 Using our Benchmark](#-using-our-benchmark)
+  * [Reproducing results from the paper](#reproducing-results-from-the-paper)
+- [📦 Simulation](#-simulation)
+  * [Simplicits simulation](#simplicits-simulation)
+  * [FEM simulation using warp.fem](#fem-simulation-using-warpfem)
+  * [FEM simulation using libuipc](#fem-simulation-using-libuipc)
+  * [Newton simulation](#newton-simulation)
+- [🤗 Credits](#-credits)
+- [📜 Citation](#-citation)
+- [License and Contact](#license-and-contact)
+
+## 🔧 Dependencies and Installation
+
+All the instructions in this README are meant to be run from the root of the repository. Running simulations requires a separate set of dependencies than this setup which we mention later in the [📦 Simulation](#-simulation) section.
+
+### Setup a Virtual Environment (Recommended)
+
+First set up the environment. We recommend using Python>=3.10, PyTorch>=2.1.0, and CUDA>=11.8. It is okay if some packages show warnings or fail to install due to version conflicts. The version conflicts are not a problem for the functionalities we use.
+
+```bash
+git clone --recursive https://github.com/nv-tlabs/VoMP
+cd VoMP
+
+# Install dependencies using the provided script (Linux only)
+chmod +x install_env.sh
+./install_env.sh
+```
+
+> [!NOTE]
+> Running install_env.sh without conda: The script includes optional conda-only steps (environment creation/activation, installing CUDA toolkit inside the env, and setting env vars). If you're using a Python `venv` and don't have conda, the script will fail when it tries to call `conda`. You can either install conda, or comment out the conda-specific lines (lines 93-115 and any `conda install` / `conda env config vars set` commands). The rest of the script relies on `pip` and standard bash commands and will work in a `venv`.
+
+### Install a Mesh Renderer (Required for Mesh Processing Only)
+
+We only need a mesh renderer so you can download any one of Isaac Sim or Blender. There is no need to install both.
+
+#### Isaac Sim
+
+For mesh material estimation, you need to install Isaac Sim or Blender manually. *This is not required for Gaussian splat processing.*
+
+Download Isaac Sim from [here](https://docs.isaacsim.omniverse.nvidia.com/5.0.0/installation/index.html) and follow the instructions to install it. On Linux, you would have a `isaac-sim.sh` file in the path you installed it. For Windows, you would have a `isaac-sim.bat` file in the path you installed it. Note the path to the `isaac-sim.sh` or `isaac-sim.bat` file.
+
+> [!NOTE]
+> You'll need to provide the Isaac Sim binary path when using mesh APIs.
+
+> [!WARNING]
+> We use Replicator in Isaac Sim to render meshes. Replicator supports USD assets. If you want to use a USD asset, since USD files can contain many things in many formats we expect you to have used [existing tools](https://openusd.org/release/toolset.html) to convert it into an explicit mesh. If you want to use a mesh asset, you can use Replicator by also having a USD file of your mesh that you can make by using [existing tools](https://openusd.org/release/toolset.html).  
+
+#### Blender
+
+For mesh material estimation, you need to install Blender 3.0+ manually. *This is not required for Gaussian splat processing.*
+
+```bash
+# Install system dependencies
+sudo apt-get update
+sudo apt-get install -y libxrender1 libxi6 libxkbcommon-x11-0 libsm6
+
+# Download and install Blender 3.0.1
+wget https://download.blender.org/release/Blender3.0/blender-3.0.1-linux-x64.tar.xz
+tar -xvf blender-3.0.1-linux-x64.tar.xz
+
+# Note the path: ./blender-3.0.1-linux-x64/blender
+```
+
+> [!NOTE]
+> You'll need to provide the Blender binary path when using mesh APIs:
+> ```python
+> results = model.get_mesh_materials("mesh.obj", blender_path="/path/to/blender")
+> ```
+
+### Setup a Conda Environment (Alternative)
+
+We also provide a conda environment file to install the dependencies. This automatically creates a new environment:
+
+```bash
+# Create and install environment from file (creates 'vomp' environment)
+conda env create -f environment.yml
+
+# Activate the environment
+conda activate vomp
+```
+
+> [!WARNING]
+> We do not recommend using this installation method. The conda environment file is accurate but it reflects the environment at its final stage and does not have the step-by-step process we use to install the dependencies.
+
+### Trained Models
+
+We provide the trained models (1.73 GB) in <a href='https://huggingface.co/nvidia/PhysicalAI-Simulation-VoMP-Model'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20-Models-yellow'></a>. 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:
+
+![Colorbar](images/youngs_modulus_colorbar_legend.png)
+
+## 🔧 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 <a href='https://huggingface.co/datasets/nvidia/PhysicalAI-Robotics-PhysicalAssets-VoMP'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20-Dataset-yellow'></a> (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/).

deps/vomp/README_train.md

@@ -0,0 +1,285 @@
+<div align="center">
+<h2>VoMP: Predicting Volumetric Mechanical Properties</h2>
+
+<a href="https://arxiv.org/abs/2510.22975"><img src='https://img.shields.io/badge/arXiv-VoMP-red' alt='Paper PDF'></a>
+<a href='https://research.nvidia.com/labs/sil/projects/vomp/'><img src='https://img.shields.io/badge/Project_Page-VoMP-green' alt='Project Page'></a>
+<a href='https://huggingface.co/nvidia/PhysicalAI-Simulation-VoMP-Model'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20-Models-yellow'></a>
+<a href='https://huggingface.co/datasets/nvidia/PhysicalAI-Robotics-PhysicalAssets-VoMP'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20-GVM%20Dataset-yellow'></a>
+</div>
+
+![](./images/teaser.png)
+
+This repository provides the implementation of **VoMP**. TL;DR: Feed-forward, fine-grained, physically based volumetric material properties from Splats, Meshes, NeRFs, etc. which can be used to produce realistic worlds.
+
+---
+
+## Contents
+
+- [🔧 Setup](#-setup)
+- [📖 Overview of the codebase](#-overview-of-the-codebase)
+- [📚 Create the dataset](#-create-the-dataset)
+  * [Preprocessed Datasets](#preprocessed-datasets)
+  * [Material Triplet Dataset (MTD)](#material-triplet-dataset--mtd-)
+  * [Geometry with Volumetric Materials (GVM)](#geometry-with-volumetric-materials--gvm-)
+  * [Preparing your own data for training the Geometry Transformer](#preparing-your-own-data-for-training-the-geometry-transformer)
+- [💻 Training](#-training)
+  * [Training the MatVAE](#training-the-matvae)
+  * [Training the Geometry Transformer](#training-the-geometry-transformer)
+  * [Training on your own data](#training-on-your-own-data)
+  * [Fine-tuning](#fine-tuning)
+- [💡 Tips](#-tips)
+
+## 🔧 Setup
+
+Follow the instructions in the [README.md](./README.md) file to set up the environment.
+
+## 📖 Overview of the codebase
+
+![](./images/method.png)
+
+The codebase is organized as follows:
+
+- `train_material_vae.py`: Main entry point for training the MatVAE.
+- `train_geometry_encoder.py`: Main entry point for training the Geometry Transformer.
+- `vomp/`: Main Python package containing all models and utilities.
+  - `models/`: Neural network architectures including MatVAE and Geometry Transformer.
+    - `geometry_encoder.py`: Geometry Transformer encoder.
+    - `material_vae/`: MatVAE model implementations.
+    - `structured_latent_vae/`: Structured latent VAE components.
+  - `trainers/`: Training frameworks for different model types.
+  - `modules/`: Neural network layer classes (sparse transformers, attention, etc.).
+  - `datasets/`: Dataset loaders (`SparseVoxelMaterials`, etc.).
+  - `representations/`: 3D representation handlers (Gaussian splats).
+  - `inference/`: Inference pipeline (`vomp.py`) and utilities.
+  - `utils/`: General utility functions and data processing tools.
+- `dataset_toolkits/`: Tools for dataset creation and preprocessing.
+  - `material_objects/`: Material property rendering, voxelization, and VLM annotation tools.
+  - `datasets/`: Dataset loaders (simready, ABO, etc.).
+- `configs/`: Configuration files for different experiments.
+  - `materials/`: MatVAE and Geometry Transformer configurations.
+- `scripts/`: Visualization and evaluation scripts.
+- `weights/`: Directory for storing pretrained model weights.
+
+## 📚 Create the dataset
+
+We provide toolkits for data preparation.
+
+![](./images/datacreation.png)
+
+### Preprocessed Datasets
+
+We provide the preprocessed datasets (with the vegetation subset removed) at: <a href='https://huggingface.co/datasets/nvidia/PhysicalAI-Robotics-PhysicalAssets-VoMP'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20-GVM%20Dataset-yellow'></a>. 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.

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"
+        }
+    }
+}

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"
+        }
+    }
+} 

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"
+        }
+    }
+} 

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"
+        }
+    }
+}

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"
+        }
+    }
+}

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
+} 

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
+} 

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
+}

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
+} 

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
+} 

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
+} 

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
+  } 

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
+  } 

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
+} 

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"
+  }
+} 

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"
+  }
+} 

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"
+  }
+} 

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"
+  }
+} 

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"
+  }
+} 

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"
+  }
+} 

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"
+  }
+} 

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"
+  }
+} 

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)

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()

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.")

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.")

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.")

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)

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())

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)

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

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)

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)

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,
+    )

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")

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()

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 <UDIM> with first available tile
+        if "<UDIM>" in texture_path:
+            # Try common UDIM tile numbers
+            for udim in ["1001", "1002", "1003", "1004"]:
+                resolved = self._resolve_texture_path(
+                    texture_path.replace("<UDIM>", udim)
+                )
+                if resolved:
+                    return resolved
+            return None
+
+        # Already absolute
+        if Path(texture_path).is_absolute():
+            p = Path(texture_path)
+            return p if p.exists() else None
+
+        # Try relative to USD directory
+        candidates = [
+            self.usd_dir / texture_path,
+            self.usd_dir / "textures" / Path(texture_path).name,
+            self.usd_dir / "Textures" / Path(texture_path).name,
+            self.usd_dir / ".." / texture_path,
+            self.usd_dir / ".." / "textures" / Path(texture_path).name,
+            self.usd_dir / ".." / "materials" / "textures" / Path(texture_path).name,
+        ]
+
+        for p in candidates:
+            if p.exists():
+                return p.resolve()
+
+        # Fuzzy matching: look for files containing the texture name
+        texture_name = Path(texture_path).stem  # e.g., "Iron_BaseColor"
+        texture_ext = Path(texture_path).suffix  # e.g., ".png"
+
+        # Search in Textures folders
+        search_dirs = [
+            self.usd_dir / "Textures",
+            self.usd_dir / "textures",
+            self.usd_dir / ".." / "Textures",
+            self.usd_dir / ".." / "textures",
+        ]
+
+        for search_dir in search_dirs:
+            if search_dir.exists():
+                for f in search_dir.iterdir():
+                    # Check if the file contains the texture name (fuzzy match)
+                    if (
+                        texture_name in f.stem
+                        and f.suffix.lower() == texture_ext.lower()
+                    ):
+                        self._log(f"    (fuzzy match: {texture_name} -> {f.name})")
+                        return f.resolve()
+
+        return None
+
+    def _categorize_input(self, name: str) -> str:
+        name_lower = name.lower()
+
+        # Check for texture/color type based on common patterns
+        if any(
+            x in name_lower for x in ["diffuse", "albedo", "basecolor", "base_color"]
+        ):
+            return "diffuse"
+        elif any(x in name_lower for x in ["normal", "bump"]):
+            return "normal"
+        elif any(x in name_lower for x in ["rough"]):
+            return "roughness"
+        elif any(x in name_lower for x in ["metal"]):
+            return "metallic"
+        elif any(x in name_lower for x in ["orm", "occlusion"]):
+            return "orm"
+        elif any(x in name_lower for x in ["opacity", "alpha"]):
+            return "opacity"
+        else:
+            return name  # Use original name if no match
+
+    def _find_fallback_textures(self, material_name: str) -> Dict[str, str]:
+        textures = {}
+
+        # Search in Textures folders
+        search_dirs = [
+            self.usd_dir / "Textures",
+            self.usd_dir / "textures",
+        ]
+
+        for search_dir in search_dirs:
+            if not search_dir.exists():
+                continue
+
+            # Find all unique texture prefixes (e.g., BlueRug from BlueRug_BaseColor.png)
+            texture_files = list(search_dir.glob("*.png")) + list(
+                search_dir.glob("*.jpg")
+            )
+            if not texture_files:
+                continue
+
+            # Group by prefix (before _BaseColor, _N, _R, etc.)
+            prefixes = set()
+            for f in texture_files:
+                stem = f.stem
+                for suffix in [
+                    "_BaseColor",
+                    "_basecolor",
+                    "_A",
+                    "_albedo",
+                    "_diffuse",
+                    "_N",
+                    "_Normal",
+                    "_normal",
+                    "_R",
+                    "_Roughness",
+                    "_roughness",
+                ]:
+                    if suffix in stem:
+                        prefix = stem.split(suffix)[0]
+                        prefixes.add(prefix)
+                        break
+
+            # Use the first available texture set
+            if prefixes:
+                prefix = sorted(prefixes)[0]  # Pick first alphabetically
+                self._log(f"    (fallback: using {prefix}_* textures)")
+
+                # Find matching textures
+                for f in texture_files:
+                    if f.stem.startswith(prefix):
+                        stem_lower = f.stem.lower()
+                        if any(
+                            x in stem_lower
+                            for x in ["basecolor", "_a", "albedo", "diffuse"]
+                        ):
+                            textures["diffuse"] = str(f.resolve())
+                            self._log(f"    ✓ fallback diffuse: {f.name}")
+                        elif any(x in stem_lower for x in ["normal", "_n"]):
+                            textures["normal"] = str(f.resolve())
+                            self._log(f"    ✓ fallback normal: {f.name}")
+                        elif any(x in stem_lower for x in ["rough", "_r"]):
+                            textures["roughness"] = str(f.resolve())
+                            self._log(f"    ✓ fallback roughness: {f.name}")
+
+                if textures:
+                    return textures
+
+        return textures
+
+    def _extract_textures_from_shader(self, shader: UsdShade.Shader) -> Dict[str, any]:
+        result = {}
+
+        for shader_input in shader.GetInputs():
+            val = shader_input.Get()
+            if val is None:
+                continue
+
+            input_name = shader_input.GetBaseName()
+            category = self._categorize_input(input_name)
+
+            # Texture path (AssetPath)
+            if isinstance(val, Sdf.AssetPath) and val.path:
+                texture_path = val.path.strip("@")
+                resolved = self._resolve_texture_path(texture_path)
+                if resolved:
+                    result[category] = str(resolved)
+                    self._log(f"    ✓ {input_name} -> {category}: {resolved.name}")
+                else:
+                    self._log(f"    ✗ {input_name}: {texture_path} (not resolved)")
+
+            # Color value (Vec3)
+            elif (
+                hasattr(val, "__len__")
+                and len(val) == 3
+                and "color" in input_name.lower()
+            ):
+                result[f"{category}_color"] = (
+                    float(val[0]),
+                    float(val[1]),
+                    float(val[2]),
+                )
+                self._log(
+                    f"    ✓ {input_name} -> {category}_color: ({val[0]:.3f}, {val[1]:.3f}, {val[2]:.3f})"
+                )
+
+        return result
+
+    def _extract_textures_from_mdl(self, mdl_path: Path) -> Dict[str, str]:
+        textures = {}
+
+        if not mdl_path.exists():
+            return textures
+
+        try:
+            content = mdl_path.read_text()
+
+            # Parse texture references
+            type_mapping = {
+                "diffuse_texture": "diffuse",
+                "normalmap_texture": "normal",
+                "reflectionroughness_texture": "roughness",
+                "metallic_texture": "metallic",
+                "ORM_texture": "orm",
+            }
+
+            for tex_type, canonical_name in type_mapping.items():
+                pattern = rf'{tex_type}:\s*texture_2d\("([^"]+)"'
+                match = re.search(pattern, content)
+                if match:
+                    rel_path = match.group(1)
+                    # MDL paths are relative to the MDL file location
+                    resolved = self._resolve_texture_path(rel_path)
+                    if not resolved:
+                        # Try relative to MDL file directory
+                        mdl_dir = mdl_path.parent
+                        candidates = [
+                            mdl_dir / rel_path,
+                            mdl_dir / "textures" / Path(rel_path).name,
+                        ]
+                        for c in candidates:
+                            if c.exists():
+                                resolved = c.resolve()
+                                break
+
+                    if resolved:
+                        textures[canonical_name] = str(resolved)
+                        self._log(f"    ✓ {canonical_name}: {resolved.name} (from MDL)")
+                    else:
+                        self._log(
+                            f"    ✗ {canonical_name}: {rel_path} (MDL, not resolved)"
+                        )
+
+        except Exception as e:
+            self._log(f"    Error parsing MDL {mdl_path}: {e}")
+
+        return textures
+
+    def _find_mdl_for_material(self, material_prim: Usd.Prim) -> Optional[Path]:
+        for child in material_prim.GetChildren():
+            if child.GetTypeName() == "Shader":
+                # Check for MDL source asset
+                mdl_attr = child.GetAttribute("info:mdl:sourceAsset")
+                if mdl_attr and mdl_attr.Get():
+                    mdl_path_val = mdl_attr.Get()
+                    if isinstance(mdl_path_val, Sdf.AssetPath) and mdl_path_val.path:
+                        mdl_rel = mdl_path_val.path.strip("@")
+
+                        # Try to resolve MDL path
+                        candidates = [
+                            self.usd_dir / mdl_rel,
+                            self.usd_dir / "materials" / Path(mdl_rel).name,
+                            self.usd_dir / ".." / "materials" / Path(mdl_rel).name,
+                        ]
+
+                        for c in candidates:
+                            if c.exists():
+                                return c.resolve()
+
+        return None
+
+    def _get_geomsubset_bindings(
+        self, mesh_prim: Usd.Prim
+    ) -> Dict[str, Tuple[str, List[int]]]:
+        bindings = {}
+
+        for child in mesh_prim.GetChildren():
+            if child.GetTypeName() == "GeomSubset":
+                subset_name = child.GetName()
+
+                # Get face indices for this subset
+                indices_attr = child.GetAttribute("indices")
+                face_indices = (
+                    list(indices_attr.Get())
+                    if indices_attr and indices_attr.Get()
+                    else []
+                )
+
+                # Get material binding
+                mat_path = None
+                binding_rel = child.GetRelationship("material:binding")
+                if binding_rel:
+                    targets = binding_rel.GetTargets()
+                    if targets:
+                        mat_path = str(targets[0])
+
+                if mat_path:
+                    bindings[subset_name] = (mat_path, face_indices)
+
+        return bindings
+
+    def extract(self) -> bool:
+        try:
+            self.stage = Usd.Stage.Open(str(self.usd_path))
+        except Exception as e:
+            print(f"ERROR: Could not open USD: {self.usd_path}")
+            print(f"  {e}")
+            return False
+
+        if not self.stage:
+            return False
+
+        self._log(f"\n=== Extracting from: {self.usd_path.name} ===")
+
+        # Step 1: Find all materials and extract textures
+        self._log("\n--- Materials ---")
+        for prim in self.stage.Traverse():
+            if prim.GetTypeName() == "Material":
+                mat_path = str(prim.GetPath())
+                self._log(f"\nMaterial: {prim.GetName()}")
+
+                textures = {}
+
+                # Try extracting from shader inputs
+                for child in prim.GetChildren():
+                    if child.GetTypeName() == "Shader":
+                        shader = UsdShade.Shader(child)
+                        textures.update(self._extract_textures_from_shader(shader))
+
+                # If no textures found, try MDL
+                if not textures:
+                    mdl_path = self._find_mdl_for_material(prim)
+                    if mdl_path:
+                        self._log(f"  Using MDL: {mdl_path.name}")
+                        textures = self._extract_textures_from_mdl(mdl_path)
+
+                # Fallback: if still no textures, search Textures folder for any available
+                if not textures:
+                    textures = self._find_fallback_textures(prim.GetName())
+
+                self.materials[mat_path] = textures
+
+        # Step 2: Find all meshes and their material bindings
+        self._log("\n--- Meshes ---")
+        for prim in self.stage.Traverse():
+            if prim.GetTypeName() == "Mesh":
+                mesh_path = str(prim.GetPath())
+                mesh_name = prim.GetName()
+
+                # Get direct material binding first
+                binding_api = UsdShade.MaterialBindingAPI(prim)
+                bound_material = binding_api.ComputeBoundMaterial()[0]
+                mat_path = str(bound_material.GetPath()) if bound_material else None
+
+                # Check for GeomSubset bindings (per-face materials)
+                geomsubset_bindings = self._get_geomsubset_bindings(prim)
+
+                # If no direct binding but has GeomSubsets, use first one as default
+                if not mat_path and geomsubset_bindings:
+                    first_subset = list(geomsubset_bindings.values())[0]
+                    mat_path = first_subset[0]
+
+                self.mesh_materials[mesh_path] = mat_path
+
+                # Get mesh geometry
+                mesh = UsdGeom.Mesh(prim)
+                points_local = mesh.GetPointsAttr().Get()
+                face_counts = mesh.GetFaceVertexCountsAttr().Get()
+                face_indices = mesh.GetFaceVertexIndicesAttr().Get()
+
+                # Apply world transform to vertices
+                xformable = UsdGeom.Xformable(prim)
+                world_transform = xformable.ComputeLocalToWorldTransform(
+                    Usd.TimeCode.Default()
+                )
+
+                # Transform points to world space
+                points = []
+                if points_local:
+                    for p in points_local:
+                        # Apply 4x4 transform matrix to point
+                        p_world = world_transform.Transform(Gf.Vec3d(p[0], p[1], p[2]))
+                        points.append(Gf.Vec3f(p_world[0], p_world[1], p_world[2]))
+                else:
+                    points = None
+
+                # Get UVs and check interpolation
+                uvs = None
+                uv_interpolation = None
+                uv_indices = None
+                for primvar_name in ["st", "uvs", "uv", "UVMap", "texCoords"]:
+                    primvar = UsdGeom.PrimvarsAPI(prim).GetPrimvar(primvar_name)
+                    if primvar and primvar.Get():
+                        uvs = primvar.Get()
+                        uv_interpolation = primvar.GetInterpolation()
+                        # For indexed primvars, get the indices
+                        if primvar.IsIndexed():
+                            uv_indices = primvar.GetIndices()
+                        break
+
+                self.meshes[mesh_path] = {
+                    "name": mesh_name,
+                    "material": mat_path,
+                    "points": points,
+                    "face_counts": face_counts,
+                    "face_indices": face_indices,
+                    "uvs": uvs,
+                    "uv_interpolation": uv_interpolation,
+                    "uv_indices": uv_indices,
+                    "geomsubsets": geomsubset_bindings,  # Store GeomSubset bindings
+                }
+
+                has_tex = bool(self.materials.get(mat_path))
+                if geomsubset_bindings:
+                    has_tex = any(
+                        self.materials.get(m) for m, _ in geomsubset_bindings.values()
+                    )
+
+                status = "✓" if has_tex else "○"
+
+                if geomsubset_bindings:
+                    self._log(
+                        f"  {status} {mesh_name} (GeomSubsets: {len(geomsubset_bindings)})"
+                    )
+                    for subset_name, (sub_mat, _) in geomsubset_bindings.items():
+                        sub_mat_name = sub_mat.split("/")[-1] if sub_mat else "none"
+                        self._log(f"      {subset_name} -> {sub_mat_name}")
+                else:
+                    self._log(
+                        f"  {status} {mesh_name} -> {mat_path or '(no material)'}"
+                    )
+
+        if not self.meshes:
+            self._log("WARNING: No meshes found in USD file")
+            return False
+
+        has_valid_mesh = any(
+            mesh_data.get("points") for mesh_data in self.meshes.values()
+        )
+        if not has_valid_mesh:
+            self._log("WARNING: No meshes with valid geometry found")
+            return False
+
+        return True
+
+    def export_obj(
+        self, output_dir: Path, normalize: bool = True
+    ) -> Tuple[Optional[Path], Optional[Path]]:
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        obj_path = output_dir / "model.obj"
+        mtl_path = output_dir / "model.mtl"
+
+        # Collect all vertices, faces, UVs
+        all_vertices = []
+        all_faces = []
+        all_uvs = []
+        face_materials = []
+        vertex_offset = 0
+        uv_offset = 0
+
+        material_list = []  # List of unique materials used
+        material_map = {}  # material_path -> index
+
+        for mesh_path, mesh_data in self.meshes.items():
+            if not mesh_data["points"]:
+                continue
+
+            points = mesh_data["points"]
+            face_counts = mesh_data["face_counts"]
+            face_indices = mesh_data["face_indices"]
+            uvs = mesh_data["uvs"]
+            uv_interpolation = mesh_data.get("uv_interpolation")
+            uv_indices = mesh_data.get("uv_indices")
+            mat_path = mesh_data["material"]
+            geomsubsets = mesh_data.get("geomsubsets", {})
+
+            # Add vertices
+            for p in points:
+                all_vertices.append((float(p[0]), float(p[1]), float(p[2])))
+
+            # Add UVs - handle different interpolation modes
+            mesh_uv_offset = len(all_uvs)
+            if uvs:
+                for uv in uvs:
+                    all_uvs.append((float(uv[0]), float(uv[1])))
+
+            # Track UV mapping for this mesh
+            # For faceVarying, we need to map face-vertex index to UV index
+            mesh_data["_uv_offset"] = mesh_uv_offset
+            mesh_data["_has_uvs"] = uvs is not None and len(uvs) > 0
+
+            # Build face-to-material mapping for GeomSubsets
+            face_to_subset_mat = {}
+            if geomsubsets:
+                for subset_name, (
+                    sub_mat_path,
+                    sub_face_indices,
+                ) in geomsubsets.items():
+                    for face_idx in sub_face_indices:
+                        face_to_subset_mat[face_idx] = sub_mat_path
+                    # Ensure this material is in our list
+                    if sub_mat_path and sub_mat_path not in material_map:
+                        material_map[sub_mat_path] = len(material_list)
+                        material_list.append(sub_mat_path)
+
+            # Track default material
+            if mat_path and mat_path not in material_map:
+                material_map[mat_path] = len(material_list)
+                material_list.append(mat_path)
+
+            # Add faces with proper UV indexing
+            idx = 0
+            face_num = 0
+            face_vertex_idx = 0  # Running index for faceVarying UVs
+
+            for count in face_counts:
+                # Determine material for this face
+                if face_num in face_to_subset_mat:
+                    face_mat = face_to_subset_mat[face_num]
+                else:
+                    face_mat = mat_path
+
+                mat_idx = material_map.get(face_mat, 0) if face_mat else 0
+
+                # Determine UV indices based on interpolation mode
+                def get_uv_idx(local_vert_idx, fv_offset):
+                    if not mesh_data["_has_uvs"]:
+                        return None
+                    vertex_idx = face_indices[local_vert_idx]
+
+                    if uv_interpolation == "faceVarying":
+                        if uv_indices is not None:
+                            # Indexed faceVarying: indices are per face-vertex
+                            return mesh_uv_offset + int(uv_indices[fv_offset])
+                        else:
+                            # Non-indexed faceVarying: sequential per face-vertex
+                            return mesh_uv_offset + fv_offset
+                    else:
+                        # vertex interpolation
+                        if uv_indices is not None:
+                            # Indexed vertex: indices are per-vertex
+                            return mesh_uv_offset + int(uv_indices[vertex_idx])
+                        else:
+                            # Non-indexed vertex: UV index matches vertex index
+                            return mesh_uv_offset + vertex_idx
+
+                if count == 3:
+                    v_indices = [
+                        face_indices[idx] + vertex_offset,
+                        face_indices[idx + 1] + vertex_offset,
+                        face_indices[idx + 2] + vertex_offset,
+                    ]
+                    uv_idxs = [
+                        get_uv_idx(idx, face_vertex_idx),
+                        get_uv_idx(idx + 1, face_vertex_idx + 1),
+                        get_uv_idx(idx + 2, face_vertex_idx + 2),
+                    ]
+                    all_faces.append((v_indices, uv_idxs))
+                    face_materials.append(mat_idx)
+                    face_vertex_idx += 3
+                elif count == 4:
+                    # Triangulate quad
+                    v_indices1 = [
+                        face_indices[idx] + vertex_offset,
+                        face_indices[idx + 1] + vertex_offset,
+                        face_indices[idx + 2] + vertex_offset,
+                    ]
+                    v_indices2 = [
+                        face_indices[idx] + vertex_offset,
+                        face_indices[idx + 2] + vertex_offset,
+                        face_indices[idx + 3] + vertex_offset,
+                    ]
+                    uv_idxs1 = [
+                        get_uv_idx(idx, face_vertex_idx),
+                        get_uv_idx(idx + 1, face_vertex_idx + 1),
+                        get_uv_idx(idx + 2, face_vertex_idx + 2),
+                    ]
+                    uv_idxs2 = [
+                        get_uv_idx(idx, face_vertex_idx),
+                        get_uv_idx(idx + 2, face_vertex_idx + 2),
+                        get_uv_idx(idx + 3, face_vertex_idx + 3),
+                    ]
+                    all_faces.append((v_indices1, uv_idxs1))
+                    all_faces.append((v_indices2, uv_idxs2))
+                    face_materials.append(mat_idx)
+                    face_materials.append(mat_idx)
+                    face_vertex_idx += 4
+                else:
+                    # Skip n-gons
+                    face_vertex_idx += count
+
+                idx += count
+                face_num += 1
+
+            vertex_offset += len(points)
+
+        if not all_vertices:
+            return None, None
+
+        # Normalize vertices to fit in [-0.5, 0.5]^3 centered at origin
+        # Use margin factor to ensure object fits fully in camera frame at all angles
+        MARGIN_FACTOR = 0.85  # Scale to 85% of unit cube to leave padding
+
+        if normalize and all_vertices:
+            # Compute bounding box
+            xs = [v[0] for v in all_vertices]
+            ys = [v[1] for v in all_vertices]
+            zs = [v[2] for v in all_vertices]
+
+            min_x, max_x = min(xs), max(xs)
+            min_y, max_y = min(ys), max(ys)
+            min_z, max_z = min(zs), max(zs)
+
+            # Compute center and scale
+            center_x = (min_x + max_x) / 2
+            center_y = (min_y + max_y) / 2
+            center_z = (min_z + max_z) / 2
+
+            extent_x = max_x - min_x
+            extent_y = max_y - min_y
+            extent_z = max_z - min_z
+            max_extent = max(extent_x, extent_y, extent_z)
+
+            if max_extent > 0:
+                # Scale to fit in unit cube, with margin for camera framing
+                scale = MARGIN_FACTOR / max_extent
+            else:
+                scale = 1.0
+
+            # Apply normalization: center then scale
+            all_vertices = [
+                (
+                    (v[0] - center_x) * scale,
+                    (v[1] - center_y) * scale,
+                    (v[2] - center_z) * scale,
+                )
+                for v in all_vertices
+            ]
+
+            self._log(f"\nNormalization applied:")
+            self._log(
+                f"  Original bounds: X[{min_x:.2f}, {max_x:.2f}], Y[{min_y:.2f}, {max_y:.2f}], Z[{min_z:.2f}, {max_z:.2f}]"
+            )
+            self._log(f"  Scale factor: {scale:.6f} (with {MARGIN_FACTOR:.0%} margin)")
+            self._log(
+                f"  Center offset: ({center_x:.2f}, {center_y:.2f}, {center_z:.2f})"
+            )
+
+        # Copy textures and write MTL
+        with open(mtl_path, "w") as f:
+            for mat_path in material_list:
+                mat_name = mat_path.split("/")[-1] if mat_path else "default_material"
+                textures = self.materials.get(mat_path, {})
+
+                f.write(f"newmtl {mat_name}\n")
+                f.write("Ka 0.2 0.2 0.2\n")
+
+                # Use diffuse color constant if available, otherwise default gray
+                if "diffuse_color" in textures:
+                    color = textures["diffuse_color"]
+                    f.write(f"Kd {color[0]:.6f} {color[1]:.6f} {color[2]:.6f}\n")
+                    self._log(
+                        f"  Material {mat_name}: using diffuse color ({color[0]:.3f}, {color[1]:.3f}, {color[2]:.3f})"
+                    )
+                else:
+                    f.write("Kd 0.8 0.8 0.8\n")
+
+                f.write("Ks 0.2 0.2 0.2\n")
+                f.write("Ns 50.0\n")
+                f.write("d 1.0\n")
+                f.write("illum 2\n")
+
+                for tex_type, tex_value in textures.items():
+                    # Skip color constants (they're tuples, not paths)
+                    if isinstance(tex_value, tuple):
+                        continue
+
+                    tex_path = tex_value
+                    if os.path.exists(tex_path):
+                        # Copy texture to output dir
+                        tex_name = os.path.basename(tex_path)
+                        dest = output_dir / tex_name
+                        if not dest.exists():
+                            shutil.copy2(tex_path, dest)
+
+                        # Write to MTL
+                        if tex_type == "diffuse":
+                            f.write(f"map_Kd {tex_name}\n")
+                        elif tex_type == "normal":
+                            f.write(f"map_Bump {tex_name}\n")
+                        elif tex_type == "roughness":
+                            f.write(f"map_Ns {tex_name}\n")
+                        elif tex_type == "metallic":
+                            f.write(f"map_Ks {tex_name}\n")
+
+                f.write("\n")
+
+        # Write OBJ
+        with open(obj_path, "w") as f:
+            f.write(f"mtllib model.mtl\n\n")
+
+            for v in all_vertices:
+                f.write(f"v {v[0]} {v[1]} {v[2]}\n")
+
+            f.write("\n")
+            for uv in all_uvs:
+                f.write(f"vt {uv[0]} {uv[1]}\n")
+
+            f.write("\n")
+
+            # Group faces by material
+            mat_faces = defaultdict(list)
+            for i, face_data in enumerate(all_faces):
+                mat_faces[face_materials[i]].append(face_data)
+
+            for mat_idx, faces in mat_faces.items():
+                mat_path = (
+                    material_list[mat_idx] if mat_idx < len(material_list) else None
+                )
+                mat_name = mat_path.split("/")[-1] if mat_path else "default_material"
+
+                f.write(f"usemtl {mat_name}\n")
+                for face_data in faces:
+                    v_indices, uv_indices = face_data
+                    # OBJ indices are 1-based
+                    if uv_indices[0] is not None:
+                        # Include UV indices
+                        f.write(
+                            f"f {v_indices[0]+1}/{uv_indices[0]+1} {v_indices[1]+1}/{uv_indices[1]+1} {v_indices[2]+1}/{uv_indices[2]+1}\n"
+                        )
+                    else:
+                        # No UVs, just vertex indices
+                        f.write(
+                            f"f {v_indices[0]+1} {v_indices[1]+1} {v_indices[2]+1}\n"
+                        )
+                f.write("\n")
+
+        return obj_path, mtl_path
+
+
+def _install_blender():
+    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_usd(
+    usd_path: str,
+    output_dir: str,
+    num_views: int = 150,
+    resolution: int = 512,
+    verbose: bool = False,
+) -> bool:
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Extract mesh and materials from USD
+    extractor = USDMaterialExtractor(usd_path, verbose=verbose)
+    if not extractor.extract():
+        print(f"Failed to extract from USD: {usd_path}")
+        return False
+
+    # Export to OBJ + MTL
+    temp_dir = tempfile.mkdtemp()
+    try:
+        # Don't normalize - let Blender's normalize_scene() handle it
+        obj_path, mtl_path = extractor.export_obj(Path(temp_dir), normalize=False)
+        if not obj_path:
+            print(f"Failed to export OBJ from USD: {usd_path}")
+            return False
+
+        if verbose:
+            print(f"\nExported to: {obj_path}")
+            # List textures
+            textures = list(Path(temp_dir).glob("*.png")) + list(
+                Path(temp_dir).glob("*.tga")
+            )
+            if textures:
+                print(f"Textures copied: {len(textures)}")
+                for t in textures:
+                    print(f"  - {t.name}")
+
+        # Generate camera views
+        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.5 ensures object corners fit in frame:
+        # Object diagonal at 0.866 from center, visible range at radius=2.5, FOV=40° is ±0.91
+        radius = [2.1] * 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)
+        ]
+
+        # Call Blender
+        blender_script = os.path.join(
+            os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
+            "dataset_toolkits",
+            "blender_script",
+            "render.py",
+        )
+
+        args = [
+            BLENDER_PATH,
+            "-b",
+            "-P",
+            blender_script,
+            "--",
+            "--views",
+            json.dumps(views),
+            "--object",
+            str(obj_path),
+            "--resolution",
+            str(resolution),
+            "--output_folder",
+            output_dir,
+            "--engine",
+            "CYCLES",
+            "--save_mesh",
+            "--use_gpu",  # Enable GPU acceleration
+        ]
+
+        if verbose:
+            print(f"\nRunning Blender...")
+
+        call(
+            args,
+            stdout=DEVNULL if not verbose else None,
+            stderr=DEVNULL if not verbose else None,
+        )
+
+        success = os.path.exists(os.path.join(output_dir, "transforms.json"))
+        return success
+
+    finally:
+        shutil.rmtree(temp_dir, ignore_errors=True)
+
+
+def _render_worker(
+    file_path: str,
+    sha256: str,
+    dataset: str,
+    output_dir: str,
+    num_views: int,
+    quiet: bool,
+) -> Optional[Dict]:
+    output_folder = os.path.join(output_dir, "renders", sha256)
+
+    # Skip if already rendered
+    if os.path.exists(os.path.join(output_folder, "transforms.json")):
+        return {"sha256": sha256, "rendered": True}
+
+    success = render_usd(
+        file_path,
+        output_folder,
+        num_views=num_views,
+        resolution=512,
+        verbose=not quiet,
+    )
+
+    if success:
+        return {"sha256": sha256, "rendered": True}
+    else:
+        if not quiet:
+            print(f"Failed to render: {file_path}")
+        return None
+
+
+def main_batch():
+    import importlib
+    import copy
+    from functools import partial
+    from easydict import EasyDict as edict
+
+    # First argument is dataset type (e.g., "allmats")
+    dataset_utils = importlib.import_module(f"dataset_toolkits.datasets.{sys.argv[1]}")
+
+    parser = argparse.ArgumentParser(
+        description="Batch render USD files with proper texture extraction"
+    )
+    parser.add_argument(
+        "--output_dir", type=str, required=True, help="Directory to save renders"
+    )
+    parser.add_argument(
+        "--filter_low_aesthetic_score",
+        type=float,
+        default=None,
+        help="Filter objects with aesthetic score lower than this value",
+    )
+    parser.add_argument(
+        "--instances",
+        type=str,
+        default=None,
+        help="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(
+        "--rank", type=int, default=0, help="Worker rank for distributed processing"
+    )
+    parser.add_argument(
+        "--world_size",
+        type=int,
+        default=1,
+        help="Total workers for distributed processing",
+    )
+    parser.add_argument(
+        "--max_workers", type=int, default=8, help="Number of parallel workers"
+    )
+    parser.add_argument("--quiet", action="store_true", help="Suppress verbose output")
+
+    # Add dataset-specific args
+    dataset_utils.add_args(parser)
+
+    opt = parser.parse_args(sys.argv[2:])
+    opt = edict(vars(opt))
+
+    os.makedirs(os.path.join(opt.output_dir, "renders"), exist_ok=True)
+
+    # Install blender
+    if not opt.quiet:
+        print("Checking blender...", flush=True)
+    _install_blender()
+
+    # Get file list from 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)
+
+    if opt.instances is None:
+        metadata = metadata[metadata["local_path"].notna()]
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[
+                metadata["aesthetic_score"] >= opt.filter_low_aesthetic_score
+            ]
+        if "rendered" in metadata.columns:
+            metadata = metadata[metadata["rendered"] == False]
+    else:
+        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)]
+
+    # Distributed processing slice
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+    records = []
+
+    # Filter already processed
+    for sha256 in copy.copy(metadata["sha256"].values):
+        if os.path.exists(
+            os.path.join(opt.output_dir, "renders", sha256, "transforms.json")
+        ):
+            records.append({"sha256": sha256, "rendered": True})
+            metadata = metadata[metadata["sha256"] != sha256]
+
+    print(f"Processing {len(metadata)} objects (rank {opt.rank}/{opt.world_size})...")
+
+    # Process objects
+    from concurrent.futures import ThreadPoolExecutor
+    from tqdm import tqdm
+
+    results = []
+    with ThreadPoolExecutor(max_workers=opt.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(
+                    _render_worker,
+                    local_path,
+                    sha256,
+                    dataset,
+                    opt.output_dir,
+                    opt.num_views,
+                    opt.quiet,
+                )
+            )
+
+        for future in tqdm(futures, desc="Rendering", disable=opt.quiet):
+            try:
+                result = future.result()
+                if result is not None:
+                    results.append(result)
+            except Exception as e:
+                if not opt.quiet:
+                    print(f"Error in worker: {e}")
+
+    # Save results
+    rendered = pd.concat(
+        [pd.DataFrame.from_records(results), pd.DataFrame.from_records(records)]
+    )
+    rendered.to_csv(
+        os.path.join(opt.output_dir, f"rendered_{opt.rank}.csv"), index=False
+    )
+
+    print(f"Done! Rendered {len(results)} objects.")
+
+
+def main_single():
+    parser = argparse.ArgumentParser(
+        description="Render a single USD file with proper texture extraction"
+    )
+    parser.add_argument("usd_file", help="Path to USD file")
+    parser.add_argument(
+        "--output_dir",
+        "-o",
+        default=None,
+        help="Output directory (default: /tmp/render_<filename>)",
+    )
+    parser.add_argument(
+        "--num_views", type=int, default=150, help="Number of views to render"
+    )
+    parser.add_argument("--resolution", type=int, default=512, help="Image resolution")
+    parser.add_argument(
+        "--verbose", "-v", action="store_true", help="Print detailed logs"
+    )
+    parser.add_argument(
+        "--extract_only",
+        action="store_true",
+        help="Only extract materials (don't render)",
+    )
+
+    args = parser.parse_args()
+
+    if not os.path.exists(args.usd_file):
+        print(f"ERROR: USD file not found: {args.usd_file}")
+        sys.exit(1)
+
+    if args.extract_only:
+        extractor = USDMaterialExtractor(args.usd_file, verbose=True)
+        extractor.extract()
+
+        print("\n=== SUMMARY ===")
+        print(f"Meshes: {len(extractor.meshes)}")
+        print(f"Materials: {len(extractor.materials)}")
+
+        total_textures = sum(len(t) for t in extractor.materials.values())
+        print(f"Total textures: {total_textures}")
+
+        for mat_path, textures in extractor.materials.items():
+            if textures:
+                mat_name = mat_path.split("/")[-1] if mat_path else "unknown"
+                print(f"\n{mat_name}:")
+                for tex_type, tex_path in textures.items():
+                    print(f"  {tex_type}: {os.path.basename(tex_path)}")
+    else:
+        _install_blender()
+
+        output_dir = args.output_dir
+        if not output_dir:
+            filename = Path(args.usd_file).stem
+            output_dir = f"/tmp/render_{filename}"
+
+        success = render_usd(
+            args.usd_file,
+            output_dir,
+            num_views=args.num_views,
+            resolution=args.resolution,
+            verbose=args.verbose,
+        )
+
+        if success:
+            print(f"\n✓ Rendered to: {output_dir}")
+        else:
+            print(f"\n✗ Rendering failed")
+            sys.exit(1)
+
+
+if __name__ == "__main__":
+    # Check if first arg is a dataset type (batch mode) or a file (single mode)
+    if (
+        len(sys.argv) > 1
+        and not sys.argv[1].startswith("-")
+        and not os.path.exists(sys.argv[1])
+    ):
+        # Batch mode: first arg is dataset type like "allmats"
+        main_batch()
+    else:
+        # Single file mode
+        main_single()

deps/vomp/dataset_toolkits/material_objects/vlm_annotations/data_subsets/commercial.py

@@ -0,0 +1,427 @@
+# 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.
+
+from dataset_toolkits.material_objects.vlm_annotations.utils.utils import (
+    COMMERCIAL_BASE_DIR,
+)
+from dataset_toolkits.material_objects.vlm_annotations.utils.render import (
+    render_sphere_with_texture,
+)
+from dataset_toolkits.material_objects.vlm_annotations.utils.vlm import (
+    analyze_material_with_vlm,
+    parse_vlm_properties,
+)
+from dataset_toolkits.material_objects.vlm_annotations.data_subsets.common import (
+    extract_materials_from_usd,
+)
+from dataset_toolkits.material_objects.vlm_annotations.data_subsets.residential import (
+    PROMPTS,
+    make_user_prompt,
+)
+import re
+from tqdm import tqdm
+import os
+import logging
+import copy
+
+# Use the centralized parser function
+parse_vlm_output = parse_vlm_properties
+
+
+def list_commercial_objects():
+    """
+    List all available commercial objects in the commercial directory.
+    """
+    usd_files = []
+    print("\nAvailable commercial objects:")
+    for root, _, files in os.walk(COMMERCIAL_BASE_DIR):
+        for file in files:
+            if file.endswith(".usd") and not os.path.basename(root).startswith("."):
+                usd_files.append(os.path.join(root, file))
+                print(f"  - {os.path.basename(root)}/{file}")
+    print()
+
+
+def process_commercial(
+    vlm_model,
+    vlm_processor,
+    limit=None,
+    processed_objects=None,
+    output_file=None,
+    existing_results=None,
+):
+    usd_files = []
+    for root, _, files in os.walk(COMMERCIAL_BASE_DIR):
+        for file in files:
+            if file.endswith(".usd") and not os.path.basename(root).startswith("."):
+                usd_files.append(os.path.join(root, file))
+
+    logging.info(f"Found {len(usd_files)} USD files in commercial dataset")
+
+    # Initialize tracking sets and results
+    processed_objects = set() if processed_objects is None else processed_objects
+    existing_results = [] if existing_results is None else existing_results
+
+    # Build a set of already processed object names from existing_results
+    existing_object_names = {
+        result.get("object_name")
+        for result in existing_results
+        if "object_name" in result
+    }
+    logging.info(
+        f"Found {len(existing_object_names)} already processed objects in existing results"
+    )
+
+    # Add names from existing_results to processed_objects to avoid reprocessing
+    processed_objects.update(existing_object_names)
+
+    # Create a copy of existing_results to avoid modifying the original
+    all_results = copy.deepcopy(existing_results)
+
+    usd_files.sort()
+
+    if limit and limit > 0:
+        usd_files = usd_files[:limit]
+
+    success_count = 0
+    failed_objects = []
+    total_segments = 0
+    unique_materials = set()
+    materials_per_object = {}
+    total_rendered_segments = 0
+    total_vlm_segments = 0
+
+    # Count total segments from existing results
+    for result in existing_results:
+        total_segments += len(result.get("segments", {}))
+
+    # Statistics for texture availability
+    segments_with_texture = 0
+    segments_without_texture = 0
+    segments_with_thumbnail_only = 0
+
+    # Track processed files to avoid duplicates from the same directory
+    processed_files = set()
+
+    for usd_file in tqdm(usd_files, desc=f"Processing commercial dataset"):
+        # Extract object name from path
+        object_name = os.path.basename(os.path.dirname(usd_file))
+
+        # Skip if we already processed this exact file
+        if usd_file in processed_files:
+            continue
+
+        # Skip objects that have already been processed
+        if object_name in processed_objects:
+            logging.info(f"Skipping already processed object: {object_name}")
+            continue
+
+        try:
+            directory = os.path.dirname(usd_file)
+
+            # Extract material information
+            result = extract_materials_from_usd(usd_file, "commercial")
+
+            if result:
+                # Add to processed_files to avoid duplicates
+                processed_files.add(usd_file)
+
+                # Track statistics
+                segments = result.get("segments", {})
+                total_segments += len(segments)
+
+                # Remove object_name and note fields from segments
+                for segment_key, segment_info in segments.items():
+                    if "object_name" in segment_info:
+                        del segment_info["object_name"]
+                    if "note" in segment_info:
+                        del segment_info["note"]
+
+                # Count unique materials for this object
+                object_materials = set()
+                for segment_name, segment_info in segments.items():
+                    material_name = segment_info.get("material_type", "unknown")
+                    unique_materials.add(material_name)
+                    object_materials.add(material_name)
+
+                # Record materials per object
+                if len(segments) > 0:
+                    materials_per_object[object_name] = len(object_materials)
+
+                # Get thumbnail path if available
+                thumb_path = None
+                # For commercial dataset, thumbnails are in .thumbs/256x256 directory
+                thumb_dir = os.path.join(
+                    os.path.dirname(usd_file), ".thumbs", "256x256"
+                )
+
+                has_thumbnail = False
+                if os.path.exists(thumb_dir):
+                    # Try to find a thumbnail matching the USD filename
+                    usd_filename = os.path.basename(usd_file)
+                    thumb_candidates = [
+                        # Regular thumbnail
+                        os.path.join(thumb_dir, f"{usd_filename}.png"),
+                        # Auto-generated thumbnail
+                        os.path.join(thumb_dir, f"{usd_filename}.auto.png"),
+                    ]
+
+                    for candidate in thumb_candidates:
+                        if os.path.exists(candidate):
+                            thumb_path = candidate
+                            has_thumbnail = True
+                            logging.info(f"Found thumbnail: {thumb_path}")
+                            break
+
+                # Process VLM for all segments if VLM model is provided
+                os.makedirs("/tmp/vlm", exist_ok=True)
+
+                if vlm_model and len(segments) > 0:
+                    for segment_key, segment_info in segments.items():
+                        textures = segment_info.get("textures", {})
+
+                        # Log texture information for diagnostics
+                        logging.info(
+                            f"Segment {segment_key} has textures: {list(textures.keys())}"
+                        )
+
+                        # Check if we have either a normal or roughness texture for rendering
+                        has_texture = (
+                            "normal" in textures
+                            or "roughness" in textures
+                            or "diffuse" in textures
+                        )
+                        if has_texture:
+                            # Has texture - render sphere and use with thumbnail
+                            segments_with_texture += 1
+                            logging.info(
+                                f"Rendering texture sphere for {object_name}, segment {segment_key}"
+                            )
+
+                            # Set up file path for this segment's rendered sphere
+                            segment_render_path = f"/tmp/vlm/texture_sphere_{object_name}_{segment_key}.png"
+
+                            # Render the textured sphere
+                            try:
+                                rgb_buffer = render_sphere_with_texture(
+                                    textures, segment_render_path
+                                )
+                                logging.info(f"RGB buffer shape: {rgb_buffer.shape}")
+                            except Exception as e:
+                                logging.error(
+                                    f"Error rendering texture for {segment_key}: {str(e)}"
+                                )
+                                segment_render_path = None
+                        else:
+                            # No texture - just use thumbnail
+                            segments_without_texture += 1
+                            segment_render_path = None
+                            logging.info(
+                                f"No texture for {object_name}, segment {segment_key}. Using thumbnail only."
+                            )
+
+                        # Always try to process with VLM, even if no texture
+                        try:
+                            # If we have a thumbnail but no texture, still run VLM with just the thumbnail
+                            if not has_texture and has_thumbnail:
+                                segments_with_thumbnail_only += 1
+                                logging.info(
+                                    f"Using thumbnail only for {object_name}, segment {segment_key}"
+                                )
+
+                            # Don't run VLM if we have neither texture nor thumbnail
+                            if not segment_render_path and not has_thumbnail:
+                                logging.warning(
+                                    f"Skipping VLM for {segment_key} - no texture or thumbnail available"
+                                )
+                                continue
+
+                            # Set semantic usage to segment name but don't store in segment data
+                            semantic_usage = segment_key
+                            temp_object_name = object_name
+
+                            # Create custom prompt based on texture availability
+                            part1 = make_user_prompt(
+                                segment_info["material_type"],
+                                semantic_usage,
+                                temp_object_name,
+                                has_texture_sphere=segment_render_path is not None,
+                            )
+
+                            # Store the custom prompt in material_info but not object_name
+                            segment_info["user_prompt"] = part1
+
+                            # Debug: Log the prompt type based on texture availability
+                            if segment_render_path is not None:
+                                logging.info(
+                                    f"Using prompt WITH texture sphere for {object_name}, segment {segment_key}"
+                                )
+                            else:
+                                logging.info(
+                                    f"Using prompt WITHOUT texture sphere for {object_name}, segment {segment_key}"
+                                )
+                                logging.info(
+                                    f"PROMPT: {part1[:100]}..."
+                                )  # Print just the beginning of the prompt
+
+                            # Create a temporary segment_info with object_name for VLM but don't save to result
+                            temp_segment_info = segment_info.copy()
+                            temp_segment_info["semantic_usage"] = semantic_usage
+                            temp_segment_info["object_name"] = temp_object_name
+
+                            vlm_analysis = analyze_material_with_vlm(
+                                segment_render_path,  # This can be None, in which case only thumbnail is used
+                                temp_segment_info,  # Use temporary copy with object_name
+                                vlm_model,
+                                vlm_processor,
+                                thumbnail_path=thumb_path,
+                                dataset_name="commercial",
+                                PROMPTS=PROMPTS,
+                                make_user_prompt=make_user_prompt,
+                                parse_vlm_output=parse_vlm_output,
+                            )
+
+                            # Add VLM analysis to segment info
+                            if vlm_analysis and "error" not in vlm_analysis:
+                                segment_info["vlm_analysis"] = vlm_analysis.get(
+                                    "vlm_analysis"
+                                )
+
+                                if vlm_analysis.get("youngs_modulus") is not None:
+                                    segment_info["youngs_modulus"] = vlm_analysis.get(
+                                        "youngs_modulus"
+                                    )
+
+                                if vlm_analysis.get("poissons_ratio") is not None:
+                                    segment_info["poissons_ratio"] = vlm_analysis.get(
+                                        "poissons_ratio"
+                                    )
+
+                                if vlm_analysis.get("density") is not None:
+                                    segment_info["density"] = vlm_analysis.get(
+                                        "density"
+                                    )
+
+                                total_vlm_segments += 1
+                                logging.info(
+                                    f"VLM analysis successful for {segment_key}:"
+                                )
+                                logging.info(
+                                    f"  Young's modulus: {vlm_analysis.get('youngs_modulus')}"
+                                )
+                                logging.info(
+                                    f"  Poisson's ratio: {vlm_analysis.get('poissons_ratio')}"
+                                )
+                                logging.info(
+                                    f"  Density: {vlm_analysis.get('density')}"
+                                )
+                            else:
+                                logging.error(
+                                    f"VLM analysis failed for {segment_key}: {vlm_analysis.get('error', 'Unknown error')}"
+                                )
+                        except Exception as e:
+                            import traceback
+
+                            logging.error(
+                                f"Error during VLM analysis for {segment_key}: {str(e)}"
+                            )
+                            logging.error(traceback.format_exc())
+
+                        total_rendered_segments += 1
+
+                all_results.append(result)  # Add to our local copy of results
+                processed_objects.add(object_name)  # Mark as processed
+
+                # Incremental save after each object if output file is provided
+                if output_file:
+                    try:
+                        with open(output_file, "w") as f:
+                            import json
+                            from dataset_toolkits.material_objects.vlm_annotations.data_subsets.common import (
+                                UsdJsonEncoder,
+                            )
+
+                            # Debug save contents
+                            logging.info(
+                                f"Saving checkpoint with {len(all_results)} objects"
+                            )
+
+                            # Ensure result types are JSON serializable
+                            for idx, item in enumerate(all_results):
+                                if "segments" in item:
+                                    for seg_key, seg_info in item["segments"].items():
+                                        # Remove object_name and note fields if they exist
+                                        if "object_name" in seg_info:
+                                            del seg_info["object_name"]
+                                        if "note" in seg_info:
+                                            del seg_info["note"]
+
+                                        if "textures" in seg_info and isinstance(
+                                            seg_info["textures"], dict
+                                        ):
+                                            # Convert any non-serializable texture paths to strings
+                                            serializable_textures = {}
+                                            for tex_type, tex_path in seg_info[
+                                                "textures"
+                                            ].items():
+                                                serializable_textures[tex_type] = str(
+                                                    tex_path
+                                                )
+                                            seg_info["textures"] = serializable_textures
+
+                            # Dump to file
+                            json.dump(all_results, f, indent=4, cls=UsdJsonEncoder)
+
+                    except Exception as e:
+                        logging.error(f"Error saving checkpoint: {str(e)}")
+                        import traceback
+
+                        logging.error(traceback.format_exc())
+
+                success_count += 1
+            else:
+                logging.warning(f"No material information extracted for {usd_file}")
+                failed_objects.append(object_name)
+        except Exception as e:
+            import traceback
+
+            logging.error(f"Error processing {usd_file}: {str(e)}")
+            logging.error(traceback.format_exc())
+            failed_objects.append(os.path.basename(os.path.dirname(usd_file)))
+
+    # Log texture statistics
+    logging.info("Texture Statistics:")
+    logging.info(f"  Total segments processed: {total_segments}")
+    logging.info(f"  Segments with textures: {segments_with_texture}")
+    logging.info(f"  Segments without textures: {segments_without_texture}")
+    logging.info(f"  Segments with thumbnail only: {segments_with_thumbnail_only}")
+    logging.info(f"  Total VLM analyses completed: {total_vlm_segments}")
+
+    # Convert materials_per_object to list format for consistency with simready
+    materials_per_object_list = []
+    for obj_name, count in materials_per_object.items():
+        materials_per_object_list.append(obj_name)
+
+    return (
+        all_results,
+        len(usd_files),
+        success_count,
+        failed_objects,
+        total_segments,
+        total_rendered_segments,
+        total_vlm_segments,
+        list(unique_materials),
+        materials_per_object_list,
+    )

deps/vomp/dataset_toolkits/material_objects/vlm_annotations/data_subsets/common.py

@@ -0,0 +1,1457 @@
+# 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 json
+import glob
+import argparse
+import numpy as np
+import logging
+import copy
+from PIL import Image
+from tqdm import tqdm
+from pxr import Usd, UsdGeom, UsdShade, Sdf, Ar
+from dataset_toolkits.material_objects.vlm_annotations.utils.utils import (
+    COMMERCIAL_BASE_DIR,
+    RESIDENTIAL_BASE_DIR,
+    VEGETATION_BASE_DIR,
+)
+import datetime
+import uuid
+
+
+class UsdJsonEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if hasattr(obj, "__dict__"):
+            return obj.__dict__
+        return str(obj)
+
+
+def find_textures_for_material(object_dir, texture_path):
+    """
+    Find textures referenced by a material in a USD file.
+
+    Args:
+        object_dir (str): Directory containing the USD file
+        texture_path (str): Texture path from the USD file
+
+    Returns:
+        dict: Dictionary mapping texture types to full paths
+    """
+    if not texture_path:
+        return {}
+
+    # Convert Sdf.AssetPath to string if needed
+    if (
+        hasattr(texture_path, "__class__")
+        and texture_path.__class__.__name__ == "AssetPath"
+    ):
+        texture_path = str(texture_path)
+
+    # Handle absolute paths
+    if os.path.isabs(texture_path):
+        if os.path.exists(texture_path):
+            return {determine_texture_type(texture_path): texture_path}
+        return {}
+
+    # Handle relative paths with various prefixes
+    clean_path = texture_path.replace("@", "").replace("./", "")
+
+    # Try direct path
+    direct_path = os.path.join(object_dir, clean_path)
+    if os.path.exists(direct_path):
+        return {determine_texture_type(direct_path): direct_path}
+
+    # Try common texture directories
+    texture_dirs = []
+    for texture_dir_name in [
+        "textures",
+        "Textures",
+        "materials/textures",
+        "Materials/Textures",
+    ]:
+        texture_dir = os.path.join(object_dir, texture_dir_name)
+        if os.path.isdir(texture_dir):
+            texture_dirs.append(texture_dir)
+
+    # Look in parent directory if object_dir doesn't have textures
+    if not texture_dirs:
+        parent_dir = os.path.dirname(object_dir)
+        for texture_dir_name in [
+            "textures",
+            "Textures",
+            "materials/textures",
+            "Materials/Textures",
+        ]:
+            texture_dir = os.path.join(parent_dir, texture_dir_name)
+            if os.path.isdir(texture_dir):
+                texture_dirs.append(texture_dir)
+
+    # Check for texture in each texture directory
+    for texture_dir in texture_dirs:
+        texture_file = os.path.join(texture_dir, os.path.basename(clean_path))
+        if os.path.exists(texture_file):
+            return {determine_texture_type(texture_file): texture_file}
+
+    return {}
+
+
+def determine_texture_type(texture_path):
+    """
+    Determine the type of texture based on its filename.
+
+    Args:
+        texture_path (str): Path to texture file
+
+    Returns:
+        str: Texture type (albedo, normal, roughness, metallic, orm)
+    """
+    filename = os.path.basename(texture_path).lower()
+
+    # Check for common texture type indicators in filename
+    if any(
+        term in filename
+        for term in ["albedo", "basecolor", "color", "_a.", "_a_", "_diffuse", "_diff"]
+    ):
+        return "albedo"
+    elif any(term in filename for term in ["normal", "nrm", "_n.", "_n_"]):
+        return "normal"
+    elif any(term in filename for term in ["roughness", "rough", "_r.", "_r_"]):
+        return "roughness"
+    elif any(term in filename for term in ["metallic", "metal", "_m.", "_m_"]):
+        return "metallic"
+    elif any(term in filename for term in ["orm", "arm", "occlusion"]):
+        return "orm"
+    elif any(term in filename for term in ["emissive", "emission", "_e."]):
+        return "emissive"
+    elif any(term in filename for term in ["opacity", "transparent", "alpha"]):
+        return "opacity"
+    elif any(term in filename for term in ["specular", "spec", "_s."]):
+        return "specular"
+    elif any(term in filename for term in ["displacement", "height", "bump"]):
+        return "displacement"
+
+    # If no specific type is identified, make an educated guess based on file extension
+    ext = os.path.splitext(filename)[1].lower()
+    if ext in [".jpg", ".jpeg", ".png", ".tga", ".tif", ".tiff"]:
+        return "albedo"  # Default to albedo for unrecognized image files
+
+    return "unknown"
+
+
+def copy_texture_to_output(
+    texture_path, output_dir, object_name, material_name, texture_type
+):
+    """
+    Copy a texture file to the output directory with a standardized name.
+
+    Args:
+        texture_path (str): Source texture path
+        output_dir (str): Output directory
+        object_name (str): Name of the object
+        material_name (str): Name of the material
+        texture_type (str): Type of texture
+
+    Returns:
+        str: Path to the copied texture file
+    """
+    if not os.path.exists(texture_path):
+        return None
+
+    # Create output subdirectory for this object if it doesn't exist
+    object_output_dir = os.path.join(output_dir, object_name)
+    os.makedirs(object_output_dir, exist_ok=True)
+
+    # Create standardized output filename
+    texture_ext = os.path.splitext(texture_path)[1]
+    output_filename = f"{material_name}_{texture_type}{texture_ext}"
+    output_path = os.path.join(object_output_dir, output_filename)
+
+    try:
+        # Copy the texture file
+        import shutil
+
+        shutil.copy2(texture_path, output_path)
+        return output_path
+    except Exception as e:
+        logging.error(f"Error copying texture {texture_path}: {str(e)}")
+        return None
+
+
+def extract_material_from_shader(shader_prim, object_dir, dataset_type=None):
+    """
+    Extract material properties and textures from a shader prim.
+
+    Args:
+        shader_prim (UsdShade.Shader): Shader prim
+        object_dir (str): Directory containing the USD file
+        dataset_type (str, optional): Type of dataset (commercial, residential, vegetation)
+
+    Returns:
+        dict: Dictionary with material properties and textures
+    """
+    material_info = {"textures": {}}
+
+    # Create a shader object from the prim
+    shader = UsdShade.Shader(shader_prim)
+    if not shader:
+        logging.warning(f"Failed to create shader from {shader_prim.GetPath()}")
+        return material_info
+
+    # Get material name from shader path
+    shader_path = str(shader_prim.GetPath())
+    material_name = None
+    if "/Looks/" in shader_path:
+        material_name = shader_path.split("/Looks/")[1].split("/")[0]
+
+    logging.info(f"Processing shader for material: {material_name}")
+
+    # For vegetation materials, try to find matching textures by material name
+    if dataset_type == "vegetation" and material_name:
+        # Find the materials/textures directory
+        object_dir_parts = object_dir.split(os.sep)
+        trees_dir = None
+        for i in range(len(object_dir_parts)):
+            if object_dir_parts[i] == "Trees":
+                trees_dir = os.sep.join(object_dir_parts[: i + 1])
+                break
+
+        if trees_dir:
+            textures_dir = os.path.join(trees_dir, "materials", "textures")
+            if os.path.exists(textures_dir):
+                material_name_lower = material_name.lower()
+                material_parts = material_name_lower.replace("_", " ").split()
+
+                # Get all texture files in the directory
+                texture_files = [
+                    f
+                    for f in os.listdir(textures_dir)
+                    if f.lower().endswith((".png", ".jpg", ".jpeg", ".tif", ".tiff"))
+                ]
+
+                # Track potential matches for different texture types
+                texture_matches = {
+                    "diffuse": [],
+                    "normal": [],
+                    "roughness": [],
+                    "metallic": [],
+                    "orm": [],
+                }
+
+                # Categorize material into types
+                material_categories = {
+                    "bark": [
+                        "bark",
+                        "trunk",
+                        "wood",
+                        "tree",
+                        "log",
+                        "stump",
+                        "stem",
+                        "branch",
+                        "twig",
+                    ],
+                    "leaf": ["leaf", "leaves", "foliage", "needle", "needles", "frond"],
+                    "flower": [
+                        "flower",
+                        "flowers",
+                        "petal",
+                        "petals",
+                        "bloom",
+                        "blossom",
+                    ],
+                    "fruit": [
+                        "fruit",
+                        "fruits",
+                        "berry",
+                        "berries",
+                        "seed",
+                        "seeds",
+                        "cone",
+                        "cones",
+                    ],
+                    "grass": [
+                        "grass",
+                        "grasses",
+                        "reed",
+                        "reeds",
+                        "sedge",
+                        "rush",
+                        "blade",
+                    ],
+                }
+
+                # Find all applicable categories
+                material_types = []
+                for category, keywords in material_categories.items():
+                    if any(keyword in material_name_lower for keyword in keywords):
+                        material_types.append(category)
+
+                # If we couldn't determine a category from material name, try from object name
+                if not material_types:
+                    object_name = os.path.splitext(os.path.basename(object_dir))[
+                        0
+                    ].lower()
+                    for category, keywords in material_categories.items():
+                        if any(keyword in object_name for keyword in keywords):
+                            material_types.append(category)
+
+                # Still no category? Add generic fallbacks
+                if not material_types:
+                    # Default to bark for most vegetation models
+                    material_types = ["bark"]
+
+                logging.info(
+                    f"Material categories for {material_name}: {material_types}"
+                )
+
+                # Scoring function for texture relevance to material name
+                def score_texture_for_material(texture_name, texture_type):
+                    score = 0
+                    texture_name_lower = texture_name.lower()
+
+                    # Direct material name match (highest priority)
+                    if material_name_lower in texture_name_lower:
+                        score += 200
+
+                    # Match individual parts of material name
+                    for part in material_parts:
+                        if len(part) > 2 and part in texture_name_lower:
+                            score += 50
+
+                    # Match material categories
+                    for material_type in material_types:
+                        # Match exact category name
+                        if material_type in texture_name_lower:
+                            score += 100
+
+                        # Match keywords for this category
+                        for keyword in material_categories.get(material_type, []):
+                            if keyword in texture_name_lower:
+                                score += 40
+
+                    # Correct type suffix
+                    type_suffixes = {
+                        "diffuse": [
+                            "basecolor",
+                            "albedo",
+                            "color",
+                            "diffuse",
+                            "_bc",
+                            "_a",
+                            "_d",
+                        ],
+                        "normal": ["normal", "nrm", "_n", "nor"],
+                        "roughness": ["roughness", "rough", "_r", "rgh"],
+                        "metallic": ["metallic", "metal", "_m", "mtl"],
+                        "orm": ["orm", "arm", "occlusion"],
+                    }
+
+                    for suffix in type_suffixes.get(texture_type, []):
+                        if suffix in texture_name_lower:
+                            score += 40
+
+                    # Boost score for more specific matches (longer texture names probably more specific)
+                    if len(texture_name_lower) > 15:
+                        score += 10
+
+                    # Exact matches for specific materials
+                    if material_name_lower == "bark" and "bark" in texture_name_lower:
+                        score += 50
+                    elif (
+                        material_name_lower == "leaves" and "leaf" in texture_name_lower
+                    ):
+                        score += 50
+                    elif (
+                        material_name_lower == "needle"
+                        and "needle" in texture_name_lower
+                    ):
+                        score += 50
+                    elif (
+                        "trunk" in material_name_lower and "bark" in texture_name_lower
+                    ):
+                        score += 30
+
+                    return score
+
+                # Process each texture file
+                for texture_file in texture_files:
+                    # Determine texture type
+                    texture_type = determine_texture_type(texture_file)
+
+                    # Don't process "unknown" textures
+                    if texture_type == "unknown":
+                        continue
+
+                    # Score this texture for this material
+                    score = score_texture_for_material(texture_file, texture_type)
+
+                    # If it's a good match (score > 0), add to potential matches
+                    if score > 0:
+                        # Convert diffuse type to match our expected naming
+                        if texture_type in ["albedo", "basecolor", "color"]:
+                            texture_type = "diffuse"
+
+                        # Add to matches with score
+                        if texture_type in texture_matches:
+                            texture_matches[texture_type].append((texture_file, score))
+
+                # Sort matches by score and select the best for each type
+                for texture_type, matches in texture_matches.items():
+                    if matches:
+                        # Sort by score (highest first)
+                        matches.sort(key=lambda x: x[1], reverse=True)
+                        best_match = matches[0][0]
+
+                        # Add to material info
+                        texture_path = os.path.join(textures_dir, best_match)
+                        material_info["textures"][texture_type] = texture_path
+                        logging.info(
+                            f"Found {texture_type} texture for {material_name}: {best_match}"
+                        )
+
+                # If we still don't have textures, use fallbacks from generic categories
+                if not any(material_info["textures"].values()):
+                    logging.info(
+                        f"No direct texture matches found for {material_name}, trying category fallbacks"
+                    )
+
+                    # Key textures we need
+                    needed_types = ["diffuse", "normal", "roughness"]
+
+                    # Generic fallbacks by category
+                    fallbacks = {
+                        "bark": {
+                            "diffuse": "pinebark1_basecolor.png",
+                            "normal": "pinebark1_normal.png",
+                            "roughness": "pinebark1_roughness.png",
+                        },
+                        "leaf": {
+                            "diffuse": "oakleaves1_basecolor.png",
+                            "normal": "oakleaves1_normal.png",
+                            "roughness": "oakleaves1_roughness.png",
+                        },
+                        "flower": {
+                            "diffuse": "goldenchain_flowers_basecolor.png",
+                            "normal": "goldenchain_flowers_normal.png",
+                            "roughness": "goldenchain_flowers_roughness.png",
+                        },
+                        "grass": {
+                            "diffuse": "ashleaves1_basecolor.png",
+                            "normal": "ashleaves1_normal.png",
+                            "roughness": "ashleaves1_roughness.png",
+                        },
+                        "needle": {
+                            "diffuse": "spruceneedles_basecolor.png",
+                            "normal": "spruceneedles_normal.png",
+                            "roughness": "spruceneedles_roughness.png",
+                        },
+                    }
+
+                    # Try each category we matched
+                    for material_type in material_types:
+                        if material_type in fallbacks:
+                            for texture_type in needed_types:
+                                if texture_type not in material_info[
+                                    "textures"
+                                ] and fallbacks[material_type].get(texture_type):
+                                    fallback_file = fallbacks[material_type][
+                                        texture_type
+                                    ]
+                                    fallback_path = os.path.join(
+                                        textures_dir, fallback_file
+                                    )
+                                    if os.path.exists(fallback_path):
+                                        material_info["textures"][
+                                            texture_type
+                                        ] = fallback_path
+                                        logging.info(
+                                            f"Using fallback {texture_type} texture for {material_name}: {fallback_file}"
+                                        )
+
+                    # If still missing textures, use bark as an ultimate fallback (most common)
+                    for texture_type in needed_types:
+                        if texture_type not in material_info["textures"]:
+                            fallback_file = fallbacks["bark"][texture_type]
+                            fallback_path = os.path.join(textures_dir, fallback_file)
+                            if os.path.exists(fallback_path):
+                                material_info["textures"][texture_type] = fallback_path
+                                logging.info(
+                                    f"Using ultimate fallback {texture_type} texture for {material_name}: {fallback_file}"
+                                )
+
+    # Check for shader attributes
+    inputs_to_check = [
+        # Common texture inputs
+        "diffuse_color_texture",
+        "inputs:diffuse_color_texture",
+        "normalmap_texture",
+        "inputs:normalmap_texture",
+        "reflectionroughness_texture",
+        "inputs:reflectionroughness_texture",
+        "diffusecolor_texture",
+        "inputs:diffusecolor_texture",
+        "normal_texture",
+        "inputs:normal_texture",
+        "roughness_texture",
+        "inputs:roughness_texture",
+        # Common material constants
+        "diffuse_color_constant",
+        "inputs:diffuse_color_constant",
+        "reflection_roughness_constant",
+        "inputs:reflection_roughness_constant",
+        "metallic_constant",
+        "inputs:metallic_constant",
+        "opacity_constant",
+        "inputs:opacity_constant",
+        "emissive_color_constant",
+        "inputs:emissive_color_constant",
+    ]
+
+    # Process each input attribute
+    for input_name in inputs_to_check:
+        # Remove "inputs:" prefix if present
+        input_name_clean = input_name.replace("inputs:", "")
+
+        # Try to get the input
+        shader_input = shader.GetInput(input_name_clean)
+        if not shader_input:
+            continue
+
+        # Get the value
+        value = shader_input.Get()
+        if value is None:
+            continue
+
+        # Format input name to standard form
+        standard_name = input_name_clean.lower()
+
+        # Check if this is a texture input
+        if "texture" in standard_name:
+            # Determine texture type
+            if "normal" in standard_name:
+                texture_type = "normal"
+            elif "rough" in standard_name:
+                texture_type = "roughness"
+            elif "diffuse" in standard_name or "color" in standard_name:
+                texture_type = "diffuse"
+            elif "specular" in standard_name:
+                texture_type = "specular"
+            elif "metallic" in standard_name:
+                texture_type = "metallic"
+            elif "opacity" in standard_name:
+                texture_type = "opacity"
+            elif "emissive" in standard_name:
+                texture_type = "emissive"
+            else:
+                texture_type = "other"
+
+            # Handle asset path values
+            if isinstance(value, Sdf.AssetPath):
+                texture_path = value.resolvedPath
+                if not texture_path:
+                    # Try to resolve relative path
+                    rel_path = value.path
+                    if rel_path.startswith("./"):
+                        rel_path = rel_path[2:]
+                    texture_path = os.path.join(object_dir, rel_path)
+
+                if os.path.exists(texture_path):
+                    # If we already found a texture through our material name matching,
+                    # don't override it for vegetation materials
+                    if (
+                        dataset_type == "vegetation"
+                        and texture_type in material_info["textures"]
+                    ):
+                        logging.info(
+                            f"Keeping already found {texture_type} texture for {material_name}"
+                        )
+                    else:
+                        material_info["textures"][texture_type] = texture_path
+
+            # For vegetation, try to find exact textures by material name
+            if (
+                dataset_type == "vegetation"
+                and not material_info["textures"].get(texture_type)
+                and material_name
+            ):
+                logging.info(
+                    f"Looking for exact vegetation texture: {texture_type} for {material_name}"
+                )
+
+                # Find the materials/textures directory
+                object_dir_parts = object_dir.split(os.sep)
+                trees_dir = None
+                for i in range(len(object_dir_parts)):
+                    if object_dir_parts[i] == "Trees":
+                        trees_dir = os.sep.join(object_dir_parts[: i + 1])
+                        break
+
+                if trees_dir:
+                    materials_dir = os.path.join(trees_dir, "materials")
+                    textures_dir = os.path.join(materials_dir, "textures")
+
+                    logging.info(f"Looking for textures in: {textures_dir}")
+
+                    if os.path.exists(textures_dir):
+                        # Look for textures with material name
+                        material_name_lower = material_name.lower()
+
+                        # Build specific patterns for this material name
+                        specific_patterns = [
+                            f"{material_name_lower}_{texture_type}.png",
+                            f"{material_name_lower.replace('_', '')}_{texture_type}.png",
+                        ]
+
+                        # Try alternate texture type names for diffuse
+                        if texture_type == "diffuse":
+                            specific_patterns.extend(
+                                [
+                                    f"{material_name_lower}_basecolor.png",
+                                    f"{material_name_lower.replace('_', '')}_basecolor.png",
+                                    f"{material_name_lower}_albedo.png",
+                                    f"{material_name_lower.replace('_', '')}_albedo.png",
+                                ]
+                            )
+
+                        # Search for exact matches only
+                        for pattern in specific_patterns:
+                            potential_file = os.path.join(textures_dir, pattern)
+                            if os.path.exists(potential_file):
+                                logging.info(
+                                    f"Found exact vegetation texture: {os.path.basename(potential_file)}"
+                                )
+                                material_info["textures"][texture_type] = potential_file
+                                break
+
+                        # If exact match not found, try partial matches
+                        if not material_info["textures"].get(texture_type):
+                            for file in os.listdir(textures_dir):
+                                file_lower = file.lower()
+                                if (
+                                    file_lower.endswith(".png")
+                                    and material_name_lower in file_lower
+                                ):
+                                    # Check for texture type in filename
+                                    if texture_type in file_lower or (
+                                        texture_type == "diffuse"
+                                        and "basecolor" in file_lower
+                                    ):
+                                        full_path = os.path.join(textures_dir, file)
+                                        logging.info(
+                                            f"Found related vegetation texture: {file}"
+                                        )
+                                        material_info["textures"][
+                                            texture_type
+                                        ] = full_path
+                                        break
+        else:
+            # Handle non-texture attributes
+            material_info[standard_name] = value
+
+    return material_info
+
+
+def apply_generic_textures_to_segments(
+    segments, object_name, object_dir, output_textures_dir=None
+):
+    """
+    Apply generic textures to mesh segments that don't have textures.
+
+    Args:
+        segments (dict): Segments dictionary to update
+        object_name (str): Name of the object
+        object_dir (str): Directory containing the USD file
+        output_textures_dir (str, optional): Directory to save extracted textures
+
+    Returns:
+        dict: Updated segments dictionary
+    """
+    # Skip if no segments
+    if not segments:
+        return segments
+
+    # Find the materials/textures directory
+    object_dir_parts = object_dir.split(os.sep)
+    trees_dir = None
+    shrub_dir = None
+    debris_dir = None
+
+    # Look for Trees directory
+    for i in range(len(object_dir_parts)):
+        if object_dir_parts[i] == "Trees":
+            trees_dir = os.sep.join(object_dir_parts[: i + 1])
+            break
+
+    # Look for Shrub directory
+    for i in range(len(object_dir_parts)):
+        if object_dir_parts[i] == "Shrub":
+            shrub_dir = os.sep.join(object_dir_parts[: i + 1])
+            break
+
+    # Look for Debris directory
+    for i in range(len(object_dir_parts)):
+        if object_dir_parts[i] == "Debris":
+            debris_dir = os.sep.join(object_dir_parts[: i + 1])
+            break
+
+    # Set up textures directory based on dataset subdirectory found
+    if trees_dir:
+        textures_dir = os.path.join(trees_dir, "materials", "textures")
+    elif shrub_dir:
+        textures_dir = os.path.join(shrub_dir, "materials", "textures")
+    elif debris_dir:
+        textures_dir = os.path.join(debris_dir, "materials", "textures")
+    else:
+        # Check for Plant_Tropical directory
+        tropical_dir = None
+        for i in range(len(object_dir_parts)):
+            if object_dir_parts[i] == "Plant_Tropical":
+                tropical_dir = os.sep.join(object_dir_parts[: i + 1])
+                break
+
+        if tropical_dir:
+            textures_dir = os.path.join(tropical_dir, "materials", "textures")
+        else:
+            # Try looking for material textures directory in current location
+            textures_dir = os.path.join(object_dir, "materials", "textures")
+            if not os.path.exists(textures_dir):
+                # Go up one directory and look there
+                parent_dir = os.path.dirname(object_dir)
+                textures_dir = os.path.join(parent_dir, "materials", "textures")
+                if not os.path.exists(textures_dir):
+                    # Try root vegetation directory as a last resort
+                    veg_root = None
+                    for i in range(len(object_dir_parts)):
+                        if object_dir_parts[i] == "vegetation":
+                            veg_root = os.sep.join(object_dir_parts[: i + 1])
+                            break
+                    if veg_root:
+                        textures_dir = os.path.join(veg_root, "materials", "textures")
+
+    # If no textures directory found, return segments unchanged
+    if not os.path.exists(textures_dir):
+        return segments
+
+    # Categorize object by name
+    object_name_lower = object_name.lower()
+    object_categories = []
+
+    # Common categories
+    category_keywords = {
+        "tree": [
+            "tree",
+            "pine",
+            "oak",
+            "maple",
+            "birch",
+            "cedar",
+            "ash",
+            "spruce",
+            "poplar",
+            "aspen",
+            "beech",
+            "dogwood",
+            "cypress",
+            "hemlock",
+        ],
+        "palm": ["palm", "frond"],
+        "flower": ["flower", "bloom", "blossom", "rose", "tulip", "lily"],
+        "grass": [
+            "grass",
+            "reed",
+            "sedge",
+            "fern",
+            "bamboo",
+            "pampas",
+            "fountain",
+            "switchgrass",
+        ],
+        "bush": [
+            "bush",
+            "shrub",
+            "boxwood",
+            "barberry",
+            "lilac",
+            "lupin",
+            "daphne",
+            "forsythia",
+            "vibernum",
+            "rhododendron",
+        ],
+    }
+
+    # Determine categories
+    for category, keywords in category_keywords.items():
+        if any(keyword in object_name_lower for keyword in keywords):
+            object_categories.append(category)
+
+    if not object_categories:
+        # Default to tree if no other category matched
+        object_categories = ["tree"]
+
+    # Define generic texture sets for each category and part
+    generic_textures = {
+        "tree": {
+            "bark": {
+                "diffuse": "pinebark1_basecolor.png",
+                "normal": "pinebark1_normal.png",
+                "roughness": "pinebark1_roughness.png",
+            },
+            "leaf": {
+                "diffuse": "oakleaves1_basecolor.png",
+                "normal": "oakleaves1_normal.png",
+                "roughness": "oakleaves1_roughness.png",
+            },
+        },
+        "palm": {
+            "bark": {
+                "diffuse": "bark1_basecolor.png",
+                "normal": "bark1_normal.png",
+                "roughness": "bark1_roughness.png",
+            },
+            "leaf": {
+                "diffuse": "palmleaves_mat_basecolor.png",
+                "normal": "palmleaves_mat_normal.png",
+                "roughness": "palmleaves_mat_roughness.png",
+            },
+        },
+        "flower": {
+            "stem": {
+                "diffuse": "bark2_basecolor.png",
+                "normal": "bark2_normal.png",
+                "roughness": "bark2_roughness.png",
+            },
+            "petal": {
+                "diffuse": "goldenchain_flowers_basecolor.png",
+                "normal": "goldenchain_flowers_normal.png",
+                "roughness": "goldenchain_flowers_roughness.png",
+            },
+        },
+        "grass": {
+            "blade": {
+                "diffuse": "ashleaves1_basecolor.png",
+                "normal": "ashleaves1_normal.png",
+                "roughness": "ashleaves1_roughness.png",
+            }
+        },
+        "bush": {
+            "branch": {
+                "diffuse": "bark3_basecolor.png",
+                "normal": "bark3_normal.png",
+                "roughness": "bark3_roughness.png",
+            },
+            "leaf": {
+                "diffuse": "dogwood_leaf_basecolor.png",
+                "normal": "dogwood_leaf_normal.png",
+                "roughness": "dogwood_leaf_roughness.png",
+            },
+        },
+    }
+
+    # Special material name to texture mappings for problematic cases
+    special_material_textures = {
+        # Special material names
+        "Lupin_m": {
+            "diffuse": "lupin_basecolor.png",
+            "normal": "lupin_normal.png",
+            "roughness": "lupin_roughness.png",
+        },
+        "Dagger_M": {
+            "diffuse": "plantatlas1_basecolor.png",
+            "normal": "plantatlas1_normal.png",
+            "roughness": "plantatlas1_roughness.png",
+        },
+        "bark3": {
+            "diffuse": "bark3_basecolor.png",
+            "normal": "bark3_normal.png",
+            "roughness": "bark3_roughness.png",
+        },
+        "Pampas_flower": {
+            "diffuse": "pampas_flower.png",
+            "normal": "fanpalm_normal.png",  # Fallback normal map
+            "roughness": "fanpalm_roughness.png",  # Fallback roughness map
+        },
+        "FountainGrass": {
+            "diffuse": "fountaingrass_basecolor.png",
+            "normal": "pampas_grass_normal.png",
+            "roughness": "pampas_grass.png",
+        },
+        "TreeBark_01": {
+            "diffuse": "tree_bark_03_diff_2k.png",
+            "normal": "bark1_normal.png",
+            "roughness": "sycamorebark2_roughness.png",
+        },
+        "Barberry": {
+            "diffuse": "barberry_basecolor.png",
+            "normal": "bark3_normal.png",  # Fallback
+            "roughness": "bark3_roughness.png",  # Fallback
+        },
+        "Century_m": {
+            "diffuse": "century_m_basecolor.png",
+            "normal": "Century_m_Normal.png",
+            "roughness": "Century_m_Roughness.png",
+        },
+        "Rhododendron": {
+            "diffuse": "rhododendron_basecolor.png",
+            "normal": "rhododendron_normal.png",
+            "roughness": "rhododendron_roughness.png",
+        },
+        # Add more problematic materials
+        "Burning_Bush": {
+            "diffuse": "burningbush_leaf_basecolor.png",
+            "normal": "burningbush_leaf_normal.png",
+            "roughness": "burningbush_leaf_roughness.png",
+        },
+        "Cedar_Shrub": {
+            "diffuse": "pinebark1_basecolor.png",
+            "normal": "pinebark1_normal.png",
+            "roughness": "pinebark1_roughness.png",
+        },
+        "Japanese_Flame": {
+            "diffuse": "japaneseflame_basecolor.png",
+            "normal": "japaneseflame_normal.png",
+            "roughness": "japaneseflame_roughness.png",
+        },
+        "Honey_Myrtle": {
+            "diffuse": "honeymyrtle_basecolor.png",
+            "normal": "hollyprivet_normal.png",  # Fallback
+            "roughness": "hollyprivet_roughness.png",  # Fallback
+        },
+        "Hurricane_Palm_bark_Mat": {
+            "diffuse": "hurricanepalm_bark_basecolor.png",
+            "normal": "hurricanepalm_bark_normal.png",
+            "roughness": "hurricanepalm_bark_roughness.png",
+        },
+        "Australian_Fern_leaves_Mat": {
+            "diffuse": "australianfern_leaves_basecolor.png",
+            "normal": "australianfern_leaves_normal.png",
+            "roughness": "australianfern_leaves_roughness.png",
+        },
+        "Australian_Fern_trunk": {
+            "diffuse": "australianfern_trunk_basecolor.png",
+            "normal": "australianfern_trunk_normal.png",
+            "roughness": "australianfern_trunk_roughness.png",
+        },
+        "Agave_mat": {
+            "diffuse": "agave_basecolor.png",
+            "normal": "agave_normal.png",
+            "roughness": "Agave_Roughness.png",
+        },
+        "Bamboo_leaf_Mat1": {
+            "diffuse": "bambooleaf_basecolor.png",
+            "normal": "bambooleaf_normal.png",
+            "roughness": "bambooleaf_roughness.png",
+        },
+        "Bamboo_shoot_Mat1": {
+            "diffuse": "bambooshoot_basecolor.png",
+            "normal": "bambooshoot_normal.png",
+            "roughness": "bambooshoot_roughness.png",
+        },
+        "CraneLily_mat": {
+            "diffuse": "cranelily_basecolor.png",
+            "normal": "cranelily_normal.png",
+            "roughness": "cranelily_roughness.png",
+        },
+        "CraneLily_mat_2": {
+            "diffuse": "cranelily_basecolor.png",
+            "normal": "cranelily_normal.png",
+            "roughness": "cranelily_roughness.png",
+        },
+        "CraneLily_mat_3": {
+            "diffuse": "cranelily_basecolor.png",
+            "normal": "cranelily_normal.png",
+            "roughness": "cranelily_roughness.png",
+        },
+        "GrassPalm_bark": {
+            "diffuse": "grasspalm_bark_basecolor.png",
+            "normal": "grasspalm_bark_normal.png",
+            "roughness": "grasspalm_bark_roughness.png",
+        },
+        "GrassPalm_leaves": {
+            "diffuse": "grasspalm_leaves_basecolor.png",
+            "normal": "grasspalm_leaves_normal.png",
+            "roughness": "grasspalm_leaves_roughness.png",
+        },
+    }
+
+    # First try to apply special material textures based on material name in each segment
+    for segment_key, segment_info in segments.items():
+        if segment_info is None:
+            continue
+
+        # Skip segments that already have textures
+        if segment_info.get("textures") and len(segment_info["textures"]) > 0:
+            continue
+
+        # Initialize textures dict if needed
+        if "textures" not in segment_info:
+            segment_info["textures"] = {}
+
+        # Get material name
+        material_name = segment_info.get("name", "")
+
+        # Check for special material name mapping
+        if material_name in special_material_textures:
+            for texture_type, texture_file in special_material_textures[
+                material_name
+            ].items():
+                texture_path = os.path.join(textures_dir, texture_file)
+                if os.path.exists(texture_path):
+                    segment_info["textures"][texture_type] = texture_path
+
+                    # Copy texture if needed
+                    if output_textures_dir:
+                        copied_path = copy_texture_to_output(
+                            texture_path,
+                            output_textures_dir,
+                            object_name,
+                            material_name,
+                            texture_type,
+                        )
+                        if copied_path:
+                            segment_info["textures"][
+                                f"{texture_type}_copied"
+                            ] = copied_path
+
+        # If we found specific textures for this segment, continue to next segment
+        if segment_info.get("textures") and len(segment_info["textures"]) > 0:
+            continue
+
+        # Apply category-based textures if specific ones weren't found
+        material_type = segment_info.get("material_type", "")
+        segment_type = "leaf"  # Default
+
+        # Determine segment type
+        if material_type in ["bark", "trunk", "stem", "branch", "stalk"]:
+            segment_type = (
+                "bark" if "bark" in generic_textures[object_categories[0]] else "branch"
+            )
+        elif material_type in ["leaf", "leaves", "foliage", "needle", "frond"]:
+            segment_type = "leaf"
+        elif material_type in ["petal", "flower", "bloom", "blossom"]:
+            segment_type = "petal"
+        elif material_type in ["blade", "grass"]:
+            segment_type = "blade"
+
+        # Get the right texture set based on object category and segment type
+        for category in object_categories:
+            if (
+                category in generic_textures
+                and segment_type in generic_textures[category]
+            ):
+                texture_set = generic_textures[category][segment_type]
+
+                # Apply textures from set
+                for texture_type, texture_file in texture_set.items():
+                    texture_path = os.path.join(textures_dir, texture_file)
+                    if os.path.exists(texture_path):
+                        segment_info["textures"][texture_type] = texture_path
+
+                        # Copy texture if needed
+                        if output_textures_dir:
+                            copied_path = copy_texture_to_output(
+                                texture_path,
+                                output_textures_dir,
+                                object_name,
+                                material_name or segment_key,
+                                texture_type,
+                            )
+                            if copied_path:
+                                segment_info["textures"][
+                                    f"{texture_type}_copied"
+                                ] = copied_path
+
+                # Break once we found a suitable texture set
+                if segment_info.get("textures") and len(segment_info["textures"]) > 0:
+                    break
+
+        # If we still don't have textures, try to find them by looking for any textures that might match
+        if not segment_info.get("textures") or len(segment_info["textures"]) == 0:
+            # Try to find any textures that might match by name
+            object_dir_lower = object_dir.lower()
+            material_name_lower = material_name.lower() if material_name else ""
+            segment_key_lower = segment_key.lower()
+            object_name_lower = object_name.lower()
+
+            # Look in the textures directory for matching textures
+            for texture_file in os.listdir(textures_dir):
+                texture_lower = texture_file.lower()
+
+                # Try to find matches by object name, material name, or segment key
+                if (
+                    object_name_lower in texture_lower
+                    or material_name_lower in texture_lower
+                    or segment_key_lower in texture_lower
+                ):
+
+                    # Determine texture type
+                    texture_type = None
+                    if "basecolor" in texture_lower or "diffuse" in texture_lower:
+                        texture_type = "diffuse"
+                    elif "normal" in texture_lower:
+                        texture_type = "normal"
+                    elif "roughness" in texture_lower:
+                        texture_type = "roughness"
+
+                    if texture_type:
+                        texture_path = os.path.join(textures_dir, texture_file)
+                        segment_info["textures"][texture_type] = texture_path
+
+                        # Copy texture if needed
+                        if output_textures_dir:
+                            copied_path = copy_texture_to_output(
+                                texture_path,
+                                output_textures_dir,
+                                object_name,
+                                material_name or segment_key,
+                                texture_type,
+                            )
+                            if copied_path:
+                                segment_info["textures"][
+                                    f"{texture_type}_copied"
+                                ] = copied_path
+
+    # If still missing textures, apply default textures
+    for segment_key, segment_info in segments.items():
+        if segment_info is None:
+            continue
+
+        if not segment_info.get("textures"):
+            segment_info["textures"] = {}
+
+        # Check if we're missing any texture types
+        missing_types = []
+        for texture_type in ["diffuse", "normal", "roughness"]:
+            if texture_type not in segment_info["textures"]:
+                missing_types.append(texture_type)
+
+        if not missing_types:
+            continue
+
+        # Determine segment type again
+        material_type = segment_info.get("material_type", "")
+        segment_type = "leaf"  # Default
+
+        if material_type in ["bark", "trunk", "stem", "branch", "stalk"]:
+            segment_type = (
+                "bark" if "bark" in generic_textures[object_categories[0]] else "branch"
+            )
+        elif material_type in ["leaf", "leaves", "foliage", "needle", "frond"]:
+            segment_type = "leaf"
+        elif material_type in ["petal", "flower", "bloom", "blossom"]:
+            segment_type = "petal"
+        elif material_type in ["blade", "grass"]:
+            segment_type = "blade"
+
+        # Apply default textures from the first applicable category
+        for category in object_categories:
+            if (
+                category in generic_textures
+                and segment_type in generic_textures[category]
+            ):
+                for texture_type in missing_types:
+                    if texture_type in generic_textures[category][segment_type]:
+                        texture_file = generic_textures[category][segment_type][
+                            texture_type
+                        ]
+                        texture_path = os.path.join(textures_dir, texture_file)
+
+                        if os.path.exists(texture_path):
+                            segment_info["textures"][texture_type] = texture_path
+
+                            # Copy texture if needed
+                            if output_textures_dir:
+                                copied_path = copy_texture_to_output(
+                                    texture_path,
+                                    output_textures_dir,
+                                    object_name,
+                                    segment_info.get("name", segment_key),
+                                    texture_type,
+                                )
+                                if copied_path:
+                                    segment_info["textures"][
+                                        f"{texture_type}_copied"
+                                    ] = copied_path
+
+                # Break once we've applied textures from a category
+                if all(
+                    texture_type in segment_info["textures"]
+                    for texture_type in missing_types
+                ):
+                    break
+
+    return segments
+
+
+def extract_materials_from_usd(
+    usd_file_path, dataset_type=None, output_textures_dir=None
+):
+    """
+    Extract material information from a USD file with improved handling of material bindings.
+
+    Args:
+        usd_file_path: Path to the USD file
+        dataset_type: Type of dataset (residential, commercial, etc.)
+
+    Returns:
+        Dictionary with material information
+    """
+    logging.info(f"Extracting materials from {usd_file_path}")
+    result = {
+        "object_name": os.path.splitext(os.path.basename(usd_file_path))[0],
+        "dataset_type": dataset_type,
+        "file_path": usd_file_path,
+        "date_processed": datetime.datetime.now().isoformat(),
+        "segments": {},
+    }
+
+    # Open the USD stage
+    try:
+        stage = Usd.Stage.Open(usd_file_path)
+        if not stage:
+            logging.error(f"Could not open USD file: {usd_file_path}")
+            return None
+    except Exception as e:
+        logging.error(f"Error opening USD file {usd_file_path}: {str(e)}")
+        return None
+
+    # Track all materials we find in the stage
+    all_materials = {}
+
+    # First pass: collect all materials and their properties
+    logging.info("First pass: collecting all materials")
+    for prim in stage.Traverse():
+        if prim.IsA(UsdShade.Material):
+            material = UsdShade.Material(prim)
+            material_path = str(prim.GetPath())
+            material_name = prim.GetName()
+
+            # Store material info with default values
+            all_materials[material_path] = {
+                "name": material_name,
+                "material_type": material_name,  # Default to name
+                "textures": {},
+            }
+
+            # Process material's shaders to find textures
+            # Correctly get all the shader prims in this material
+            shader_prims = []
+            for child_prim in Usd.PrimRange(prim):
+                if child_prim.IsA(UsdShade.Shader):
+                    shader_prims.append(child_prim)
+
+            for shader_prim in shader_prims:
+                shader = UsdShade.Shader(shader_prim)
+                if not shader:
+                    continue
+
+                # Inspect shader inputs for textures
+                for input in shader.GetInputs():
+                    input_name = input.GetBaseName()
+
+                    # Check if this input has a connected source that's an asset
+                    if input.HasConnectedSource():
+                        source = input.GetConnectedSource()
+                        if source:
+                            source_shader, source_output, _ = source
+                            source_prim = source_shader.GetPrim()
+
+                            # Check if the source is a texture
+                            if source_prim.IsA(UsdShade.Shader):
+                                source_shader_id = UsdShade.Shader(
+                                    source_prim
+                                ).GetShaderId()
+                                if "texture" in str(source_shader_id).lower():
+                                    # Try to find the file asset path
+                                    for source_input in UsdShade.Shader(
+                                        source_prim
+                                    ).GetInputs():
+                                        if source_input.GetBaseName() in [
+                                            "file",
+                                            "filename",
+                                            "filePath",
+                                            "varname",
+                                        ]:
+                                            asset_path = source_input.Get()
+                                            if asset_path:
+                                                # Determine texture type from connection patterns
+                                                tex_type = "unknown"
+                                                if (
+                                                    "diffuse" in input_name.lower()
+                                                    or "albedo" in input_name.lower()
+                                                    or "color" in input_name.lower()
+                                                ):
+                                                    tex_type = "diffuse"
+                                                elif "normal" in input_name.lower():
+                                                    tex_type = "normal"
+                                                elif "roughness" in input_name.lower():
+                                                    tex_type = "roughness"
+                                                elif "metallic" in input_name.lower():
+                                                    tex_type = "metallic"
+                                                elif "specular" in input_name.lower():
+                                                    tex_type = "specular"
+                                                elif (
+                                                    "displacement" in input_name.lower()
+                                                ):
+                                                    tex_type = "displacement"
+
+                                                # Store texture path
+                                                logging.info(
+                                                    f"Found texture: {tex_type} = {asset_path} for material {material_name}"
+                                                )
+                                                all_materials[material_path][
+                                                    "textures"
+                                                ][tex_type] = str(asset_path)
+
+                    # Direct asset inputs (not connected through other shaders)
+                    elif input.GetTypeName() == "asset":
+                        asset_path = input.Get()
+                        if asset_path:
+                            # Determine texture type from input name
+                            tex_type = "unknown"
+                            if (
+                                "diffuse" in input_name.lower()
+                                or "albedo" in input_name.lower()
+                                or "color" in input_name.lower()
+                            ):
+                                tex_type = "diffuse"
+                            elif "normal" in input_name.lower():
+                                tex_type = "normal"
+                            elif "roughness" in input_name.lower():
+                                tex_type = "roughness"
+                            elif "metallic" in input_name.lower():
+                                tex_type = "metallic"
+                            elif "specular" in input_name.lower():
+                                tex_type = "specular"
+                            elif "displacement" in input_name.lower():
+                                tex_type = "displacement"
+
+                            # Store texture path
+                            logging.info(
+                                f"Found direct texture: {tex_type} = {asset_path} for material {material_name}"
+                            )
+                            all_materials[material_path]["textures"][tex_type] = str(
+                                asset_path
+                            )
+
+    # Second pass: find all material bindings
+    logging.info("Second pass: finding material bindings")
+
+    # Process meshes and their subsets
+    for prim in stage.Traverse():
+        if prim.IsA(UsdGeom.Mesh):
+            mesh = UsdGeom.Mesh(prim)
+            mesh_name = prim.GetName()
+            logging.info(f"Processing mesh: {mesh_name}")
+
+            # First check direct binding on the mesh
+            binding_api = UsdShade.MaterialBindingAPI(prim)
+            direct_binding = binding_api.GetDirectBinding()
+            direct_material = None
+
+            if direct_binding.GetMaterial():
+                direct_material = direct_binding.GetMaterial()
+                material_path = str(direct_material.GetPath())
+                logging.info(f"  Found direct material binding: {material_path}")
+
+                if material_path in all_materials:
+                    # Create segment for the whole mesh
+                    segment_key = f"{mesh_name}_whole"
+                    material_info = all_materials[material_path].copy()
+                    material_info["semantic_usage"] = mesh_name
+
+                    result["segments"][segment_key] = material_info
+                    logging.info(
+                        f"  Created segment {segment_key} with material {material_path}"
+                    )
+
+            # Then check GeomSubsets - these are more specific material assignments
+            imageable = UsdGeom.Imageable(prim)
+            subsets = UsdGeom.Subset.GetGeomSubsets(imageable)
+
+            if subsets:
+                logging.info(f"  Found {len(subsets)} geom subsets for {mesh_name}")
+                for subset in subsets:
+                    subset_prim = subset.GetPrim()
+                    subset_name = subset_prim.GetName()
+                    family = (
+                        subset.GetFamilyNameAttr().Get()
+                        if subset.GetFamilyNameAttr()
+                        else "unknown"
+                    )
+
+                    logging.info(
+                        f"  Processing subset: {subset_name} (Family: {family})"
+                    )
+
+                    # Check material binding on subset
+                    subset_binding_api = UsdShade.MaterialBindingAPI(subset_prim)
+                    subset_direct_binding = subset_binding_api.GetDirectBinding()
+
+                    if subset_direct_binding.GetMaterial():
+                        subset_material = subset_direct_binding.GetMaterial()
+                        subset_material_path = str(subset_material.GetPath())
+                        logging.info(
+                            f"    Found subset material binding: {subset_material_path}"
+                        )
+
+                        if subset_material_path in all_materials:
+                            # Create segment for this subset
+                            segment_key = subset_name
+                            material_info = all_materials[subset_material_path].copy()
+                            material_info["semantic_usage"] = subset_name
+
+                            result["segments"][segment_key] = material_info
+                            logging.info(
+                                f"    Created segment {segment_key} with material {subset_material_path}"
+                            )
+
+            # If no subsets but we have a direct material, use that
+            if not subsets and direct_material:
+                material_path = str(direct_material.GetPath())
+
+                if material_path in all_materials:
+                    # Create segment for the whole mesh
+                    segment_key = mesh_name
+                    material_info = all_materials[material_path].copy()
+                    material_info["semantic_usage"] = mesh_name
+
+                    result["segments"][segment_key] = material_info
+                    logging.info(
+                        f"  No subsets, created segment {segment_key} with material {material_path}"
+                    )
+
+    # Final check - make sure we have segments
+    if not result["segments"]:
+        logging.warning(f"No material segments found in {usd_file_path}")
+
+        # Last resort - add all materials as segments
+        for material_path, material_info in all_materials.items():
+            material_name = material_info["name"]
+            segment_key = f"material_{material_name}"
+
+            result["segments"][segment_key] = material_info.copy()
+            result["segments"][segment_key]["semantic_usage"] = material_name
+
+            logging.info(
+                f"Added material {material_name} as segment {segment_key} (last resort)"
+            )
+
+    logging.info(
+        f"Extracted {len(result['segments'])} material segments from {usd_file_path}"
+    )
+    return result

deps/vomp/dataset_toolkits/material_objects/vlm_annotations/data_subsets/residential.py

@@ -0,0 +1,582 @@
+# 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.
+
+from dataset_toolkits.material_objects.vlm_annotations.utils.utils import (
+    RESIDENTIAL_BASE_DIR,
+)
+from dataset_toolkits.material_objects.vlm_annotations.utils.render import (
+    render_sphere_with_texture,
+)
+from dataset_toolkits.material_objects.vlm_annotations.utils.vlm import (
+    analyze_material_with_vlm,
+    parse_vlm_properties,
+)
+from dataset_toolkits.material_objects.vlm_annotations.data_subsets.common import (
+    extract_materials_from_usd,
+)
+import re
+from tqdm import tqdm
+import os
+import logging
+import copy
+
+PROMPTS = {
+    "few_shot_examples": (
+        """
+Example 1:
+Material: metal
+Usage: structural component
+Object name: SteelBeam
+
+Analysis: 
+Step 1: Based on the images, this appears to be a standard structural steel with a matte gray finish.
+Step 2: The surface has medium roughness with some subtle texture visible in the reflection pattern.
+Step 3: Considering its usage as a structural component, this is likely a carbon steel alloy.
+Step 4: Comparing with reference materials, standard structural steel typically has:
+   - High stiffness (Young's modulus ~200 GPa)
+   - Medium Poisson's ratio typical of metals
+   - High density consistent with iron-based alloys
+
+Young's modulus: 2.0e11 Pa
+Poisson's ratio: 0.29
+Density: 7800 kg/m^3
+
+Example 2:
+Material: plastic
+Usage: household container
+Object name: PlasticContainer
+
+Analysis:
+Step 1: The material shows the characteristic smooth, uniform appearance of a consumer plastic.
+Step 2: It has moderate gloss with some translucency and a slight texture.
+Step 3: Given its household container application, this is likely polypropylene.
+Step 4: The visual and contextual evidence suggests:
+   - Medium-low stiffness typical of polyolefin plastics
+   - Higher Poisson's ratio indicating good lateral deformation
+   - Low-medium density typical of consumer thermoplastics
+
+Young's modulus: 1.3e9 Pa
+Poisson's ratio: 0.42
+Density: 950 kg/m^3
+
+Example 3:
+Material: fabric
+Usage: furniture covering
+Object name: FabricCouch
+
+Analysis:
+Step 1: The material shows a woven textile structure with visible fibers.
+Step 2: The surface has significant texture with a matte appearance and no specular highlights.
+Step 3: As furniture upholstery, this is likely a synthetic or natural fiber blend.
+Step 4: Based on the visual characteristics and usage:
+   - Low stiffness as expected for flexible textiles
+   - Medium-high Poisson's ratio from the woven structure
+   - Low density typical of fibrous materials
+
+Young's modulus: 1.2e8 Pa
+Poisson's ratio: 0.38
+Density: 300 kg/m^3
+
+Example 4:
+Material: organic
+Usage: decorative element
+Object name: DriedLeaf
+
+Analysis:
+Step 1: This is an organic material with the characteristic shape and structure of a dried leaf.
+Step 2: The surface shows visible veins, a matte finish, and a brittle, thin structure.
+Step 3: As a dried leaf, it's a natural cellulose-based composite material.
+Step 4: Considering similar organic materials like paper and dried plant fibers:
+   - Low-medium stiffness in the fiber direction
+   - Medium Poisson's ratio reflecting the fibrous structure
+   - Low density typical of dried plant matter
+
+Young's modulus: 2.5e9 Pa
+Poisson's ratio: 0.30
+Density: 400 kg/m^3
+"""
+    ),
+    "query_prompt": (
+        """
+Based on the provided images and context information, analyze the material properties.
+Note: The material segment might be internal to the object and not visible from the outside.
+
+Respond using EXACTLY the following format (do not deviate from this structure):
+
+Analysis: 
+Step 1: Identify the material class/type based on visual appearance
+Step 2: Describe the surface characteristics (texture, reflectivity, color)
+Step 3: Determine the specific material subtype considering its usage
+Step 4: Reason through each property estimate based on visual and contextual clues
+
+Young's modulus: <value in scientific notation> Pa
+Poisson's ratio: <single decimal value between 0.0 and 0.5>
+Density: <value in scientific notation> kg/m^3
+
+Critical Instructions:
+1. You MUST provide numerical estimates for ALL materials, including organic or unusual materials
+2. For natural materials like leaves, wood, or paper, provide estimates based on similar materials with known properties
+3. Never use "N/A", "unknown", or any non-numeric responses for the material properties
+4. For Poisson's ratio, provide a simple decimal number (like 0.3 or 0.42)
+5. Each property should be on its own line with exact formatting shown above
+"""
+    ),
+}
+
+
+def make_user_prompt(
+    material_type, semantic_usage, object_name, has_texture_sphere=True
+):
+    intro_text = (
+        """
+You are a materials science expert analyzing two images:
+1. A photo of the full object (showing how the material appears in context).
+2. A sphere with the material's texture (showing color/roughness/reflectivity in isolation).
+
+Using both images and the information below, identify the real-world material and estimate its mechanical properties.
+"""
+        if has_texture_sphere
+        else """
+You are a materials science expert analyzing an image of the full object (showing how the material appears in context).
+
+Using this image and the information below, identify the real-world material and estimate its mechanical properties.
+"""
+    )
+
+    return f"""{intro_text}
+Material context:
+  * Material type: {material_type}
+  * Usage: {semantic_usage}
+  * Object name: {object_name}
+
+Your task is to provide three specific properties:
+1. Young's modulus (in Pa using scientific notation)
+2. Poisson's ratio (a value between 0.0 and 0.5)
+3. Density (in kg/m^3 using scientific notation)
+"""
+
+
+# Use the centralized parser function from vlm.py instead
+parse_vlm_output = parse_vlm_properties
+
+
+def list_residential_objects():
+    """
+    List all available residential objects in the residential directory.
+    """
+    usd_files = []
+    print("\nAvailable residential objects:")
+    for root, _, files in os.walk(RESIDENTIAL_BASE_DIR):
+        for file in files:
+            if file.endswith(".usd") and not os.path.basename(root).startswith("."):
+                usd_files.append(os.path.join(root, file))
+                print(f"  - {os.path.basename(root)}/{file}")
+    print()
+
+
+def process_residential(
+    vlm_model,
+    vlm_processor,
+    limit=None,
+    processed_objects=None,
+    output_file=None,
+    existing_results=None,
+):
+    usd_files = []
+    for root, _, files in os.walk(RESIDENTIAL_BASE_DIR):
+        for file in files:
+            if file.endswith(".usd") and not os.path.basename(root).startswith("."):
+                usd_files.append(os.path.join(root, file))
+
+    logging.info(f"Found {len(usd_files)} USD files in residential dataset")
+
+    # Initialize tracking sets and results
+    processed_objects = set() if processed_objects is None else processed_objects
+    existing_results = [] if existing_results is None else existing_results
+
+    # Build a set of already processed object names from existing_results
+    existing_object_names = {
+        result.get("object_name")
+        for result in existing_results
+        if "object_name" in result
+    }
+    logging.info(
+        f"Found {len(existing_object_names)} already processed objects in existing results"
+    )
+
+    # Add names from existing_results to processed_objects to avoid reprocessing
+    processed_objects.update(existing_object_names)
+
+    # Create a copy of existing_results to avoid modifying the original
+    all_results = copy.deepcopy(existing_results)
+
+    usd_files.sort()
+
+    if limit and limit > 0:
+        usd_files = usd_files[:limit]
+
+    success_count = 0
+    failed_objects = []
+    total_segments = 0
+    unique_materials = set()
+    materials_per_object = {}
+    total_rendered_segments = 0
+    total_vlm_segments = 0
+
+    # Count total segments from existing results
+    for result in existing_results:
+        total_segments += len(result.get("segments", {}))
+
+    # Statistics for texture availability
+    segments_with_texture = 0
+    segments_without_texture = 0
+    segments_with_thumbnail_only = 0
+
+    # Track processed files to avoid duplicates from the same directory
+    processed_files = set()
+
+    for usd_file in tqdm(usd_files, desc=f"Processing residential dataset"):
+        # Extract object name from path
+        object_name = os.path.basename(os.path.dirname(usd_file))
+
+        # Skip if we already processed this exact file
+        if usd_file in processed_files:
+            continue
+
+        # Skip objects that have already been processed
+        if object_name in processed_objects:
+            logging.info(f"Skipping already processed object: {object_name}")
+            continue
+
+        try:
+            directory = os.path.dirname(usd_file)
+
+            # Extract material information
+            result = extract_materials_from_usd(usd_file, "residential")
+
+            if result:
+                # Add to processed_files to avoid duplicates
+                processed_files.add(usd_file)
+
+                # Track statistics
+                segments = result.get("segments", {})
+                total_segments += len(segments)
+
+                # Remove object_name and note fields from segments
+                for segment_key, segment_info in segments.items():
+                    if "object_name" in segment_info:
+                        del segment_info["object_name"]
+                    if "note" in segment_info:
+                        del segment_info["note"]
+
+                # Count unique materials for this object
+                object_materials = set()
+                for segment_name, segment_info in segments.items():
+                    material_name = segment_info.get("material_type", "unknown")
+                    unique_materials.add(material_name)
+                    object_materials.add(material_name)
+
+                # Record materials per object
+                if len(segments) > 0:
+                    materials_per_object[object_name] = len(object_materials)
+
+                # Get thumbnail path if available
+                thumb_path = None
+                # For residential dataset, thumbnails are in .thumbs/256x256 directory
+                thumb_dir = os.path.join(
+                    os.path.dirname(usd_file), ".thumbs", "256x256"
+                )
+
+                has_thumbnail = False
+                if os.path.exists(thumb_dir):
+                    # Try to find a thumbnail matching the USD filename
+                    usd_filename = os.path.basename(usd_file)
+                    thumb_candidates = [
+                        # Regular thumbnail
+                        os.path.join(thumb_dir, f"{usd_filename}.png"),
+                        # Auto-generated thumbnail
+                        os.path.join(thumb_dir, f"{usd_filename}.auto.png"),
+                    ]
+
+                    for candidate in thumb_candidates:
+                        if os.path.exists(candidate):
+                            thumb_path = candidate
+                            has_thumbnail = True
+                            logging.info(f"Found thumbnail: {thumb_path}")
+                            break
+
+                # Process VLM for all segments if VLM model is provided
+                os.makedirs("/tmp/vlm", exist_ok=True)
+
+                if vlm_model and len(segments) > 0:
+                    for segment_key, segment_info in segments.items():
+                        textures = segment_info.get("textures", {})
+
+                        # Log texture information for diagnostics
+                        logging.info(
+                            f"Segment {segment_key} has textures: {list(textures.keys())}"
+                        )
+
+                        # Check if we have either a normal or roughness texture for rendering
+                        has_texture = (
+                            "normal" in textures
+                            or "roughness" in textures
+                            or "diffuse" in textures
+                        )
+                        if has_texture:
+                            # Has texture - render sphere and use with thumbnail
+                            segments_with_texture += 1
+                            logging.info(
+                                f"Rendering texture sphere for {object_name}, segment {segment_key}"
+                            )
+
+                            # Set up file path for this segment's rendered sphere
+                            segment_render_path = f"/tmp/vlm/texture_sphere_{object_name}_{segment_key}.png"
+
+                            # Render the textured sphere
+                            try:
+                                rgb_buffer = render_sphere_with_texture(
+                                    textures, segment_render_path
+                                )
+                                logging.info(f"RGB buffer shape: {rgb_buffer.shape}")
+                            except Exception as e:
+                                logging.error(
+                                    f"Error rendering texture for {segment_key}: {str(e)}"
+                                )
+                                segment_render_path = None
+                        else:
+                            # No texture - just use thumbnail
+                            segments_without_texture += 1
+                            segment_render_path = None
+                            logging.info(
+                                f"No texture for {object_name}, segment {segment_key}. Using thumbnail only."
+                            )
+
+                        # Always try to process with VLM, even if no texture
+                        try:
+                            # If we have a thumbnail but no texture, still run VLM with just the thumbnail
+                            if not has_texture and has_thumbnail:
+                                segments_with_thumbnail_only += 1
+                                logging.info(
+                                    f"Using thumbnail only for {object_name}, segment {segment_key}"
+                                )
+
+                            # Don't run VLM if we have neither texture nor thumbnail
+                            if not segment_render_path and not has_thumbnail:
+                                logging.warning(
+                                    f"Skipping VLM for {segment_key} - no texture or thumbnail available"
+                                )
+                                continue
+
+                            # Set semantic usage to segment name but don't store in segment data
+                            semantic_usage = segment_key
+                            temp_object_name = object_name
+
+                            # Create custom prompt based on texture availability
+                            custom_prompt = make_user_prompt(
+                                segment_info["material_type"],
+                                semantic_usage,
+                                temp_object_name,
+                                has_texture_sphere=segment_render_path is not None,
+                            )
+
+                            # Store the custom prompt in material_info but not object_name
+                            segment_info["user_prompt"] = custom_prompt
+
+                            # Debug: Log the prompt type based on texture availability
+                            if segment_render_path is not None:
+                                logging.info(
+                                    f"Using prompt WITH texture sphere for {object_name}, segment {segment_key}"
+                                )
+                            else:
+                                logging.info(
+                                    f"Using prompt WITHOUT texture sphere for {object_name}, segment {segment_key}"
+                                )
+                                logging.info(
+                                    f"PROMPT: {custom_prompt[:100]}..."
+                                )  # Print just the beginning of the prompt
+
+                            # Create a temporary segment_info with object_name for VLM but don't save to result
+                            temp_segment_info = segment_info.copy()
+                            temp_segment_info["semantic_usage"] = semantic_usage
+                            temp_segment_info["object_name"] = temp_object_name
+
+                            vlm_analysis = analyze_material_with_vlm(
+                                segment_render_path,  # This can be None, in which case only thumbnail is used
+                                temp_segment_info,  # Use temporary copy with object_name
+                                vlm_model,
+                                vlm_processor,
+                                thumbnail_path=thumb_path,
+                                dataset_name="residential",
+                                PROMPTS=PROMPTS,
+                                make_user_prompt=make_user_prompt,
+                                parse_vlm_output=parse_vlm_output,
+                            )
+
+                            # Add VLM analysis to segment info
+                            if vlm_analysis and "error" not in vlm_analysis:
+                                segment_info["vlm_analysis"] = vlm_analysis.get(
+                                    "vlm_analysis"
+                                )
+
+                                if vlm_analysis.get("youngs_modulus") is not None:
+                                    segment_info["youngs_modulus"] = vlm_analysis.get(
+                                        "youngs_modulus"
+                                    )
+
+                                if vlm_analysis.get("poissons_ratio") is not None:
+                                    segment_info["poissons_ratio"] = vlm_analysis.get(
+                                        "poissons_ratio"
+                                    )
+
+                                if vlm_analysis.get("density") is not None:
+                                    segment_info["density"] = vlm_analysis.get(
+                                        "density"
+                                    )
+
+                                total_vlm_segments += 1
+                                logging.info(
+                                    f"VLM analysis successful for {segment_key}:"
+                                )
+                                logging.info(
+                                    f"  Young's modulus: {vlm_analysis.get('youngs_modulus')}"
+                                )
+                                logging.info(
+                                    f"  Poisson's ratio: {vlm_analysis.get('poissons_ratio')}"
+                                )
+                                logging.info(
+                                    f"  Density: {vlm_analysis.get('density')}"
+                                )
+                            else:
+                                logging.error(
+                                    f"VLM analysis failed for {segment_key}: {vlm_analysis.get('error', 'Unknown error')}"
+                                )
+                        except Exception as e:
+                            import traceback
+
+                            logging.error(
+                                f"Error during VLM analysis for {segment_key}: {str(e)}"
+                            )
+                            logging.error(traceback.format_exc())
+
+                        total_rendered_segments += 1
+
+                all_results.append(result)  # Add to our local copy of results
+                processed_objects.add(object_name)  # Mark as processed
+
+                # Incremental save after each object if output file is provided
+                if output_file:
+                    try:
+                        with open(output_file, "w") as f:
+                            import json
+                            from dataset_toolkits.material_objects.vlm_annotations.data_subsets.common import (
+                                UsdJsonEncoder,
+                            )
+
+                            # Debug save contents
+                            logging.info(
+                                f"Saving checkpoint with {len(all_results)} objects"
+                            )
+
+                            # Ensure result types are JSON serializable
+                            for idx, item in enumerate(all_results):
+                                if "segments" in item:
+                                    for seg_key, seg_info in item["segments"].items():
+                                        # Remove object_name and note fields if they exist
+                                        if "object_name" in seg_info:
+                                            del seg_info["object_name"]
+                                        if "note" in seg_info:
+                                            del seg_info["note"]
+
+                                        if "textures" in seg_info and isinstance(
+                                            seg_info["textures"], dict
+                                        ):
+                                            # Convert any non-serializable texture paths to strings
+                                            serializable_textures = {}
+                                            for tex_type, tex_path in seg_info[
+                                                "textures"
+                                            ].items():
+                                                serializable_textures[tex_type] = str(
+                                                    tex_path
+                                                )
+                                            seg_info["textures"] = serializable_textures
+
+                            # Try to serialize to a string first to check for issues
+                            try:
+                                json_str = json.dumps(
+                                    all_results, cls=UsdJsonEncoder, indent=4
+                                )
+                                logging.info(
+                                    f"JSON serialization successful, string length: {len(json_str)}"
+                                )
+
+                                # Now write to file
+                                f.write(json_str)
+
+                            except Exception as json_err:
+                                logging.error(
+                                    f"JSON serialization error: {str(json_err)}"
+                                )
+                                # Try to identify problematic objects
+                                for i, item in enumerate(all_results):
+                                    try:
+                                        json.dumps(item, cls=UsdJsonEncoder)
+                                    except Exception as e:
+                                        logging.error(
+                                            f"Error serializing object {i}: {str(e)}"
+                                        )
+                                raise json_err  # Re-raise to be caught by outer exception handler
+
+                    except Exception as e:
+                        logging.error(f"Error saving checkpoint: {str(e)}")
+                        import traceback
+
+                        logging.error(traceback.format_exc())
+
+                success_count += 1
+            else:
+                logging.warning(f"No material information extracted for {usd_file}")
+                failed_objects.append(object_name)
+        except Exception as e:
+            import traceback
+
+            logging.error(f"Error processing {usd_file}: {str(e)}")
+            logging.error(traceback.format_exc())
+            failed_objects.append(os.path.basename(os.path.dirname(usd_file)))
+
+    # Convert materials_per_object to list format for consistency with simready
+    materials_per_object_list = []
+    for obj_name, count in materials_per_object.items():
+        materials_per_object_list.append(obj_name)
+
+    # Log texture statistics
+    logging.info("Texture Statistics:")
+    logging.info(f"  Total segments processed: {total_segments}")
+    logging.info(f"  Segments with textures: {segments_with_texture}")
+    logging.info(f"  Segments without textures: {segments_without_texture}")
+    logging.info(f"  Segments with thumbnail only: {segments_with_thumbnail_only}")
+    logging.info(f"  Total VLM analyses completed: {total_vlm_segments}")
+
+    return (
+        all_results,
+        len(usd_files),
+        success_count,
+        failed_objects,
+        total_segments,
+        total_rendered_segments,
+        total_vlm_segments,
+        list(unique_materials),
+        materials_per_object_list,
+    )