From 54e5892607fbc3b843fdcb5383014cee0afcd20c Mon Sep 17 00:00:00 2001 From: Xinjie Date: Tue, 28 Oct 2025 11:54:13 +0800 Subject: [PATCH] feat(convert): Add script for infinigen assets convertion (#47) --- apps/__init__.py | 0 apps/app_style.py | 27 ++ apps/common.py | 33 +- apps/image_to_3d.py | 4 +- apps/text_to_3d.py | 4 +- apps/texture_edit.py | 4 +- apps/visualize_asset.py | 419 +++++++++++++++++++++++++ embodied_gen/data/asset_converter.py | 313 +++++++++++++++--- embodied_gen/data/convex_decomposer.py | 2 +- embodied_gen/scripts/compose_layout.py | 2 + 10 files changed, 719 insertions(+), 89 deletions(-) create mode 100644 apps/__init__.py create mode 100644 apps/app_style.py create mode 100644 apps/visualize_asset.py diff --git a/apps/__init__.py b/apps/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/apps/app_style.py b/apps/app_style.py new file mode 100644 index 0000000..cf27056 --- /dev/null +++ b/apps/app_style.py @@ -0,0 +1,27 @@ +from gradio.themes import Soft +from gradio.themes.utils.colors import gray, neutral, slate, stone, teal, zinc + +lighting_css = """ + +""" + +image_css = """ + +""" + +custom_theme = Soft( + primary_hue=stone, + secondary_hue=gray, + radius_size="md", + text_size="sm", + spacing_size="sm", +) diff --git a/apps/common.py b/apps/common.py index 9399ab3..c3891f8 100644 --- a/apps/common.py +++ b/apps/common.py @@ -30,8 +30,6 @@ import torch import torch.nn.functional as F import trimesh from easydict import EasyDict as edict -from gradio.themes import Soft -from gradio.themes.utils.colors import gray, neutral, slate, stone, teal, zinc from PIL import Image from embodied_gen.data.backproject_v2 import entrypoint as backproject_api from embodied_gen.data.differentiable_render import entrypoint as render_api @@ -151,6 +149,7 @@ if os.getenv("GRADIO_APP") == "imageto3d": TMP_DIR = os.path.join( os.path.dirname(os.path.abspath(__file__)), "sessions/imageto3d" ) + os.makedirs(TMP_DIR, exist_ok=True) elif os.getenv("GRADIO_APP") == "textto3d": RBG_REMOVER = RembgRemover() RBG14_REMOVER = BMGG14Remover() @@ -168,6 +167,7 @@ elif os.getenv("GRADIO_APP") == "textto3d": TMP_DIR = os.path.join( os.path.dirname(os.path.abspath(__file__)), "sessions/textto3d" ) + os.makedirs(TMP_DIR, exist_ok=True) elif os.getenv("GRADIO_APP") == "texture_edit": PIPELINE_IP = build_texture_gen_pipe( base_ckpt_dir="./weights", @@ -182,34 +182,7 @@ elif os.getenv("GRADIO_APP") == "texture_edit": TMP_DIR = os.path.join( os.path.dirname(os.path.abspath(__file__)), "sessions/texture_edit" ) - -os.makedirs(TMP_DIR, exist_ok=True) - - -lighting_css = """ - -""" - -image_css = """ - -""" - -custom_theme = Soft( - primary_hue=stone, - secondary_hue=gray, - radius_size="md", - text_size="sm", - spacing_size="sm", -) + os.makedirs(TMP_DIR, exist_ok=True) def start_session(req: gr.Request) -> None: diff --git a/apps/image_to_3d.py b/apps/image_to_3d.py index 752d031..3840ebb 100644 --- a/apps/image_to_3d.py +++ b/apps/image_to_3d.py @@ -21,18 +21,16 @@ os.environ["GRADIO_APP"] = "imageto3d" from glob import glob import gradio as gr +from app_style import custom_theme, image_css, lighting_css from common import ( MAX_SEED, VERSION, active_btn_by_content, - custom_theme, end_session, extract_3d_representations_v2, extract_urdf, get_seed, - image_css, image_to_3d, - lighting_css, preprocess_image_fn, preprocess_sam_image_fn, select_point, diff --git a/apps/text_to_3d.py b/apps/text_to_3d.py index 7bf8380..ab41e3d 100644 --- a/apps/text_to_3d.py +++ b/apps/text_to_3d.py @@ -21,20 +21,18 @@ os.environ["GRADIO_APP"] = "textto3d" import gradio as gr +from app_style import custom_theme, image_css, lighting_css from common import ( MAX_SEED, VERSION, active_btn_by_text_content, - custom_theme, end_session, extract_3d_representations_v2, extract_urdf, get_cached_image, get_seed, get_selected_image, - image_css, image_to_3d, - lighting_css, start_session, text2image_fn, ) diff --git a/apps/texture_edit.py b/apps/texture_edit.py index e505082..d36905e 100644 --- a/apps/texture_edit.py +++ b/apps/texture_edit.py @@ -19,18 +19,16 @@ import os os.environ["GRADIO_APP"] = "texture_edit" import gradio as gr +from app_style import custom_theme, image_css, lighting_css from common import ( MAX_SEED, VERSION, backproject_texture_v2, - custom_theme, end_session, generate_condition, generate_texture_mvimages, get_seed, get_selected_image, - image_css, - lighting_css, render_result_video, start_session, ) diff --git a/apps/visualize_asset.py b/apps/visualize_asset.py new file mode 100644 index 0000000..9712997 --- /dev/null +++ b/apps/visualize_asset.py @@ -0,0 +1,419 @@ +import os +import shutil +import xml.etree.ElementTree as ET + +import gradio as gr +import pandas as pd +from app_style import custom_theme, lighting_css + +# --- Configuration & Data Loading --- +VERSION = "v0.1.5" +RUNNING_MODE = "local" # local or hf_remote +CSV_FILE = "dataset_index.csv" + +if RUNNING_MODE == "local": + DATA_ROOT = "/horizon-bucket/robot_lab/datasets/embodiedgen/assets" +elif RUNNING_MODE == "hf_remote": + from huggingface_hub import snapshot_download + + snapshot_download( + repo_id="HorizonRobotics/EmbodiedGenData", + repo_type="dataset", + allow_patterns=f"dataset/**", + local_dir="EmbodiedGenData", + local_dir_use_symlinks=False, + ) + DATA_ROOT = "EmbodiedGenData/dataset" +else: + raise ValueError( + f"Unknown RUNNING_MODE: {RUNNING_MODE}, must be 'local' or 'hf_remote'." + ) + +csv_path = os.path.join(DATA_ROOT, CSV_FILE) +df = pd.read_csv(csv_path) +TMP_DIR = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "sessions/asset_viewer" +) +os.makedirs(TMP_DIR, exist_ok=True) + +# --- Custom CSS for Styling --- +css = """ +.gradio-container .gradio-group { box-shadow: 0 2px 4px rgba(0,0,0,0.05) !important; } +#asset-gallery { border: 1px solid #E5E7EB; border-radius: 8px; padding: 8px; background-color: #F9FAFB; } +""" + +lighting_css = """ + +""" + + +# --- Helper Functions --- +def get_primary_categories(): + return sorted(df["primary_category"].dropna().unique()) + + +def get_secondary_categories(primary): + if not primary: + return [] + return sorted( + df[df["primary_category"] == primary]["secondary_category"] + .dropna() + .unique() + ) + + +def get_categories(primary, secondary): + if not primary or not secondary: + return [] + return sorted( + df[ + (df["primary_category"] == primary) + & (df["secondary_category"] == secondary) + ]["category"] + .dropna() + .unique() + ) + + +def get_assets(primary, secondary, category): + if not primary or not secondary: + return [], gr.update(interactive=False) + + subset = df[ + (df["primary_category"] == primary) + & (df["secondary_category"] == secondary) + ] + if category: + subset = subset[subset["category"] == category] + + items = [] + for row in subset.itertuples(): + asset_dir = os.path.join(DATA_ROOT, row.asset_dir) + video_path = None + if pd.notna(asset_dir) and os.path.exists(asset_dir): + for f in os.listdir(asset_dir): + if f.lower().endswith(".mp4"): + video_path = os.path.join(asset_dir, f) + break + items.append( + video_path + if video_path + else "https://dummyimage.com/512x512/cccccc/000000&text=No+Preview" + ) + + return items, gr.update(interactive=True) + + +def show_asset_from_gallery( + evt: gr.SelectData, primary: str, secondary: str, category: str +): + index = evt.index + subset = df[ + (df["primary_category"] == primary) + & (df["secondary_category"] == secondary) + ] + if category: + subset = subset[subset["category"] == category] + + est_type_text = "N/A" + est_height_text = "N/A" + est_mass_text = "N/A" + est_mu_text = "N/A" + + if index >= len(subset): + return ( + None, + "Error: Selection index is out of bounds.", + None, + None, + est_type_text, + est_height_text, + est_mass_text, + est_mu_text, + ) + + row = subset.iloc[index] + desc = row["description"] + urdf_path = os.path.join(DATA_ROOT, row["urdf_path"]) + asset_dir = os.path.join(DATA_ROOT, row["asset_dir"]) + mesh_to_display = None + if pd.notna(urdf_path) and os.path.exists(urdf_path): + try: + tree = ET.parse(urdf_path) + root = tree.getroot() + + mesh_element = root.find('.//visual/geometry/mesh') + if mesh_element is not None: + mesh_filename = mesh_element.get('filename') + if mesh_filename: + glb_filename = os.path.splitext(mesh_filename)[0] + ".glb" + potential_path = os.path.join(asset_dir, glb_filename) + if os.path.exists(potential_path): + mesh_to_display = potential_path + + category_elem = root.find('.//extra_info/category') + if category_elem is not None and category_elem.text: + est_type_text = category_elem.text.strip() + + height_elem = root.find('.//extra_info/real_height') + if height_elem is not None and height_elem.text: + est_height_text = height_elem.text.strip() + + mass_elem = root.find('.//extra_info/min_mass') + if mass_elem is not None and mass_elem.text: + est_mass_text = mass_elem.text.strip() + + mu_elem = root.find('.//collision/gazebo/mu2') + if mu_elem is not None and mu_elem.text: + est_mu_text = mu_elem.text.strip() + + except ET.ParseError: + desc = f"Error: Failed to parse URDF at {urdf_path}. {desc}" + except Exception as e: + desc = f"An error occurred while processing URDF: {str(e)}. {desc}" + + return ( + gr.update(value=mesh_to_display), + desc, + asset_dir, + urdf_path, + est_type_text, + est_height_text, + est_mass_text, + est_mu_text, + ) + + +def create_asset_zip(asset_dir: str, req: gr.Request): + user_dir = os.path.join(TMP_DIR, str(req.session_hash)) + os.makedirs(user_dir, exist_ok=True) + + asset_folder_name = os.path.basename(os.path.normpath(asset_dir)) + zip_path_base = os.path.join(user_dir, asset_folder_name) + + archive_path = shutil.make_archive( + base_name=zip_path_base, format='zip', root_dir=asset_dir + ) + gr.Info(f"✅ {asset_folder_name}.zip is ready and can be downloaded.") + + return archive_path + + +def start_session(req: gr.Request) -> None: + user_dir = os.path.join(TMP_DIR, str(req.session_hash)) + os.makedirs(user_dir, exist_ok=True) + + +def end_session(req: gr.Request) -> None: + user_dir = os.path.join(TMP_DIR, str(req.session_hash)) + if os.path.exists(user_dir): + shutil.rmtree(user_dir) + + +# --- Gradio UI Definition --- +with gr.Blocks( + theme=custom_theme, + css=css, + title="3D Asset Library", +) as demo: + gr.HTML(lighting_css, visible=False) + gr.Markdown( + """ + ## 🏛️ ***EmbodiedGen***: 3D Asset Gallery Explorer + + **🔖 Version**: {VERSION} +

+ + 🌐 Project Page + + + 📄 arXiv + + + 💻 GitHub + + + 🎥 Video + +

+ + Browse and visualize the EmbodiedGen 3D asset database. Select categories to filter and click on a preview to load the model. + + """.format( + VERSION=VERSION + ), + elem_classes=["header"], + ) + + primary_list = get_primary_categories() + primary_val = primary_list[0] if primary_list else None + secondary_list = get_secondary_categories(primary_val) + secondary_val = secondary_list[0] if secondary_list else None + category_list = get_categories(primary_val, secondary_val) + category_val = category_list[0] if category_list else None + asset_folder = gr.State(value=None) + + with gr.Row(equal_height=False): + with gr.Column(scale=1, min_width=350): + with gr.Group(): + gr.Markdown("### Select Asset Category") + primary = gr.Dropdown( + choices=primary_list, + value=primary_val, + label="🗂️ Primary Category", + ) + secondary = gr.Dropdown( + choices=secondary_list, + value=secondary_val, + label="📂 Secondary Category", + ) + category = gr.Dropdown( + choices=category_list, + value=category_val, + label="🏷️ Asset Category", + ) + + with gr.Group(): + gallery = gr.Gallery( + value=get_assets(primary_val, secondary_val, category_val)[ + 0 + ], + label="🖼️ Asset Previews", + columns=3, + height="auto", + allow_preview=True, + elem_id="asset-gallery", + interactive=bool(category_val), + ) + + with gr.Column(scale=2, min_width=500): + with gr.Group(): + viewer = gr.Model3D( + label="🧊 3D Model Viewer", + height=500, + clear_color=[0.95, 0.95, 0.95], + elem_id="lighter_mesh", + ) + with gr.Row(): + # TODO: Add more asset details if needed + est_type_text = gr.Textbox( + label="Asset category", interactive=False + ) + est_height_text = gr.Textbox( + label="Real height(.m)", interactive=False + ) + est_mass_text = gr.Textbox( + label="Mass(.kg)", interactive=False + ) + est_mu_text = gr.Textbox( + label="Friction coefficient", interactive=False + ) + with gr.Accordion(label="Asset Details", open=False): + desc_box = gr.Textbox( + label="📝 Asset Description", interactive=False + ) + urdf_file = gr.Textbox( + label="URDF File Path", interactive=False, lines=2 + ) + with gr.Row(): + extract_btn = gr.Button( + "📥 Extract Asset", + variant="primary", + interactive=False, + ) + download_btn = gr.DownloadButton( + label="⬇️ Download Asset", + variant="primary", + interactive=False, + ) + + def update_on_primary_change(p): + s_choices = get_secondary_categories(p) + return ( + gr.update(choices=s_choices, value=None), + gr.update(choices=[], value=None), + [], + gr.update(interactive=False), + ) + + def update_on_secondary_change(p, s): + c_choices = get_categories(p, s) + return ( + gr.update(choices=c_choices, value=None), + [], + gr.update(interactive=False), + ) + + def update_on_secondary_change(p, s): + c_choices = get_categories(p, s) + asset_previews, gallery_update = get_assets(p, s, None) + return ( + gr.update(choices=c_choices, value=None), + asset_previews, + gallery_update, + ) + + primary.change( + fn=update_on_primary_change, + inputs=[primary], + outputs=[secondary, category, gallery, gallery], + ) + + secondary.change( + fn=update_on_secondary_change, + inputs=[primary, secondary], + outputs=[category, gallery, gallery], + ) + + category.change( + fn=get_assets, + inputs=[primary, secondary, category], + outputs=[gallery, gallery], + ) + + gallery.select( + fn=show_asset_from_gallery, + inputs=[primary, secondary, category], + outputs=[ + viewer, + desc_box, + asset_folder, + urdf_file, + est_type_text, + est_height_text, + est_mass_text, + est_mu_text, + ], + ).success( + lambda: tuple( + [ + gr.Button(interactive=True), + gr.Button(interactive=False), + ] + ), + outputs=[extract_btn, download_btn], + ) + + extract_btn.click( + fn=create_asset_zip, inputs=[asset_folder], outputs=[download_btn] + ).success( + fn=lambda: gr.update(interactive=True), + outputs=download_btn, + ) + + demo.load(start_session) + demo.unload(end_session) + + +if __name__ == "__main__": + demo.launch( + server_name="10.34.8.82", + server_port=8088, + allowed_paths=[ + "/horizon-bucket/robot_lab/datasets/embodiedgen/assets" + ], + ) diff --git a/embodied_gen/data/asset_converter.py b/embodied_gen/data/asset_converter.py index 2b92907..546a75b 100644 --- a/embodied_gen/data/asset_converter.py +++ b/embodied_gen/data/asset_converter.py @@ -5,7 +5,8 @@ import os import xml.etree.ElementTree as ET from abc import ABC, abstractmethod from dataclasses import dataclass -from shutil import copy +from glob import glob +from shutil import copy, copytree, rmtree import trimesh from scipy.spatial.transform import Rotation @@ -44,14 +45,23 @@ class AssetConverterBase(ABC): self, input_mesh: str, output_mesh: str, mesh_origin: ET.Element ) -> None: """Apply transform to the mesh based on the origin element in URDF.""" - mesh = trimesh.load(input_mesh) + mesh = trimesh.load(input_mesh, group_material=False) rpy = list(map(float, mesh_origin.get("rpy").split(" "))) rotation = Rotation.from_euler("xyz", rpy, degrees=False) offset = list(map(float, mesh_origin.get("xyz").split(" "))) - mesh.vertices = (mesh.vertices @ rotation.as_matrix().T) + offset - os.makedirs(os.path.dirname(output_mesh), exist_ok=True) - _ = mesh.export(output_mesh) + + if isinstance(mesh, trimesh.Scene): + combined = trimesh.Scene() + for mesh_part in mesh.geometry.values(): + mesh_part.vertices = ( + mesh_part.vertices @ rotation.as_matrix().T + ) + offset + combined.add_geometry(mesh_part) + _ = combined.export(output_mesh) + else: + mesh.vertices = (mesh.vertices @ rotation.as_matrix().T) + offset + _ = mesh.export(output_mesh) return @@ -95,31 +105,48 @@ class MeshtoMJCFConverter(AssetConverterBase): mesh = geometry.find("mesh") filename = mesh.get("filename") scale = mesh.get("scale", "1.0 1.0 1.0") - - mesh_asset = ET.SubElement( - mujoco_element, "mesh", name=mesh_name, file=filename, scale=scale - ) - geom = ET.SubElement(body, "geom", type="mesh", mesh=mesh_name) - - self._copy_asset_file( - f"{input_dir}/{filename}", - f"{output_dir}/{filename}", - ) - - # Preprocess the mesh by applying rotation. input_mesh = f"{input_dir}/{filename}" output_mesh = f"{output_dir}/{filename}" + self._copy_asset_file(input_mesh, output_mesh) + mesh_origin = element.find("origin") if mesh_origin is not None: self.transform_mesh(input_mesh, output_mesh, mesh_origin) - if material is not None: - geom.set("material", material.get("name")) - if is_collision: - geom.set("contype", "1") - geom.set("conaffinity", "1") - geom.set("rgba", "1 1 1 0") + mesh_parts = trimesh.load( + output_mesh, group_material=False, force="scene" + ) + mesh_parts = mesh_parts.geometry.values() + else: + mesh_parts = [trimesh.load(output_mesh, force="mesh")] + for idx, mesh_part in enumerate(mesh_parts): + if is_collision: + idx_mesh_name = f"{mesh_name}_{idx}" + base, ext = os.path.splitext(filename) + idx_filename = f"{base}_{idx}{ext}" + base_outdir = os.path.dirname(output_mesh) + mesh_part.export(os.path.join(base_outdir, '..', idx_filename)) + geom_attrs = { + "contype": "1", + "conaffinity": "1", + "rgba": "1 1 1 0", + } + else: + idx_mesh_name, idx_filename = mesh_name, filename + geom_attrs = {"contype": "0", "conaffinity": "0"} + + ET.SubElement( + mujoco_element, + "mesh", + name=idx_mesh_name, + file=idx_filename, + scale=scale, + ) + geom = ET.SubElement(body, "geom", type="mesh", mesh=idx_mesh_name) + geom.attrib.update(geom_attrs) + if material is not None: + geom.set("material", material.get("name")) def add_materials( self, @@ -137,26 +164,36 @@ class MeshtoMJCFConverter(AssetConverterBase): mesh = geometry.find("mesh") filename = mesh.get("filename") dirname = os.path.dirname(filename) + material = None + for path in glob(f"{input_dir}/{dirname}/*.png"): + file_name = os.path.basename(path) + if "keep_materials" in self.kwargs: + find_flag = False + for keep_key in self.kwargs["keep_materials"]: + if keep_key in file_name.lower(): + find_flag = True + if find_flag is False: + continue - material = ET.SubElement( - mujoco_element, - "material", - name=f"material_{name}", - texture=f"texture_{name}", - reflectance=str(reflectance), - ) - ET.SubElement( - mujoco_element, - "texture", - name=f"texture_{name}", - type="2d", - file=f"{dirname}/material_0.png", - ) - - self._copy_asset_file( - f"{input_dir}/{dirname}/material_0.png", - f"{output_dir}/{dirname}/material_0.png", - ) + self._copy_asset_file( + path, + f"{output_dir}/{dirname}/{file_name}", + ) + texture_name = f"texture_{name}_{os.path.splitext(file_name)[0]}" + material = ET.SubElement( + mujoco_element, + "material", + name=f"material_{name}", + texture=texture_name, + reflectance=str(reflectance), + ) + ET.SubElement( + mujoco_element, + "texture", + name=texture_name, + type="2d", + file=f"{dirname}/{file_name}", + ) return material @@ -213,6 +250,93 @@ class MeshtoMJCFConverter(AssetConverterBase): logger.info(f"Successfully converted {urdf_path} → {mjcf_path}") +class URDFtoMJCFConverter(MeshtoMJCFConverter): + """Convert URDF files with joints to MJCF format, handling transformations from joints.""" + + def convert(self, urdf_path: str, mjcf_path: str, **kwargs) -> str: + """Convert a URDF file with joints to MJCF format.""" + tree = ET.parse(urdf_path) + root = tree.getroot() + + mujoco_struct = ET.Element("mujoco") + mujoco_struct.set("model", root.get("name")) + mujoco_asset = ET.SubElement(mujoco_struct, "asset") + mujoco_worldbody = ET.SubElement(mujoco_struct, "worldbody") + + input_dir = os.path.dirname(urdf_path) + output_dir = os.path.dirname(mjcf_path) + os.makedirs(output_dir, exist_ok=True) + + body_dict = {} + for idx, link in enumerate(root.findall("link")): + link_name = link.get("name", f"unnamed_link_{idx}") + body = ET.SubElement(mujoco_worldbody, "body", name=link_name) + body_dict[link_name] = body + if link.find("visual") is not None: + material = self.add_materials( + mujoco_asset, + link, + "visual", + input_dir, + output_dir, + name=str(idx), + ) + self.add_geometry( + mujoco_asset, + link, + body, + "visual", + input_dir, + output_dir, + f"visual_mesh_{idx}", + material, + ) + if link.find("collision") is not None: + self.add_geometry( + mujoco_asset, + link, + body, + "collision", + input_dir, + output_dir, + f"collision_mesh_{idx}", + is_collision=True, + ) + + # Process joints to set transformations and hierarchy + for joint in root.findall("joint"): + joint_type = joint.get("type") + if joint_type != "fixed": + logger.warning("Only support fixed joints in conversion now.") + continue + + parent_link = joint.find("parent").get("link") + child_link = joint.find("child").get("link") + origin = joint.find("origin") + if parent_link not in body_dict or child_link not in body_dict: + logger.warning( + f"Parent or child link not found for joint: {joint.get('name')}" + ) + continue + + child_body = body_dict[child_link] + mujoco_worldbody.remove(child_body) + parent_body = body_dict[parent_link] + parent_body.append(child_body) + if origin is not None: + xyz = origin.get("xyz", "0 0 0") + rpy = origin.get("rpy", "0 0 0") + child_body.set("pos", xyz) + child_body.set("euler", rpy) + + tree = ET.ElementTree(mujoco_struct) + ET.indent(tree, space=" ", level=0) + tree.write(mjcf_path, encoding="utf-8", xml_declaration=True) + logger.info(f"Successfully converted {urdf_path} → {mjcf_path}") + + return mjcf_path + + class MeshtoUSDConverter(AssetConverterBase): """Convert Mesh file from URDF into USD format.""" @@ -289,6 +413,7 @@ class MeshtoUSDConverter(AssetConverterBase): ) urdf_converter = MeshConverter(cfg) usd_path = urdf_converter.usd_path + rmtree(os.path.dirname(output_mesh)) stage = Usd.Stage.Open(usd_path) layer = stage.GetRootLayer() @@ -335,6 +460,75 @@ class MeshtoUSDConverter(AssetConverterBase): logger.info(f"Successfully converted {urdf_path} → {usd_path}") +class PhysicsUSDAdder(MeshtoUSDConverter): + DEFAULT_BIND_APIS = [ + "MaterialBindingAPI", + "PhysicsMeshCollisionAPI", + "PhysicsCollisionAPI", + "PhysxCollisionAPI", + "PhysicsRigidBodyAPI", + ] + + def convert(self, usd_path: str, output_file: str = None): + from pxr import PhysxSchema, Sdf, Usd, UsdGeom, UsdPhysics + + if output_file is None: + output_file = usd_path + else: + dst_dir = os.path.dirname(output_file) + src_dir = os.path.dirname(usd_path) + copytree(src_dir, dst_dir, dirs_exist_ok=True) + + stage = Usd.Stage.Open(output_file) + layer = stage.GetRootLayer() + with Usd.EditContext(stage, layer): + for prim in stage.Traverse(): + if prim.IsA(UsdGeom.Xform): + for child in prim.GetChildren(): + if not child.IsA(UsdGeom.Mesh): + continue + + # Skip the lightfactory in Infinigen + if "lightfactory" in prim.GetName().lower(): + continue + + approx_attr = prim.GetAttribute( + "physics:approximation" + ) + if not approx_attr: + approx_attr = prim.CreateAttribute( + "physics:approximation", + Sdf.ValueTypeNames.Token, + ) + approx_attr.Set("convexDecomposition") + physx_conv_api = PhysxSchema.PhysxConvexDecompositionCollisionAPI.Apply( + prim + ) + physx_conv_api.GetShrinkWrapAttr().Set(True) + + rigid_body_api = UsdPhysics.RigidBodyAPI.Apply(prim) + rigid_body_api.CreateKinematicEnabledAttr().Set(True) + if prim.GetAttribute("physics:mass"): + prim.RemoveProperty("physics:mass") + if prim.GetAttribute("physics:velocity"): + prim.RemoveProperty("physics:velocity") + + api_schemas = prim.GetMetadata("apiSchemas") + if api_schemas is None: + api_schemas = Sdf.TokenListOp() + + api_list = list(api_schemas.GetAddedOrExplicitItems()) + for api in self.DEFAULT_BIND_APIS: + if api not in api_list: + api_list.append(api) + + api_schemas.appendedItems = api_list + prim.SetMetadata("apiSchemas", api_schemas) + + layer.Save() + logger.info(f"Successfully converted {usd_path} to {output_file}") + + class URDFtoUSDConverter(MeshtoUSDConverter): """Convert URDF files into USD format. @@ -440,12 +634,14 @@ class AssetConverterFactory: target_type: AssetType, source_type: AssetType = "urdf", **kwargs ) -> AssetConverterBase: """Create an asset converter instance based on target and source types.""" - if target_type == AssetType.MJCF and source_type == AssetType.URDF: + if target_type == AssetType.MJCF and source_type == AssetType.MESH: converter = MeshtoMJCFConverter(**kwargs) - elif target_type == AssetType.USD and source_type == AssetType.URDF: - converter = URDFtoUSDConverter(**kwargs) + elif target_type == AssetType.MJCF and source_type == AssetType.URDF: + converter = URDFtoMJCFConverter(**kwargs) elif target_type == AssetType.USD and source_type == AssetType.MESH: converter = MeshtoUSDConverter(**kwargs) + elif target_type == AssetType.USD and source_type == AssetType.URDF: + converter = URDFtoUSDConverter(**kwargs) else: raise ValueError( f"Unsupported converter type: {source_type} -> {target_type}." @@ -455,8 +651,8 @@ class AssetConverterFactory: if __name__ == "__main__": - # target_asset_type = AssetType.MJCF - target_asset_type = AssetType.USD + target_asset_type = AssetType.MJCF + # target_asset_type = AssetType.USD urdf_paths = [ "outputs/embodiedgen_assets/demo_assets/remote_control/result/remote_control.urdf", @@ -464,11 +660,11 @@ if __name__ == "__main__": if target_asset_type == AssetType.MJCF: output_files = [ - "outputs/embodiedgen_assets/demo_assets/remote_control/mjcf/remote_control.mjcf", + "outputs/embodiedgen_assets/demo_assets/demo_assets/remote_control/mjcf/remote_control.mjcf", ] asset_converter = AssetConverterFactory.create( target_type=AssetType.MJCF, - source_type=AssetType.URDF, + source_type=AssetType.MESH, ) elif target_asset_type == AssetType.USD: @@ -495,3 +691,22 @@ if __name__ == "__main__": # with asset_converter: # asset_converter.convert(urdf_path, output_file) + + # # Convert infinigen urdf to mjcf + # urdf_path = "/home/users/xinjie.wang/xinjie/infinigen/outputs/exports/kitchen_i_urdf/export_scene/scene.urdf" + # output_file = "/home/users/xinjie.wang/xinjie/infinigen/outputs/exports/kitchen_i_urdf/mjcf/scene.mjcf" + # asset_converter = AssetConverterFactory.create( + # target_type=AssetType.MJCF, + # source_type=AssetType.URDF, + # keep_materials=["diffuse"], + # ) + # with asset_converter: + # asset_converter.convert(urdf_path, output_file) + + # # Convert infinigen usdc to physics usdc + # converter = PhysicsUSDAdder() + # with converter: + # converter.convert( + # usd_path="/home/users/xinjie.wang/xinjie/infinigen/outputs/usdc/export_scene/export_scene.usdc", + # output_file="/home/users/xinjie.wang/xinjie/infinigen/outputs/usdc_p3/export_scene/export_scene.usdc", + # ) diff --git a/embodied_gen/data/convex_decomposer.py b/embodied_gen/data/convex_decomposer.py index f1c4e97..88e6084 100644 --- a/embodied_gen/data/convex_decomposer.py +++ b/embodied_gen/data/convex_decomposer.py @@ -27,7 +27,7 @@ logger = logging.getLogger(__name__) __all__ = [ "decompose_convex_coacd", "decompose_convex_mesh", - "decompose_convex_process", + "decompose_convex_mp", ] diff --git a/embodied_gen/scripts/compose_layout.py b/embodied_gen/scripts/compose_layout.py index 0dae279..8a9e21e 100644 --- a/embodied_gen/scripts/compose_layout.py +++ b/embodied_gen/scripts/compose_layout.py @@ -56,6 +56,8 @@ def entrypoint(**kwargs): for key in layout_info.assets: src = f"{origin_dir}/{layout_info.assets[key]}" dst = f"{output_dir}/{layout_info.assets[key]}" + if src == dst: + continue shutil.copytree(src, dst, dirs_exist_ok=True) with open(out_layout_path, "w") as f: