diff --git a/src/multiverse_parser/exporter/urdf_exporter.py b/src/multiverse_parser/exporter/urdf_exporter.py index e8565d4..f0d1c0b 100644 --- a/src/multiverse_parser/exporter/urdf_exporter.py +++ b/src/multiverse_parser/exporter/urdf_exporter.py @@ -1,10 +1,12 @@ #!/usr/bin/env python3 - +import glob import os import shutil from typing import Tuple, Optional, Union import numpy +import trimesh +from trimesh.decomposition import convex_decomposition from scipy.spatial.transform import Rotation from urdf_parser_py import urdf @@ -18,7 +20,7 @@ from ..utils import xform_cache from pxr import UsdUrdf, Gf, UsdPhysics, UsdGeom, Usd, UsdShade - +import lxml.etree as ET def get_robot_name(world_builder: WorldBuilder) -> str: usd_urdf = UsdUrdf.Urdf.Get(world_builder.stage, "/urdf") @@ -560,6 +562,141 @@ def export(self, keep_usd: bool = True): self.factory.save_tmp_model(usd_file_path=self.file_path.replace(".urdf", ".usda"), excludes=["usd", ".usda"]) + + def export_with_convex_collisions(self, keep_usd: bool = True): + """ + Export the URDF with convex decomposed collision meshes. + + :param keep_usd: If True, keeps the original USD file after decomposition. + """ + # 1. Export the original URDF + self.export(keep_usd=keep_usd) + original_urdf_path = self.file_path + decomposed_urdf_path = original_urdf_path.replace(".urdf", "_decomposed.urdf") + + # 2. Parse original URDF + parser = ET.XMLParser(remove_blank_text=True) + tree = ET.parse(original_urdf_path, parser) + root = tree.getroot() + + # 3. Find all collision mesh references and decompose + for mesh in root.findall(".//collision/geometry/mesh"): + mesh_path = mesh.attrib.get("filename") + if not mesh_path or not mesh_path.endswith(".stl"): + continue # skip non-STL + + if mesh_path.startswith("package://"): + print(f"[SKIP] Cannot resolve package URI: {mesh_path}") + continue + + # Remove "file://" prefix if present + mesh_path_clean = mesh_path.replace("file://", "") + full_input_path = os.path.abspath(os.path.join(os.path.dirname(self.file_path), mesh_path_clean)) + # Build path to stl file + mesh_root = os.path.dirname(os.path.dirname(full_input_path)) + obj_dir = os.path.join(mesh_root, "obj") + os.makedirs(obj_dir, exist_ok=True) + + # Prepare input and output paths + vhacd_input_basename = os.path.basename(full_input_path).replace(".stl", "_convex.obj") + tmp_obj_path = os.path.join("/tmp", vhacd_input_basename) + vhacd_output_basename = vhacd_input_basename.replace("_convex.obj", "_convex") + vhacd_output_dir = os.path.join(obj_dir, vhacd_output_basename) + + # Run VHACD if not already done + if not os.path.exists(vhacd_output_dir): + # convert stl to obj + mesh = trimesh.load(full_input_path) + mesh.export(tmp_obj_path) + + self._run_vhacd_trimesh(tmp_obj_path, obj_dir) + else: + print(f"[SKIP] Already decomposed: {vhacd_output_dir}") + + # Add multiple tags for each decomposed mesh + decomposed_parts = sorted(glob.glob(os.path.join(vhacd_output_dir, "*.obj"))) + if not decomposed_parts: + raise FileNotFoundError(f"No decomposed mesh parts found in {vhacd_output_dir}") + + collision_elem = mesh.getparent().getparent() # mesh → geometry → collision + link_elem = collision_elem.getparent() + link_elem.remove(collision_elem) + + # Add new tags to the + for part_path in decomposed_parts: + part_relpath = os.path.relpath(part_path, os.path.dirname(original_urdf_path)) + + collision = ET.Element("collision") + geometry = ET.SubElement(collision, "geometry") + mesh_elem = ET.SubElement(geometry, "mesh") + mesh_elem.attrib["filename"] = part_relpath + + link_elem.append(collision) + + # 4. Save the decomposed URDF + tree.write(decomposed_urdf_path, pretty_print=True, xml_declaration=True, encoding="UTF-8") + print(f"[✔] Decomposed URDF saved: {decomposed_urdf_path}") + + def _run_vhacd_trimesh(self, input_stl: str, obj_dir: str, max_convex_hulls: int = 16): + """ + Run VHACD convex decomposition on the input STL file and save the results as OBJ files. + + :param input_stl: Path to the input STL file. + :param obj_dir: Directory to save the decomposed OBJ files. + :param max_convex_hulls: Maximum number of convex hulls to generate. + """ + print(f"[VHACD] Running convex decomposition on {input_stl} (max hulls: {max_convex_hulls})") + + mesh = trimesh.load_mesh(input_stl) + + if not isinstance(mesh, trimesh.Trimesh): + raise ValueError(f"Expected Trimesh, got {type(mesh)}") + + # Skip if mesh is already convex + if mesh.is_convex: + print("[INFO] Mesh is already convex, skipping decomposition.") + base_name = os.path.basename(input_stl).replace(".obj", "") + output_base_path = os.path.join(obj_dir, base_name) + os.makedirs(output_base_path, exist_ok=True) + + part_filename = f"{base_name}_convex.obj" + part_path = os.path.join(output_base_path, part_filename) + mesh.export(part_path) + print(f"[✔] Exported original convex mesh: {part_path}") + return + + # Otherwise, run convex decomposition + mesh_args_list = convex_decomposition( + mesh, + maxConvexHulls=max_convex_hulls, + resolution=10_000, + minimumVolumePercentErrorAllowed=1.0, + maxRecursionDepth=10, + shrinkWrap=True, + fillMode="flood", + maxNumVerticesPerCH=64, + asyncACD=True, + minEdgeLength=2, + findBestPlane=False + ) + + if not mesh_args_list: + raise RuntimeError("VHACD decomposition returned no convex parts") + + base_name = os.path.basename(input_stl).replace(".obj", "") + output_base_path = os.path.join(obj_dir, base_name) + os.makedirs(output_base_path, exist_ok=True) + + for i, mesh_args in enumerate(mesh_args_list): + part = trimesh.Trimesh(**mesh_args) + part_filename = f"{base_name}_{i}.obj" + part_path = os.path.join(output_base_path, part_filename) + part.export(part_path) + print(f"[✔] Exported convex part {i}: {part_path}") + + print(f"[✔] VHACD complete: {len(mesh_args_list)} parts saved to {output_base_path}") + + @property def file_path(self) -> str: return self._file_path