feat(convert): Add script for infinigen assets convertion (#47)
This commit is contained in:
Xinjie
GitHub
parent
0fb691bd22
commit
54e5892607
0
apps/__init__.py
Normal file
0
apps/__init__.py
Normal file
27
apps/app_style.py
Normal file
27
apps/app_style.py
Normal file
@ -0,0 +1,27 @@
|
||||
from gradio.themes import Soft
|
||||
from gradio.themes.utils.colors import gray, neutral, slate, stone, teal, zinc
|
||||
|
||||
lighting_css = """
|
||||
<style>
|
||||
#lighter_mesh canvas {
|
||||
filter: brightness(1.9) !important;
|
||||
}
|
||||
</style>
|
||||
"""
|
||||
|
||||
image_css = """
|
||||
<style>
|
||||
.image_fit .image-frame {
|
||||
object-fit: contain !important;
|
||||
height: 100% !important;
|
||||
}
|
||||
</style>
|
||||
"""
|
||||
|
||||
custom_theme = Soft(
|
||||
primary_hue=stone,
|
||||
secondary_hue=gray,
|
||||
radius_size="md",
|
||||
text_size="sm",
|
||||
spacing_size="sm",
|
||||
)
|
||||
@ -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 = """
|
||||
<style>
|
||||
#lighter_mesh canvas {
|
||||
filter: brightness(1.9) !important;
|
||||
}
|
||||
</style>
|
||||
"""
|
||||
|
||||
image_css = """
|
||||
<style>
|
||||
.image_fit .image-frame {
|
||||
object-fit: contain !important;
|
||||
height: 100% !important;
|
||||
}
|
||||
</style>
|
||||
"""
|
||||
|
||||
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:
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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,
|
||||
)
|
||||
|
||||
@ -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,
|
||||
)
|
||||
|
||||
419
apps/visualize_asset.py
Normal file
419
apps/visualize_asset.py
Normal file
@ -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 = """
|
||||
<style>
|
||||
#lighter_mesh canvas {
|
||||
filter: brightness(2.2) !important;
|
||||
}
|
||||
</style>
|
||||
"""
|
||||
|
||||
|
||||
# --- 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}
|
||||
<p style="display: flex; gap: 10px; flex-wrap: nowrap;">
|
||||
<a href="https://horizonrobotics.github.io/robot_lab/embodied_gen/index.html">
|
||||
<img alt="🌐 Project Page" src="https://img.shields.io/badge/🌐-Project_Page-blue">
|
||||
</a>
|
||||
<a href="https://arxiv.org/abs/2506.10600">
|
||||
<img alt="📄 arXiv" src="https://img.shields.io/badge/📄-arXiv-b31b1b">
|
||||
</a>
|
||||
<a href="https://github.com/HorizonRobotics/EmbodiedGen">
|
||||
<img alt="💻 GitHub" src="https://img.shields.io/badge/GitHub-000000?logo=github">
|
||||
</a>
|
||||
<a href="https://www.youtube.com/watch?v=rG4odybuJRk">
|
||||
<img alt="🎥 Video" src="https://img.shields.io/badge/🎥-Video-red">
|
||||
</a>
|
||||
</p>
|
||||
|
||||
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"
|
||||
],
|
||||
)
|
||||
@ -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",
|
||||
# )
|
||||
|
||||
@ -27,7 +27,7 @@ logger = logging.getLogger(__name__)
|
||||
__all__ = [
|
||||
"decompose_convex_coacd",
|
||||
"decompose_convex_mesh",
|
||||
"decompose_convex_process",
|
||||
"decompose_convex_mp",
|
||||
]
|
||||
|
||||
|
||||
|
||||
@ -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:
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user