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

VoMP: Predicting Volumetric Mechanical Properties

Feed-forward, fine-grained, physically based volumetric material properties from Splats, Meshes, NeRFs, and more.

Paper PDF Project Page
""" 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)