app.py

import os
import shutil
import tempfile
from typing import Any, 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 LazyLoadDino, 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 _preload_dino(model: Vomp) -> None:
    print("Preloading DINO model...")
    dino = LazyLoadDino(
        model_name="dinov2_vitl14_reg",
        device=model.device,
        use_trt=getattr(model, "use_trt", False),
    )
    _ = dino.get_model()
    _ = dino.get_transform()


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!")
_preload_dino(model)


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


_SH_C0 = 0.28209479177387814
_SPLAT_POINT_SCALE = 0.0015


def _write_property_splat_ply(
    coords: np.ndarray,
    values: np.ndarray,
    output_path: str,
    colormap: str = "viridis",
    point_scale: float = _SPLAT_POINT_SCALE,
) -> str:
    """Write a property-colored point cloud as a 3D Gaussian Splatting .ply.
    """
    coords = np.asarray(coords, dtype=np.float32)
    values = np.asarray(values, dtype=np.float32).reshape(-1)
    if coords.ndim != 2 or coords.shape[1] != 3:
        raise ValueError(f"coords must be (N,3), got {coords.shape}")
    if values.shape[0] != coords.shape[0]:
        raise ValueError(
            f"values must be (N,), got {values.shape} for coords {coords.shape}"
        )

    vmin, vmax = float(values.min()), float(values.max())
    if vmax - vmin > 1e-12:
        norm = (values - vmin) / (vmax - vmin)
    else:
        norm = np.zeros_like(values)
    rgb = plt.cm.get_cmap(colormap)(norm)[:, :3].astype(np.float32)  # 0..1

    n = coords.shape[0]
    fields = [
        "x", "y", "z", "nx", "ny", "nz",
        "f_dc_0", "f_dc_1", "f_dc_2",
        "opacity", "scale_0", "scale_1", "scale_2",
        "rot_0", "rot_1", "rot_2", "rot_3",
    ]
    arr = np.zeros((n, len(fields)), dtype=np.float32)
    arr[:, 0:3] = coords
    arr[:, 6:9] = (rgb - 0.5) / _SH_C0
    arr[:, 9] = 6.0
    arr[:, 10:13] = np.log(point_scale)
    arr[:, 13] = 1.0

    header = (
        "ply\nformat binary_little_endian 1.0\n"
        f"element vertex {n}\n"
        + "".join(f"property float {f}\n" for f in fields)
        + "end_header\n"
    )
    with open(output_path, "wb") as fp:
        fp.write(header.encode("ascii"))
        fp.write(arr.tobytes())
    return output_path


def _create_material_visualizations(
    material_file: str, output_dir: str
) -> Dict[str, Tuple[Any, 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:
            ply_path = os.path.join(output_dir, f"{prop_name}_cloud.ply")
            _write_property_splat_ply(coords_normalized, prop_data, ply_path)
            colorbar_path = os.path.join(output_dir, f"{prop_name}_colorbar.png")
            _create_colorbar(prop_data, prop_name, colorbar_path)
            result[prop_name] = (ply_path, colorbar_path)
            print(f"Created point cloud for {prop_name}")

    return result


@spaces.GPU(duration=60)
@torch.no_grad()
def process_3d_model(input_file):
    empty_result = (
        None,  # youngs_cloud
        None,  # youngs_colorbar
        None,  # poissons_cloud
        None,  # poissons_colorbar
        None,  # density_cloud
        None,  # density_colorbar
        None,  # materials file
    )

    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,
                output_dir=output_dir,
                seed=42,
            )
        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}")

        visualizations = _create_material_visualizations(material_file, output_dir)

        youngs_cloud = visualizations.get("youngs_modulus", (None, None))[0]
        youngs_colorbar = visualizations.get("youngs_modulus", (None, None))[1]

        poissons_cloud = visualizations.get("poissons_ratio", (None, None))[0]
        poissons_colorbar = visualizations.get("poissons_ratio", (None, None))[1]

        density_cloud = visualizations.get("density", (None, None))[0]
        density_colorbar = visualizations.get("density", (None, None))[1]

        return (
            youngs_cloud,
            youngs_colorbar,
            poissons_cloud,
            poissons_colorbar,
            density_cloud,
            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)}")


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

.main-content {
    display: flex;
    flex-direction: column-reverse;
    gap: 16px;
}
"""

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) to predict volumetric mechanical properties (Young's modulus, Poisson's ratio, density) for realistic physics simulation.

### Tips for Best Results

- Make sure the input asset has textures
- Make sure the input asset is oriented y-up for best results
- Download the `.npz` artifact to inspect the full-resolution outputs with the viewer we provide in the code release
- The demo is not representative of the time taken to run inference, please run our codebase for significantlybetter performance
- We support other representations in our codebase, inlcudig arbitrary custom representations
"""


with gr.Blocks(title="VoMP") as demo:
    gr.HTML(title_md)
    gr.Markdown(description_md)

    with gr.Column(elem_classes="main-content"):
        with gr.Row():
            with gr.Column(scale=1, elem_classes="input-column"):
                gr.Markdown("### 📤 Input")
                input_model = gr.Model3D(
                    label="Upload 3D Model",
                    clear_color=[1.0, 1.0, 1.0, 1.0],
                )

                submit_btn = gr.Button(
                    "🚀 Generate Materials", variant="primary", size="lg"
                )

            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_cloud = gr.Model3D(
                            label="Young's Modulus",
                            clear_color=[0.1, 0.1, 0.1, 1.0],
                            height=400,
                        )
                        youngs_colorbar = gr.Image(height=50, show_label=False)

                    with gr.Column(scale=1, min_width=200):
                        poissons_cloud = gr.Model3D(
                            label="Poisson's Ratio",
                            clear_color=[0.1, 0.1, 0.1, 1.0],
                            height=400,
                        )
                        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_cloud = gr.Model3D(
                            label="Density",
                            clear_color=[0.1, 0.1, 0.1, 1.0],
                            height=400,
                        )
                        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.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=[
                youngs_cloud,
                youngs_colorbar,
                poissons_cloud,
                poissons_colorbar,
                density_cloud,
                density_colorbar,
                output_file,
            ],
            fn=process_3d_model,
            cache_examples=False,
        )

    # Event handlers
    submit_btn.click(
        fn=process_3d_model,
        inputs=[input_model],
        outputs=[
            youngs_cloud,
            youngs_colorbar,
            poissons_cloud,
            poissons_colorbar,
            density_cloud,
            density_colorbar,
            output_file,
        ],
    )

if __name__ == "__main__":
    demo.launch(css=css)