Data API

embodied_gen.data.asset_converter

AssetConverterBase

Bases: ABC

Abstract base class for asset converters.

Provides context management and mesh transformation utilities.

__enter__

__enter__()

Context manager entry.

Source code in embodied_gen/data/asset_converter.py

def __enter__(self):
    """Context manager entry."""
    return self

__exit__

__exit__(exc_type, exc_val, exc_tb)

Context manager exit.

Source code in embodied_gen/data/asset_converter.py

def __exit__(self, exc_type, exc_val, exc_tb):
    """Context manager exit."""
    return False

convert abstractmethod

convert(urdf_path: str, output_path: str, **kwargs) -> str

Convert an asset file.

Parameters:

Name	Type	Description	Default
urdf_path	str	Path to input URDF file.	required
output_path	str	Path to output file.	required
**kwargs		Additional arguments.	{}

Returns:

Name	Type	Description
str	str	Path to converted asset.

Source code in embodied_gen/data/asset_converter.py

@abstractmethod
def convert(self, urdf_path: str, output_path: str, **kwargs) -> str:
    """Convert an asset file.

    Args:
        urdf_path (str): Path to input URDF file.
        output_path (str): Path to output file.
        **kwargs: Additional arguments.

    Returns:
        str: Path to converted asset.
    """
    pass

transform_mesh

transform_mesh(input_mesh: str, output_mesh: str, mesh_origin: Element) -> None

Apply transform to mesh based on URDF origin element.

Parameters:

Name	Type	Description	Default
input_mesh	str	Path to input mesh.	required
output_mesh	str	Path to output mesh.	required
mesh_origin	Element	Origin element from URDF.	required

Source code in embodied_gen/data/asset_converter.py

def transform_mesh(
    self, input_mesh: str, output_mesh: str, mesh_origin: ET.Element
) -> None:
    """Apply transform to mesh based on URDF origin element.

    Args:
        input_mesh (str): Path to input mesh.
        output_mesh (str): Path to output mesh.
        mesh_origin (ET.Element): Origin element from URDF.
    """
    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(" ")))
    os.makedirs(os.path.dirname(output_mesh), exist_ok=True)

    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

AssetConverterFactory

Factory for creating asset converters based on target and source types.

Example

from embodied_gen.data.asset_converter import AssetConverterFactory
from embodied_gen.utils.enum import AssetType

converter = AssetConverterFactory.create(
    target_type=AssetType.USD, source_type=AssetType.MESH
)
with converter:
    for urdf_path, output_file in zip(urdf_paths, output_files):
        converter.convert(urdf_path, output_file)

create staticmethod

create(target_type: AssetType, source_type: AssetType = 'urdf', **kwargs) -> AssetConverterBase

Creates an asset converter instance.

Parameters:

Name	Type	Description	Default
target_type	AssetType	Target asset type.	required
source_type	AssetType	Source asset type.	'urdf'
**kwargs		Additional arguments.	{}

Returns:

Name	Type	Description
AssetConverterBase	AssetConverterBase	Converter instance.

Source code in embodied_gen/data/asset_converter.py

@staticmethod
def create(
    target_type: AssetType, source_type: AssetType = "urdf", **kwargs
) -> AssetConverterBase:
    """Creates an asset converter instance.

    Args:
        target_type (AssetType): Target asset type.
        source_type (AssetType, optional): Source asset type.
        **kwargs: Additional arguments.

    Returns:
        AssetConverterBase: Converter instance.
    """
    if target_type == AssetType.MJCF and source_type == AssetType.MESH:
        converter = MeshtoMJCFConverter(**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}."
        )

    return converter

MeshtoMJCFConverter

MeshtoMJCFConverter(**kwargs)

Bases: AssetConverterBase

Converts mesh-based URDF files to MJCF format.

Handles geometry, materials, and asset copying.

Source code in embodied_gen/data/asset_converter.py

def __init__(
    self,
    **kwargs,
) -> None:
    self.kwargs = kwargs

add_geometry

add_geometry(mujoco_element: Element, link: Element, body: Element, tag: str, input_dir: str, output_dir: str, mesh_name: str, material: Element | None = None, is_collision: bool = False) -> None

Adds geometry to MJCF body from URDF link.

Parameters:

Name	Type	Description	Default
mujoco_element	Element	MJCF asset element.	required
link	Element	URDF link element.	required
body	Element	MJCF body element.	required
tag	str	Tag name ("visual" or "collision").	required
input_dir	str	Input directory.	required
output_dir	str	Output directory.	required
mesh_name	str	Mesh name.	required
material	Element	Material element.	None
is_collision	bool	If True, treat as collision geometry.	False

Source code in embodied_gen/data/asset_converter.py

def add_geometry(
    self,
    mujoco_element: ET.Element,
    link: ET.Element,
    body: ET.Element,
    tag: str,
    input_dir: str,
    output_dir: str,
    mesh_name: str,
    material: ET.Element | None = None,
    is_collision: bool = False,
) -> None:
    """Adds geometry to MJCF body from URDF link.

    Args:
        mujoco_element (ET.Element): MJCF asset element.
        link (ET.Element): URDF link element.
        body (ET.Element): MJCF body element.
        tag (str): Tag name ("visual" or "collision").
        input_dir (str): Input directory.
        output_dir (str): Output directory.
        mesh_name (str): Mesh name.
        material (ET.Element, optional): Material element.
        is_collision (bool, optional): If True, treat as collision geometry.
    """
    element = link.find(tag)
    geometry = element.find("geometry")
    mesh = geometry.find("mesh")
    filename = mesh.get("filename")
    scale = mesh.get("scale", "1.0 1.0 1.0")
    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 is_collision:
        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"))

add_materials

add_materials(mujoco_element: Element, link: Element, tag: str, input_dir: str, output_dir: str, name: str, reflectance: float = 0.2) -> ET.Element

Adds materials to MJCF asset from URDF link.

Parameters:

Name	Type	Description	Default
mujoco_element	Element	MJCF asset element.	required
link	Element	URDF link element.	required
tag	str	Tag name.	required
input_dir	str	Input directory.	required
output_dir	str	Output directory.	required
name	str	Material name.	required
reflectance	float	Reflectance value.	0.2

Returns:

Type	Description
Element	ET.Element: Material element.

Source code in embodied_gen/data/asset_converter.py

def add_materials(
    self,
    mujoco_element: ET.Element,
    link: ET.Element,
    tag: str,
    input_dir: str,
    output_dir: str,
    name: str,
    reflectance: float = 0.2,
) -> ET.Element:
    """Adds materials to MJCF asset from URDF link.

    Args:
        mujoco_element (ET.Element): MJCF asset element.
        link (ET.Element): URDF link element.
        tag (str): Tag name.
        input_dir (str): Input directory.
        output_dir (str): Output directory.
        name (str): Material name.
        reflectance (float, optional): Reflectance value.

    Returns:
        ET.Element: Material element.
    """
    element = link.find(tag)
    geometry = element.find("geometry")
    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

        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

convert

convert(urdf_path: str, mjcf_path: str)

Converts a URDF file to MJCF format.

Parameters:

Name	Type	Description	Default
urdf_path	str	Path to URDF file.	required
mjcf_path	str	Path to output MJCF file.	required

Source code in embodied_gen/data/asset_converter.py

def convert(self, urdf_path: str, mjcf_path: str):
    """Converts a URDF file to MJCF format.

    Args:
        urdf_path (str): Path to URDF file.
        mjcf_path (str): Path to output MJCF file.
    """
    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)
    for idx, link in enumerate(root.findall("link")):
        link_name = link.get("name", "unnamed_link")
        body = ET.SubElement(mujoco_worldbody, "body", name=link_name)

        material = self.add_materials(
            mujoco_asset,
            link,
            "visual",
            input_dir,
            output_dir,
            name=str(idx),
        )
        joint = ET.SubElement(body, "joint", attrib={"type": "free"})
        self.add_geometry(
            mujoco_asset,
            link,
            body,
            "visual",
            input_dir,
            output_dir,
            f"visual_mesh_{idx}",
            material,
        )
        self.add_geometry(
            mujoco_asset,
            link,
            body,
            "collision",
            input_dir,
            output_dir,
            f"collision_mesh_{idx}",
            is_collision=True,
        )

    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}")

MeshtoUSDConverter

MeshtoUSDConverter(force_usd_conversion: bool = True, make_instanceable: bool = False, simulation_app=None, **kwargs)

Bases: AssetConverterBase

Converts mesh-based URDF files to USD format.

Adds physics APIs and post-processes collision meshes.

Initializes the converter.

Parameters:

Name	Type	Description	Default
force_usd_conversion	bool	Force USD conversion.	True
make_instanceable	bool	Make prims instanceable.	False
simulation_app	optional	Simulation app instance.	None
**kwargs		Additional arguments.	{}

Source code in embodied_gen/data/asset_converter.py

def __init__(
    self,
    force_usd_conversion: bool = True,
    make_instanceable: bool = False,
    simulation_app=None,
    **kwargs,
):
    """Initializes the converter.

    Args:
        force_usd_conversion (bool, optional): Force USD conversion.
        make_instanceable (bool, optional): Make prims instanceable.
        simulation_app (optional): Simulation app instance.
        **kwargs: Additional arguments.
    """
    if simulation_app is not None:
        self.simulation_app = simulation_app

    if "exit_close" in kwargs:
        self.exit_close = kwargs.pop("exit_close")
    else:
        self.exit_close = True

    self.usd_parms = dict(
        force_usd_conversion=force_usd_conversion,
        make_instanceable=make_instanceable,
        **kwargs,
    )

__enter__

__enter__()

Context manager entry, launches simulation app if needed.

Source code in embodied_gen/data/asset_converter.py

def __enter__(self):
    """Context manager entry, launches simulation app if needed."""
    from isaaclab.app import AppLauncher

    if not hasattr(self, "simulation_app"):
        if "launch_args" not in self.usd_parms:
            launch_args = dict(
                headless=True,
                no_splash=True,
                fast_shutdown=True,
                disable_gpu=True,
            )
        else:
            launch_args = self.usd_parms.pop("launch_args")
        self.app_launcher = AppLauncher(launch_args)
        self.simulation_app = self.app_launcher.app

    return self

__exit__

__exit__(exc_type, exc_val, exc_tb)

Context manager exit, closes simulation app if created.

Source code in embodied_gen/data/asset_converter.py

def __exit__(self, exc_type, exc_val, exc_tb):
    """Context manager exit, closes simulation app if created."""
    # Close the simulation app if it was created here
    if hasattr(self, "app_launcher") and self.exit_close:
        self.simulation_app.close()

    if exc_val is not None:
        logger.error(f"Exception occurred: {exc_val}.")

    return False

convert

convert(urdf_path: str, output_file: str)

Converts a URDF file to USD and post-processes collision meshes.

Parameters:

Name	Type	Description	Default
urdf_path	str	Path to URDF file.	required
output_file	str	Path to output USD file.	required

Source code in embodied_gen/data/asset_converter.py

def convert(self, urdf_path: str, output_file: str):
    """Converts a URDF file to USD and post-processes collision meshes.

    Args:
        urdf_path (str): Path to URDF file.
        output_file (str): Path to output USD file.
    """
    from isaaclab.sim.converters import MeshConverter, MeshConverterCfg
    from pxr import PhysxSchema, Sdf, Usd, UsdShade

    tree = ET.parse(urdf_path)
    root = tree.getroot()
    mesh_file = root.find("link/visual/geometry/mesh").get("filename")
    input_mesh = os.path.join(os.path.dirname(urdf_path), mesh_file)
    output_dir = os.path.abspath(os.path.dirname(output_file))
    output_mesh = f"{output_dir}/mesh/{os.path.basename(mesh_file)}"
    mesh_origin = root.find("link/visual/origin")
    if mesh_origin is not None:
        self.transform_mesh(input_mesh, output_mesh, mesh_origin)

    cfg = MeshConverterCfg(
        asset_path=output_mesh,
        usd_dir=output_dir,
        usd_file_name=os.path.basename(output_file),
        **self.usd_parms,
    )
    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()
    with Usd.EditContext(stage, layer):
        for prim in stage.Traverse():
            # Change texture path to relative path.
            if prim.GetName() == "material_0":
                shader = UsdShade.Shader(prim).GetInput("diffuse_texture")
                if shader.Get() is not None:
                    relative_path = shader.Get().path.replace(
                        f"{output_dir}/", ""
                    )
                    shader.Set(Sdf.AssetPath(relative_path))

            # Add convex decomposition collision and set ShrinkWrap.
            elif prim.GetName() == "mesh":
                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)

                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 {urdf_path} → {usd_path}")

PhysicsUSDAdder

PhysicsUSDAdder(force_usd_conversion: bool = True, make_instanceable: bool = False, simulation_app=None, **kwargs)

Bases: MeshtoUSDConverter

Adds physics APIs and collision properties to USD assets.

Useful for post-processing USD files for simulation.

Source code in embodied_gen/data/asset_converter.py

def __init__(
    self,
    force_usd_conversion: bool = True,
    make_instanceable: bool = False,
    simulation_app=None,
    **kwargs,
):
    """Initializes the converter.

    Args:
        force_usd_conversion (bool, optional): Force USD conversion.
        make_instanceable (bool, optional): Make prims instanceable.
        simulation_app (optional): Simulation app instance.
        **kwargs: Additional arguments.
    """
    if simulation_app is not None:
        self.simulation_app = simulation_app

    if "exit_close" in kwargs:
        self.exit_close = kwargs.pop("exit_close")
    else:
        self.exit_close = True

    self.usd_parms = dict(
        force_usd_conversion=force_usd_conversion,
        make_instanceable=make_instanceable,
        **kwargs,
    )

convert

convert(usd_path: str, output_file: str = None)

Adds physics APIs and collision properties to a USD file.

Parameters:

Name	Type	Description	Default
usd_path	str	Path to input USD file.	required
output_file	str	Path to output USD file.	None

Source code in embodied_gen/data/asset_converter.py

def convert(self, usd_path: str, output_file: str = None):
    """Adds physics APIs and collision properties to a USD file.

    Args:
        usd_path (str): Path to input USD file.
        output_file (str, optional): Path to output USD file.
    """
    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}")

URDFtoMJCFConverter

URDFtoMJCFConverter(**kwargs)

Bases: MeshtoMJCFConverter

Converts URDF files with joints to MJCF format, handling joint transformations.

Handles fixed joints and hierarchical body structure.

Source code in embodied_gen/data/asset_converter.py

def __init__(
    self,
    **kwargs,
) -> None:
    self.kwargs = kwargs

convert

convert(urdf_path: str, mjcf_path: str, **kwargs) -> str

Converts a URDF file with joints to MJCF format.

Parameters:

Name	Type	Description	Default
urdf_path	str	Path to URDF file.	required
mjcf_path	str	Path to output MJCF file.	required
**kwargs		Additional arguments.	{}

Returns:

Name	Type	Description
str	str	Path to converted MJCF file.

Source code in embodied_gen/data/asset_converter.py

def convert(self, urdf_path: str, mjcf_path: str, **kwargs) -> str:
    """Converts a URDF file with joints to MJCF format.

    Args:
        urdf_path (str): Path to URDF file.
        mjcf_path (str): Path to output MJCF file.
        **kwargs: Additional arguments.

    Returns:
        str: Path to converted MJCF file.
    """
    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

URDFtoUSDConverter

URDFtoUSDConverter(fix_base: bool = False, merge_fixed_joints: bool = False, make_instanceable: bool = True, force_usd_conversion: bool = True, collision_from_visuals: bool = True, joint_drive=None, rotate_wxyz: tuple[float] | None = None, simulation_app=None, **kwargs)

Bases: MeshtoUSDConverter

Converts URDF files to USD format.

Parameters:

Name	Type	Description	Default
fix_base	bool	Fix the base link.	False
merge_fixed_joints	bool	Merge fixed joints.	False
make_instanceable	bool	Make prims instanceable.	True
force_usd_conversion	bool	Force conversion to USD.	True
collision_from_visuals	bool	Generate collisions from visuals.	True
joint_drive	optional	Joint drive configuration.	None
rotate_wxyz	tuple[float]	Quaternion for rotation.	None
simulation_app	optional	Simulation app instance.	None
**kwargs		Additional arguments.	{}

Initializes the converter.

Parameters:

Name	Type	Description	Default
fix_base	bool	Fix the base link.	False
merge_fixed_joints	bool	Merge fixed joints.	False
make_instanceable	bool	Make prims instanceable.	True
force_usd_conversion	bool	Force conversion to USD.	True
collision_from_visuals	bool	Generate collisions from visuals.	True
joint_drive	optional	Joint drive configuration.	None
rotate_wxyz	tuple[float]	Quaternion for rotation.	None
simulation_app	optional	Simulation app instance.	None
**kwargs		Additional arguments.	{}

Source code in embodied_gen/data/asset_converter.py

def __init__(
    self,
    fix_base: bool = False,
    merge_fixed_joints: bool = False,
    make_instanceable: bool = True,
    force_usd_conversion: bool = True,
    collision_from_visuals: bool = True,
    joint_drive=None,
    rotate_wxyz: tuple[float] | None = None,
    simulation_app=None,
    **kwargs,
):
    """Initializes the converter.

    Args:
        fix_base (bool, optional): Fix the base link.
        merge_fixed_joints (bool, optional): Merge fixed joints.
        make_instanceable (bool, optional): Make prims instanceable.
        force_usd_conversion (bool, optional): Force conversion to USD.
        collision_from_visuals (bool, optional): Generate collisions from visuals.
        joint_drive (optional): Joint drive configuration.
        rotate_wxyz (tuple[float], optional): Quaternion for rotation.
        simulation_app (optional): Simulation app instance.
        **kwargs: Additional arguments.
    """
    self.usd_parms = dict(
        fix_base=fix_base,
        merge_fixed_joints=merge_fixed_joints,
        make_instanceable=make_instanceable,
        force_usd_conversion=force_usd_conversion,
        collision_from_visuals=collision_from_visuals,
        joint_drive=joint_drive,
        **kwargs,
    )
    self.rotate_wxyz = rotate_wxyz
    if simulation_app is not None:
        self.simulation_app = simulation_app

convert

convert(urdf_path: str, output_file: str)

Converts a URDF file to USD and post-processes collision meshes.

Parameters:

Name	Type	Description	Default
urdf_path	str	Path to URDF file.	required
output_file	str	Path to output USD file.	required

Source code in embodied_gen/data/asset_converter.py

def convert(self, urdf_path: str, output_file: str):
    """Converts a URDF file to USD and post-processes collision meshes.

    Args:
        urdf_path (str): Path to URDF file.
        output_file (str): Path to output USD file.
    """
    from isaaclab.sim.converters import UrdfConverter, UrdfConverterCfg
    from pxr import Gf, PhysxSchema, Sdf, Usd, UsdGeom

    cfg = UrdfConverterCfg(
        asset_path=urdf_path,
        usd_dir=os.path.abspath(os.path.dirname(output_file)),
        usd_file_name=os.path.basename(output_file),
        **self.usd_parms,
    )

    urdf_converter = UrdfConverter(cfg)
    usd_path = urdf_converter.usd_path

    stage = Usd.Stage.Open(usd_path)
    layer = stage.GetRootLayer()
    with Usd.EditContext(stage, layer):
        for prim in stage.Traverse():
            if prim.GetName() == "collisions":
                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)

                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)

    if self.rotate_wxyz is not None:
        inner_prim = next(
            p
            for p in stage.GetDefaultPrim().GetChildren()
            if p.IsA(UsdGeom.Xform)
        )
        xformable = UsdGeom.Xformable(inner_prim)
        xformable.ClearXformOpOrder()
        orient_op = xformable.AddOrientOp(UsdGeom.XformOp.PrecisionDouble)
        orient_op.Set(Gf.Quatd(*self.rotate_wxyz))

    layer.Save()
    logger.info(f"Successfully converted {urdf_path} → {usd_path}")

cvt_embodiedgen_asset_to_anysim

cvt_embodiedgen_asset_to_anysim(urdf_files: list[str], target_dirs: list[str], target_type: AssetType, source_type: AssetType, overwrite: bool = False, **kwargs) -> dict[str, str]

Convert URDF files generated by EmbodiedGen into formats required by simulators.

Supported simulators include SAPIEN, Isaac Sim, MuJoCo, Isaac Gym, Genesis, and Pybullet. Converting to the USD format requires isaacsim to be installed.

Example

from embodied_gen.data.asset_converter import cvt_embodiedgen_asset_to_anysim
from embodied_gen.utils.enum import AssetType

dst_asset_path = cvt_embodiedgen_asset_to_anysim(
    urdf_files=[
        "path1_to_embodiedgen_asset/asset.urdf",
        "path2_to_embodiedgen_asset/asset.urdf",
    ],
    target_dirs=[
        "path1_to_target_dir/asset.usd",
        "path2_to_target_dir/asset.usd",
    ],
    target_type=AssetType.USD,
    source_type=AssetType.MESH,
)

Parameters:

Name	Type	Description	Default
urdf_files	list[str]	List of URDF file paths.	required
target_dirs	list[str]	List of target directories.	required
target_type	AssetType	Target asset type.	required
source_type	AssetType	Source asset type.	required
overwrite	bool	Overwrite existing files.	False
**kwargs		Additional converter arguments.	{}

Returns:

Type	Description
dict[str, str]	dict[str, str]: Mapping from URDF file to converted asset file.

Source code in embodied_gen/data/asset_converter.py

def cvt_embodiedgen_asset_to_anysim(
    urdf_files: list[str],
    target_dirs: list[str],
    target_type: AssetType,
    source_type: AssetType,
    overwrite: bool = False,
    **kwargs,
) -> dict[str, str]:
    """Convert URDF files generated by EmbodiedGen into formats required by simulators.

    Supported simulators include SAPIEN, Isaac Sim, MuJoCo, Isaac Gym, Genesis, and Pybullet.
    Converting to the `USD` format requires `isaacsim` to be installed.

    Example:
        ```py
        from embodied_gen.data.asset_converter import cvt_embodiedgen_asset_to_anysim
        from embodied_gen.utils.enum import AssetType

        dst_asset_path = cvt_embodiedgen_asset_to_anysim(
            urdf_files=[
                "path1_to_embodiedgen_asset/asset.urdf",
                "path2_to_embodiedgen_asset/asset.urdf",
            ],
            target_dirs=[
                "path1_to_target_dir/asset.usd",
                "path2_to_target_dir/asset.usd",
            ],
            target_type=AssetType.USD,
            source_type=AssetType.MESH,
        )
        ```

    Args:
        urdf_files (list[str]): List of URDF file paths.
        target_dirs (list[str]): List of target directories.
        target_type (AssetType): Target asset type.
        source_type (AssetType): Source asset type.
        overwrite (bool, optional): Overwrite existing files.
        **kwargs: Additional converter arguments.

    Returns:
        dict[str, str]: Mapping from URDF file to converted asset file.
    """

    if isinstance(urdf_files, str):
        urdf_files = [urdf_files]
    if isinstance(target_dirs, str):
        urdf_files = [target_dirs]

    # If the target type is URDF, no conversion is needed.
    if target_type == AssetType.URDF:
        return {key: key for key in urdf_files}

    asset_converter = AssetConverterFactory.create(
        target_type=target_type,
        source_type=source_type,
        **kwargs,
    )
    asset_paths = dict()

    with asset_converter:
        for urdf_file, target_dir in zip(urdf_files, target_dirs):
            filename = os.path.basename(urdf_file).replace(".urdf", "")
            if target_type == AssetType.MJCF:
                target_file = f"{target_dir}/{filename}.xml"
            elif target_type == AssetType.USD:
                target_file = f"{target_dir}/{filename}.usd"
            else:
                raise NotImplementedError(
                    f"Target type {target_type} not supported."
                )
            if not os.path.exists(target_file) or overwrite:
                asset_converter.convert(urdf_file, target_file)

            asset_paths[urdf_file] = target_file

    return asset_paths

embodied_gen.data.datasets

PanoGSplatDataset

PanoGSplatDataset(data_dir: str, split: str = Literal['train', 'eval'], data_name: str = 'gs_data.pt', max_sample_num: int = None)

Bases: Dataset

A PyTorch Dataset for loading panorama-based 3D Gaussian Splatting data.

This dataset is designed to be compatible with train and eval pipelines that use COLMAP-style camera conventions.

Parameters:

Name	Type	Description	Default
data_dir	str	Root directory where the dataset file is located.	required
split	str	Dataset split to use, either "train" or "eval".	Literal['train', 'eval']
data_name	str	Name of the dataset file (default: "gs_data.pt").	'gs_data.pt'
max_sample_num	int	Maximum number of samples to load. If None, all available samples in the split will be used.	None

Source code in embodied_gen/data/datasets.py

def __init__(
    self,
    data_dir: str,
    split: str = Literal["train", "eval"],
    data_name: str = "gs_data.pt",
    max_sample_num: int = None,
) -> None:
    self.data_path = os.path.join(data_dir, data_name)
    self.split = split
    self.max_sample_num = max_sample_num
    if not os.path.exists(self.data_path):
        raise FileNotFoundError(
            f"Dataset file {self.data_path} not found. Please provide the correct path."
        )
    self.data = torch.load(self.data_path, weights_only=False)
    self.frames = self.data[split]
    if max_sample_num is not None:
        self.frames = self.frames[:max_sample_num]
    self.points = self.data.get("points", None)
    self.points_rgb = self.data.get("points_rgb", None)

embodied_gen.data.differentiable_render

ImageRender

ImageRender(render_items: list[RenderItems], camera_params: CameraSetting, recompute_vtx_normal: bool = True, with_mtl: bool = False, gen_color_gif: bool = False, gen_color_mp4: bool = False, gen_viewnormal_mp4: bool = False, gen_glonormal_mp4: bool = False, no_index_file: bool = False, light_factor: float = 1.0)

Bases: object

Differentiable mesh renderer supporting multi-view rendering.

This class wraps differentiable rasterization using nvdiffrast to render mesh geometry to various maps (normal, depth, alpha, albedo, etc.) and supports saving images and videos.

Parameters:

Name	Type	Description	Default
render_items	list[RenderItems]	List of rendering targets.	required
camera_params	CameraSetting	Camera parameters for rendering.	required
recompute_vtx_normal	bool	Recompute vertex normals. Defaults to True.	True
with_mtl	bool	Load mesh material files. Defaults to False.	False
gen_color_gif	bool	Generate GIF of color images. Defaults to False.	False
gen_color_mp4	bool	Generate MP4 of color images. Defaults to False.	False
gen_viewnormal_mp4	bool	Generate MP4 of view-space normals. Defaults to False.	False
gen_glonormal_mp4	bool	Generate MP4 of global-space normals. Defaults to False.	False
no_index_file	bool	Skip saving index file. Defaults to False.	False
light_factor	float	PBR light intensity multiplier. Defaults to 1.0.	1.0

Example

from embodied_gen.data.differentiable_render import ImageRender
from embodied_gen.data.utils import CameraSetting
from embodied_gen.utils.enum import RenderItems

camera_params = CameraSetting(
    num_images=6,
    elevation=[20, -10],
    distance=5,
    resolution_hw=(512,512),
    fov=math.radians(30),
    device='cuda',
)
render_items = [RenderItems.IMAGE.value, RenderItems.DEPTH.value]
renderer = ImageRender(
    render_items,
    camera_params,
    with_mtl=args.with_mtl,
    gen_color_mp4=True,
)
renderer.render_mesh(mesh_path='mesh.obj', output_root='./renders')

Source code in embodied_gen/data/differentiable_render.py

def __init__(
    self,
    render_items: list[RenderItems],
    camera_params: CameraSetting,
    recompute_vtx_normal: bool = True,
    with_mtl: bool = False,
    gen_color_gif: bool = False,
    gen_color_mp4: bool = False,
    gen_viewnormal_mp4: bool = False,
    gen_glonormal_mp4: bool = False,
    no_index_file: bool = False,
    light_factor: float = 1.0,
) -> None:
    camera = init_kal_camera(camera_params)
    self.camera = camera

    # Setup MVP matrix and renderer.
    mv = camera.view_matrix()  # (n 4 4) world2cam
    p = camera.intrinsics.projection_matrix()
    # NOTE: add a negative sign at P[0, 2] as the y axis is flipped in `nvdiffrast` output.  # noqa
    p[:, 1, 1] = -p[:, 1, 1]
    # mvp = torch.bmm(p, mv) # camera.view_projection_matrix()
    self.mv = mv
    self.p = p

    renderer = DiffrastRender(
        p_matrix=p,
        mv_matrix=mv,
        resolution_hw=camera_params.resolution_hw,
        context=dr.RasterizeCudaContext(),
        mask_thresh=0.5,
        grad_db=False,
        device=camera_params.device,
        antialias_mask=True,
    )
    self.renderer = renderer
    self.recompute_vtx_normal = recompute_vtx_normal
    self.render_items = render_items
    self.device = camera_params.device
    self.with_mtl = with_mtl
    self.gen_color_gif = gen_color_gif
    self.gen_color_mp4 = gen_color_mp4
    self.gen_viewnormal_mp4 = gen_viewnormal_mp4
    self.gen_glonormal_mp4 = gen_glonormal_mp4
    self.light_factor = light_factor
    self.no_index_file = no_index_file

__call__

__call__(mesh_path: str, output_dir: str, prompt: str = None) -> dict[str, str]

Renders a single mesh and returns output paths.

Parameters:

Name	Type	Description	Default
mesh_path	str	Path to mesh file.	required
output_dir	str	Directory to save outputs.	required
prompt	str	Caption prompt for MP4 metadata.	None

Returns:

Type	Description
dict[str, str]	dict[str, str]: Mapping of render types to saved image paths.

Source code in embodied_gen/data/differentiable_render.py

def __call__(
    self, mesh_path: str, output_dir: str, prompt: str = None
) -> dict[str, str]:
    """Renders a single mesh and returns output paths.

    Args:
        mesh_path (str): Path to mesh file.
        output_dir (str): Directory to save outputs.
        prompt (str, optional): Caption prompt for MP4 metadata.

    Returns:
        dict[str, str]: Mapping of render types to saved image paths.
    """
    try:
        mesh = import_kaolin_mesh(mesh_path, self.with_mtl)
    except Exception as e:
        logger.error(f"[ERROR MESH LOAD]: {e}, skip {mesh_path}")
        return

    mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)
    if self.recompute_vtx_normal:
        mesh.vertex_normals = calc_vertex_normals(
            mesh.vertices, mesh.faces
        )

    mesh = mesh.to(self.device)
    vertices, faces, vertex_normals = (
        mesh.vertices,
        mesh.faces,
        mesh.vertex_normals,
    )

    # Perform rendering.
    data_dict = defaultdict(list)
    if RenderItems.ALPHA.value in self.render_items:
        masks, _ = self.renderer.render_rast_alpha(vertices, faces)
        render_paths = save_images(
            masks, f"{output_dir}/{RenderItems.ALPHA}"
        )
        data_dict[RenderItems.ALPHA.value] = render_paths

    if RenderItems.GLOBAL_NORMAL.value in self.render_items:
        rendered_normals, masks = self.renderer.render_global_normal(
            vertices, faces, vertex_normals
        )
        if self.gen_glonormal_mp4:
            if isinstance(rendered_normals, torch.Tensor):
                rendered_normals = rendered_normals.detach().cpu().numpy()
            create_mp4_from_images(
                rendered_normals,
                output_path=f"{output_dir}/normal.mp4",
                fps=15,
                prompt=prompt,
            )
        else:
            render_paths = save_images(
                rendered_normals,
                f"{output_dir}/{RenderItems.GLOBAL_NORMAL}",
                cvt_color=cv2.COLOR_BGR2RGB,
            )
            data_dict[RenderItems.GLOBAL_NORMAL.value] = render_paths

        if RenderItems.VIEW_NORMAL.value in self.render_items:
            assert (
                RenderItems.GLOBAL_NORMAL in self.render_items
            ), f"Must render global normal firstly, got render_items: {self.render_items}."  # noqa
            rendered_view_normals = self.renderer.transform_normal(
                rendered_normals, self.mv, masks, to_view=True
            )

            if self.gen_viewnormal_mp4:
                create_mp4_from_images(
                    rendered_view_normals,
                    output_path=f"{output_dir}/view_normal.mp4",
                    fps=15,
                    prompt=prompt,
                )
            else:
                render_paths = save_images(
                    rendered_view_normals,
                    f"{output_dir}/{RenderItems.VIEW_NORMAL}",
                    cvt_color=cv2.COLOR_BGR2RGB,
                )
                data_dict[RenderItems.VIEW_NORMAL.value] = render_paths

    if RenderItems.POSITION_MAP.value in self.render_items:
        rendered_position, masks = self.renderer.render_position(
            vertices, faces
        )
        norm_position = self.renderer.normalize_map_by_mask(
            rendered_position, masks
        )
        render_paths = save_images(
            norm_position,
            f"{output_dir}/{RenderItems.POSITION_MAP}",
            cvt_color=cv2.COLOR_BGR2RGB,
        )
        data_dict[RenderItems.POSITION_MAP.value] = render_paths

    if RenderItems.DEPTH.value in self.render_items:
        rendered_depth, masks = self.renderer.render_depth(vertices, faces)
        norm_depth = self.renderer.normalize_map_by_mask(
            rendered_depth, masks
        )
        render_paths = save_images(
            norm_depth,
            f"{output_dir}/{RenderItems.DEPTH}",
        )
        data_dict[RenderItems.DEPTH.value] = render_paths

        render_paths = save_images(
            rendered_depth,
            f"{output_dir}/{RenderItems.DEPTH}_exr",
            to_uint8=False,
            format=".exr",
        )
        data_dict[f"{RenderItems.DEPTH.value}_exr"] = render_paths

    if RenderItems.IMAGE.value in self.render_items:
        images = []
        albedos = []
        diffuses = []
        masks, _ = self.renderer.render_rast_alpha(vertices, faces)
        try:
            for idx, cam in enumerate(self.camera):
                image, albedo, diffuse, _ = render_pbr(
                    mesh, cam, light_factor=self.light_factor
                )
                image = torch.cat([image[0], masks[idx]], axis=-1)
                images.append(image.detach().cpu().numpy())

                if RenderItems.ALBEDO.value in self.render_items:
                    albedo = torch.cat([albedo[0], masks[idx]], axis=-1)
                    albedos.append(albedo.detach().cpu().numpy())

                if RenderItems.DIFFUSE.value in self.render_items:
                    diffuse = torch.cat([diffuse[0], masks[idx]], axis=-1)
                    diffuses.append(diffuse.detach().cpu().numpy())

        except Exception as e:
            logger.error(f"[ERROR pbr render]: {e}, skip {mesh_path}")
            return

        if self.gen_color_gif:
            create_gif_from_images(
                images,
                output_path=f"{output_dir}/color.gif",
                fps=15,
            )

        if self.gen_color_mp4:
            create_mp4_from_images(
                images,
                output_path=f"{output_dir}/color.mp4",
                fps=15,
                prompt=prompt,
            )

        if self.gen_color_mp4 or self.gen_color_gif:
            return data_dict

        render_paths = save_images(
            images,
            f"{output_dir}/{RenderItems.IMAGE}",
            cvt_color=cv2.COLOR_BGRA2RGBA,
        )
        data_dict[RenderItems.IMAGE.value] = render_paths

        render_paths = save_images(
            albedos,
            f"{output_dir}/{RenderItems.ALBEDO}",
            cvt_color=cv2.COLOR_BGRA2RGBA,
        )
        data_dict[RenderItems.ALBEDO.value] = render_paths

        render_paths = save_images(
            diffuses,
            f"{output_dir}/{RenderItems.DIFFUSE}",
            cvt_color=cv2.COLOR_BGRA2RGBA,
        )
        data_dict[RenderItems.DIFFUSE.value] = render_paths

    data_dict["status"] = "success"

    logger.info(f"Finish rendering in {output_dir}")

    return data_dict

render_mesh

render_mesh(mesh_path: Union[str, List[str]], output_root: str, uuid: Union[str, List[str]] = None, prompts: List[str] = None) -> None

Renders one or more meshes and saves outputs.

Parameters:

Name	Type	Description	Default
mesh_path	Union[str, List[str]]	Path(s) to mesh files.	required
output_root	str	Directory to save outputs.	required
uuid	Union[str, List[str]]	Unique IDs for outputs.	None
prompts	List[str]	Text prompts for videos.	None

Source code in embodied_gen/data/differentiable_render.py

def render_mesh(
    self,
    mesh_path: Union[str, List[str]],
    output_root: str,
    uuid: Union[str, List[str]] = None,
    prompts: List[str] = None,
) -> None:
    """Renders one or more meshes and saves outputs.

    Args:
        mesh_path (Union[str, List[str]]): Path(s) to mesh files.
        output_root (str): Directory to save outputs.
        uuid (Union[str, List[str]], optional): Unique IDs for outputs.
        prompts (List[str], optional): Text prompts for videos.
    """
    mesh_path = as_list(mesh_path)
    if uuid is None:
        uuid = [os.path.basename(p).split(".")[0] for p in mesh_path]
    uuid = as_list(uuid)
    assert len(mesh_path) == len(uuid)
    os.makedirs(output_root, exist_ok=True)

    meta_info = dict()
    for idx, (path, uid) in tqdm(
        enumerate(zip(mesh_path, uuid)), total=len(mesh_path)
    ):
        output_dir = os.path.join(output_root, uid)
        os.makedirs(output_dir, exist_ok=True)
        prompt = prompts[idx] if prompts else None
        data_dict = self(path, output_dir, prompt)
        meta_info[uid] = data_dict

    if self.no_index_file:
        return

    index_file = os.path.join(output_root, "index.json")
    with open(index_file, "w") as fout:
        json.dump(meta_info, fout)

    logger.info(f"Rendering meta info logged in {index_file}")

create_gif_from_images

create_gif_from_images(images: list[ndarray], output_path: str, fps: int = 10) -> None

Creates a GIF animation from a list of images.

Parameters:

Name	Type	Description	Default
images	list[ndarray]	List of images as numpy arrays.	required
output_path	str	Path to save the GIF file.	required
fps	int	Frames per second. Defaults to 10.	10

Source code in embodied_gen/data/differentiable_render.py

def create_gif_from_images(
    images: list[np.ndarray], output_path: str, fps: int = 10
) -> None:
    """Creates a GIF animation from a list of images.

    Args:
        images (list[np.ndarray]): List of images as numpy arrays.
        output_path (str): Path to save the GIF file.
        fps (int, optional): Frames per second. Defaults to 10.
    """
    pil_images = []
    for image in images:
        image = image.clip(min=0, max=1)
        image = (255.0 * image).astype(np.uint8)
        image = Image.fromarray(image, mode="RGBA")
        pil_images.append(image.convert("RGB"))

    duration = 1000 // fps
    pil_images[0].save(
        output_path,
        save_all=True,
        append_images=pil_images[1:],
        duration=duration,
        loop=0,
    )

    logger.info(f"GIF saved to {output_path}")

create_mp4_from_images

create_mp4_from_images(images: list[ndarray], output_path: str, fps: int = 10, prompt: str = None)

Creates an MP4 video from a list of images.

Parameters:

Name	Type	Description	Default
images	list[ndarray]	List of images as numpy arrays.	required
output_path	str	Path to save the MP4 file.	required
fps	int	Frames per second. Defaults to 10.	10
prompt	str	Optional text prompt overlay.	None

Source code in embodied_gen/data/differentiable_render.py

def create_mp4_from_images(
    images: list[np.ndarray],
    output_path: str,
    fps: int = 10,
    prompt: str = None,
):
    """Creates an MP4 video from a list of images.

    Args:
        images (list[np.ndarray]): List of images as numpy arrays.
        output_path (str): Path to save the MP4 file.
        fps (int, optional): Frames per second. Defaults to 10.
        prompt (str, optional): Optional text prompt overlay.
    """
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 0.5
    font_thickness = 1
    color = (255, 255, 255)
    position = (20, 25)

    with imageio.get_writer(output_path, fps=fps) as writer:
        for image in images:
            image = image.clip(min=0, max=1)
            image = (255.0 * image).astype(np.uint8)
            image = image[..., :3]
            if prompt is not None:
                cv2.putText(
                    image,
                    prompt,
                    position,
                    font,
                    font_scale,
                    color,
                    font_thickness,
                )

            writer.append_data(image)

    logger.info(f"MP4 video saved to {output_path}")

embodied_gen.data.mesh_operator

MeshFixer

MeshFixer(vertices: Union[Tensor, ndarray], faces: Union[Tensor, ndarray], device: str = 'cuda')

Bases: object

MeshFixer simplifies and repairs 3D triangle meshes by TSDF.

Attributes:

Name	Type	Description
vertices	Tensor	A tensor of shape (V, 3) representing vertex positions.
faces	Tensor	A tensor of shape (F, 3) representing face indices.
device	str	Device to run computations on, typically "cuda" or "cpu".

Main logic reference: https://github.com/microsoft/TRELLIS/blob/main/trellis/utils/postprocessing_utils.py#L22

Source code in embodied_gen/data/mesh_operator.py

def __init__(
    self,
    vertices: Union[torch.Tensor, np.ndarray],
    faces: Union[torch.Tensor, np.ndarray],
    device: str = "cuda",
) -> None:
    self.device = device
    if isinstance(vertices, np.ndarray):
        vertices = torch.tensor(vertices)
    self.vertices = vertices

    if isinstance(faces, np.ndarray):
        faces = torch.tensor(faces)
    self.faces = faces

__call__

__call__(filter_ratio: float, max_hole_size: float, resolution: int, num_views: int, norm_mesh_ratio: float = 1.0) -> Tuple[np.ndarray, np.ndarray]

Post-process the mesh by simplifying and filling holes.

This method performs a two-step process: 1. Simplifies mesh by reducing faces using quadric edge decimation. 2. Fills holes by removing invisible faces, repairing small boundaries.

Parameters:

Name	Type	Description	Default
filter_ratio	float	Ratio of faces to simplify out. Must be in the range (0, 1).	required
max_hole_size	float	Maximum area of a hole to fill. Connected components of holes larger than this size will not be repaired.	required
resolution	int	Resolution of the rasterization buffer.	required
num_views	int	Number of viewpoints to sample for rasterization.	required
norm_mesh_ratio	float	A scaling factor applied to the vertices of the mesh during processing.	1.0

Returns:

Type	Description
Tuple[ndarray, ndarray]	Tuple[np.ndarray, np.ndarray]: - vertices: Simplified and repaired vertex array of (V, 3). - faces: Simplified and repaired face array of (F, 3).

Source code in embodied_gen/data/mesh_operator.py

@spaces.GPU
def __call__(
    self,
    filter_ratio: float,
    max_hole_size: float,
    resolution: int,
    num_views: int,
    norm_mesh_ratio: float = 1.0,
) -> Tuple[np.ndarray, np.ndarray]:
    """Post-process the mesh by simplifying and filling holes.

    This method performs a two-step process:
    1. Simplifies mesh by reducing faces using quadric edge decimation.
    2. Fills holes by removing invisible faces, repairing small boundaries.

    Args:
        filter_ratio (float): Ratio of faces to simplify out.
            Must be in the range (0, 1).
        max_hole_size (float): Maximum area of a hole to fill. Connected
            components of holes larger than this size will not be repaired.
        resolution (int): Resolution of the rasterization buffer.
        num_views (int): Number of viewpoints to sample for rasterization.
        norm_mesh_ratio (float, optional): A scaling factor applied to the
            vertices of the mesh during processing.

    Returns:
        Tuple[np.ndarray, np.ndarray]:
            - vertices: Simplified and repaired vertex array of (V, 3).
            - faces: Simplified and repaired face array of (F, 3).
    """
    self.vertices = self.vertices.to(self.device)
    self.faces = self.faces.to(self.device)

    self.simplify(ratio=filter_ratio)
    self.fill_holes(
        max_hole_size=max_hole_size,
        max_hole_nbe=int(250 * np.sqrt(1 - filter_ratio)),
        resolution=resolution,
        num_views=num_views,
        norm_mesh_ratio=norm_mesh_ratio,
    )

    return self.vertices_np, self.faces_np

simplify

simplify(ratio: float) -> None

Simplify the mesh using quadric edge collapse decimation.

Parameters:

Name	Type	Description	Default
ratio	float	Ratio of faces to filter out.	required

Source code in embodied_gen/data/mesh_operator.py

@log_mesh_changes
def simplify(self, ratio: float) -> None:
    """Simplify the mesh using quadric edge collapse decimation.

    Args:
        ratio (float): Ratio of faces to filter out.
    """
    if ratio <= 0 or ratio >= 1:
        raise ValueError("Simplify ratio must be between 0 and 1.")

    # Convert to PyVista format for simplification
    mesh = pv.PolyData(
        self.vertices_np,
        np.hstack([np.full((self.faces.shape[0], 1), 3), self.faces_np]),
    )
    mesh.clean(inplace=True)
    mesh.clear_data()
    mesh = mesh.triangulate()
    mesh = mesh.decimate(ratio, progress_bar=True)

    # Update vertices and faces
    self.vertices = torch.tensor(
        mesh.points, device=self.device, dtype=torch.float32
    )
    self.faces = torch.tensor(
        mesh.faces.reshape(-1, 4)[:, 1:],
        device=self.device,
        dtype=torch.int32,
    )

embodied_gen.data.backproject_v2

TextureBacker

TextureBacker(camera_params: CameraSetting, view_weights: list[float], render_wh: tuple[int, int] = (2048, 2048), texture_wh: tuple[int, int] = (2048, 2048), bake_angle_thresh: int = 75, mask_thresh: float = 0.5, smooth_texture: bool = True, inpaint_smooth: bool = False)

Texture baking pipeline for multi-view projection and fusion.

This class generates UV-based textures for a 3D mesh using multi-view images, depth, and normal information. It includes mesh normalization, UV unwrapping, visibility-aware back-projection, confidence-weighted fusion, and inpainting.

Parameters:

Name	Type	Description	Default
camera_params	CameraSetting	Camera intrinsics and extrinsics.	required
view_weights	list[float]	Weights for each view in texture fusion.	required
render_wh	tuple[int, int]	Intermediate rendering resolution.	(2048, 2048)
texture_wh	tuple[int, int]	Output texture resolution.	(2048, 2048)
bake_angle_thresh	int	Max angle for valid projection.	75
mask_thresh	float	Threshold for visibility masks.	0.5
smooth_texture	bool	Apply post-processing to texture.	True
inpaint_smooth	bool	Apply inpainting smoothing.	False

Example

from embodied_gen.data.backproject_v2 import TextureBacker
from embodied_gen.data.utils import CameraSetting
import trimesh
from PIL import Image

camera_params = CameraSetting(
    num_images=6,
    elevation=[20, -10],
    distance=5,
    resolution_hw=(2048,2048),
    fov=math.radians(30),
    device='cuda',
)
view_weights = [1, 0.1, 0.02, 0.1, 1, 0.02]
mesh = trimesh.load('mesh.obj')
images = [Image.open(f'view_{i}.png') for i in range(6)]
texture_backer = TextureBacker(camera_params, view_weights)
textured_mesh = texture_backer(images, mesh, 'output.obj')

Source code in embodied_gen/data/backproject_v2.py

def __init__(
    self,
    camera_params: CameraSetting,
    view_weights: list[float],
    render_wh: tuple[int, int] = (2048, 2048),
    texture_wh: tuple[int, int] = (2048, 2048),
    bake_angle_thresh: int = 75,
    mask_thresh: float = 0.5,
    smooth_texture: bool = True,
    inpaint_smooth: bool = False,
) -> None:
    self.camera_params = camera_params
    self.renderer = None
    self.view_weights = view_weights
    self.device = camera_params.device
    self.render_wh = render_wh
    self.texture_wh = texture_wh
    self.mask_thresh = mask_thresh
    self.smooth_texture = smooth_texture
    self.inpaint_smooth = inpaint_smooth

    self.bake_angle_thresh = bake_angle_thresh
    self.bake_unreliable_kernel_size = int(
        (2 / 512) * max(self.render_wh[0], self.render_wh[1])
    )

__call__

__call__(colors: list[Image], mesh: Trimesh, output_path: str) -> trimesh.Trimesh

Runs the texture baking and exports the textured mesh.

Parameters:

Name	Type	Description	Default
colors	list[Image]	List of input view images.	required
mesh	Trimesh	Input mesh to be textured.	required
output_path	str	Path to save the output textured mesh.	required

Returns:

Type	Description
Trimesh	trimesh.Trimesh: The textured mesh with UV and texture image.

Source code in embodied_gen/data/backproject_v2.py

def __call__(
    self,
    colors: list[Image.Image],
    mesh: trimesh.Trimesh,
    output_path: str,
) -> trimesh.Trimesh:
    """Runs the texture baking and exports the textured mesh.

    Args:
        colors (list[Image.Image]): List of input view images.
        mesh (trimesh.Trimesh): Input mesh to be textured.
        output_path (str): Path to save the output textured mesh.

    Returns:
        trimesh.Trimesh: The textured mesh with UV and texture image.
    """
    mesh = self.load_mesh(mesh)
    texture_np, mask_np = self.compute_texture(colors, mesh)

    texture_np = self.uv_inpaint(mesh, texture_np, mask_np)
    if self.smooth_texture:
        texture_np = post_process_texture(texture_np)

    vertices, faces, uv_map = self.get_mesh_np_attrs(
        mesh, self.scale, self.center
    )
    textured_mesh = save_mesh_with_mtl(
        vertices, faces, uv_map, texture_np, output_path
    )

    return textured_mesh

back_project

back_project(image, vis_mask, depth, normal, uv) -> tuple[torch.Tensor, torch.Tensor]

Back-projects image and confidence to UV texture space.

Parameters:

Name	Type	Description	Default
image	Image or ndarray	Input image.	required
vis_mask	Tensor	Visibility mask.	required
depth	Tensor	Depth map.	required
normal	Tensor	Normal map.	required
uv	Tensor	UV coordinates.	required

Returns:

Type	Description
tuple[Tensor, Tensor]	tuple[torch.Tensor, torch.Tensor]: Texture and confidence map.

Source code in embodied_gen/data/backproject_v2.py

def back_project(
    self, image, vis_mask, depth, normal, uv
) -> tuple[torch.Tensor, torch.Tensor]:
    """Back-projects image and confidence to UV texture space.

    Args:
        image (PIL.Image or np.ndarray): Input image.
        vis_mask (torch.Tensor): Visibility mask.
        depth (torch.Tensor): Depth map.
        normal (torch.Tensor): Normal map.
        uv (torch.Tensor): UV coordinates.

    Returns:
        tuple[torch.Tensor, torch.Tensor]: Texture and confidence map.
    """
    image = np.array(image)
    image = torch.as_tensor(image, device=self.device, dtype=torch.float32)
    if image.ndim == 2:
        image = image.unsqueeze(-1)
    image = image / 255

    depth_inv = (1.0 - depth) * vis_mask
    sketch_image = self._render_depth_edges(depth_inv)

    cos = F.cosine_similarity(
        torch.tensor([[0, 0, 1]], device=self.device),
        normal.view(-1, 3),
    ).view_as(normal[..., :1])
    cos[cos < np.cos(np.radians(self.bake_angle_thresh))] = 0

    k = self.bake_unreliable_kernel_size * 2 + 1
    kernel = torch.ones((1, 1, k, k), device=self.device)

    vis_mask = vis_mask.permute(2, 0, 1).unsqueeze(0).float()
    vis_mask = F.conv2d(
        1.0 - vis_mask,
        kernel,
        padding=k // 2,
    )
    vis_mask = 1.0 - (vis_mask > 0).float()
    vis_mask = vis_mask.squeeze(0).permute(1, 2, 0)

    sketch_image = sketch_image.permute(2, 0, 1).unsqueeze(0)
    sketch_image = F.conv2d(sketch_image, kernel, padding=k // 2)
    sketch_image = (sketch_image > 0).float()
    sketch_image = sketch_image.squeeze(0).permute(1, 2, 0)
    vis_mask = vis_mask * (sketch_image < 0.5)

    cos[vis_mask == 0] = 0
    valid_pixels = (vis_mask != 0).view(-1)

    return (
        self._scatter_texture(uv, image, valid_pixels),
        self._scatter_texture(uv, cos, valid_pixels),
    )

compute_enhanced_viewnormal

compute_enhanced_viewnormal(mv_mtx: Tensor, vertices: Tensor, faces: Tensor) -> torch.Tensor

Computes enhanced view normals for mesh faces.

Parameters:

Name	Type	Description	Default
mv_mtx	Tensor	View matrices.	required
vertices	Tensor	Mesh vertices.	required
faces	Tensor	Mesh faces.	required

Returns:

Type	Description
Tensor	torch.Tensor: View normals.

Source code in embodied_gen/data/backproject_v2.py

def compute_enhanced_viewnormal(
    self, mv_mtx: torch.Tensor, vertices: torch.Tensor, faces: torch.Tensor
) -> torch.Tensor:
    """Computes enhanced view normals for mesh faces.

    Args:
        mv_mtx (torch.Tensor): View matrices.
        vertices (torch.Tensor): Mesh vertices.
        faces (torch.Tensor): Mesh faces.

    Returns:
        torch.Tensor: View normals.
    """
    rast, _ = self.renderer.compute_dr_raster(vertices, faces)
    rendered_view_normals = []
    for idx in range(len(mv_mtx)):
        pos_cam = _transform_vertices(mv_mtx[idx], vertices, keepdim=True)
        pos_cam = pos_cam[:, :3] / pos_cam[:, 3:]
        v0, v1, v2 = (pos_cam[faces[:, i]] for i in range(3))
        face_norm = F.normalize(
            torch.cross(v1 - v0, v2 - v0, dim=-1), dim=-1
        )
        vertex_norm = (
            torch.from_numpy(
                trimesh.geometry.mean_vertex_normals(
                    len(pos_cam), faces.cpu(), face_norm.cpu()
                )
            )
            .to(vertices.device)
            .contiguous()
        )
        im_base_normals, _ = dr.interpolate(
            vertex_norm[None, ...].float(),
            rast[idx : idx + 1],
            faces.to(torch.int32),
        )
        rendered_view_normals.append(im_base_normals)

    rendered_view_normals = torch.cat(rendered_view_normals, dim=0)

    return rendered_view_normals

compute_texture

compute_texture(colors: list[Image], mesh: Trimesh) -> trimesh.Trimesh

Computes the fused texture for the mesh from multi-view images.

Parameters:

Name	Type	Description	Default
colors	list[Image]	List of view images.	required
mesh	Trimesh	Mesh to texture.	required

Returns:

Type	Description
Trimesh	tuple[np.ndarray, np.ndarray]: Texture and mask.

Source code in embodied_gen/data/backproject_v2.py

@spaces.GPU
def compute_texture(
    self,
    colors: list[Image.Image],
    mesh: trimesh.Trimesh,
) -> trimesh.Trimesh:
    """Computes the fused texture for the mesh from multi-view images.

    Args:
        colors (list[Image.Image]): List of view images.
        mesh (trimesh.Trimesh): Mesh to texture.

    Returns:
        tuple[np.ndarray, np.ndarray]: Texture and mask.
    """
    self._lazy_init_render(self.camera_params, self.mask_thresh)

    vertices = torch.from_numpy(mesh.vertices).to(self.device).float()
    faces = torch.from_numpy(mesh.faces).to(self.device).to(torch.int)
    uv_map = torch.from_numpy(mesh.visual.uv).to(self.device).float()

    rendered_depth, masks = self.renderer.render_depth(vertices, faces)
    norm_deps = self.renderer.normalize_map_by_mask(rendered_depth, masks)
    render_uvs, _ = self.renderer.render_uv(vertices, faces, uv_map)
    view_normals = self.compute_enhanced_viewnormal(
        self.renderer.mv_mtx, vertices, faces
    )

    textures, weighted_cos_maps = [], []
    for color, mask, dep, normal, uv, weight in zip(
        colors,
        masks,
        norm_deps,
        view_normals,
        render_uvs,
        self.view_weights,
    ):
        texture, cos_map = self.back_project(color, mask, dep, normal, uv)
        textures.append(texture)
        weighted_cos_maps.append(weight * (cos_map**4))

    texture, mask = self.fast_bake_texture(textures, weighted_cos_maps)

    texture_np = texture.cpu().numpy()
    mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype(np.uint8)

    return texture_np, mask_np

fast_bake_texture

fast_bake_texture(textures: list[Tensor], confidence_maps: list[Tensor]) -> tuple[torch.Tensor, torch.Tensor]

Fuses multiple textures and confidence maps.

Parameters:

Name	Type	Description	Default
textures	list[Tensor]	List of textures.	required
confidence_maps	list[Tensor]	List of confidence maps.	required

Returns:

Type	Description
tuple[Tensor, Tensor]	tuple[torch.Tensor, torch.Tensor]: Fused texture and mask.

Source code in embodied_gen/data/backproject_v2.py

@torch.no_grad()
def fast_bake_texture(
    self, textures: list[torch.Tensor], confidence_maps: list[torch.Tensor]
) -> tuple[torch.Tensor, torch.Tensor]:
    """Fuses multiple textures and confidence maps.

    Args:
        textures (list[torch.Tensor]): List of textures.
        confidence_maps (list[torch.Tensor]): List of confidence maps.

    Returns:
        tuple[torch.Tensor, torch.Tensor]: Fused texture and mask.
    """
    channel = textures[0].shape[-1]
    texture_merge = torch.zeros(self.texture_wh + [channel]).to(
        self.device
    )
    trust_map_merge = torch.zeros(self.texture_wh + [1]).to(self.device)
    for texture, cos_map in zip(textures, confidence_maps):
        view_sum = (cos_map > 0).sum()
        painted_sum = ((cos_map > 0) * (trust_map_merge > 0)).sum()
        if painted_sum / view_sum > 0.99:
            continue
        texture_merge += texture * cos_map
        trust_map_merge += cos_map
    texture_merge = texture_merge / torch.clamp(trust_map_merge, min=1e-8)

    return texture_merge, trust_map_merge > 1e-8

get_mesh_np_attrs

get_mesh_np_attrs(mesh: Trimesh, scale: float = None, center: ndarray = None) -> tuple[np.ndarray, np.ndarray, np.ndarray]

Gets mesh attributes as numpy arrays.

Parameters:

Name	Type	Description	Default
mesh	Trimesh	Input mesh.	required
scale	float	Scale factor.	None
center	ndarray	Center offset.	None

Returns:

Name	Type	Description
tuple	tuple[ndarray, ndarray, ndarray]	(vertices, faces, uv_map)

Source code in embodied_gen/data/backproject_v2.py

def get_mesh_np_attrs(
    self,
    mesh: trimesh.Trimesh,
    scale: float = None,
    center: np.ndarray = None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Gets mesh attributes as numpy arrays.

    Args:
        mesh (trimesh.Trimesh): Input mesh.
        scale (float, optional): Scale factor.
        center (np.ndarray, optional): Center offset.

    Returns:
        tuple: (vertices, faces, uv_map)
    """
    vertices = mesh.vertices.copy()
    faces = mesh.faces.copy()
    uv_map = mesh.visual.uv.copy()
    uv_map[:, 1] = 1.0 - uv_map[:, 1]

    if scale is not None:
        vertices = vertices / scale
    if center is not None:
        vertices = vertices + center

    return vertices, faces, uv_map

load_mesh

load_mesh(mesh: Trimesh) -> trimesh.Trimesh

Normalizes mesh and unwraps UVs.

Parameters:

Name	Type	Description	Default
mesh	Trimesh	Input mesh.	required

Returns:

Type	Description
Trimesh	trimesh.Trimesh: Mesh with normalized vertices and UVs.

Source code in embodied_gen/data/backproject_v2.py

def load_mesh(self, mesh: trimesh.Trimesh) -> trimesh.Trimesh:
    """Normalizes mesh and unwraps UVs.

    Args:
        mesh (trimesh.Trimesh): Input mesh.

    Returns:
        trimesh.Trimesh: Mesh with normalized vertices and UVs.
    """
    mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)
    self.scale, self.center = scale, center

    vmapping, indices, uvs = xatlas.parametrize(mesh.vertices, mesh.faces)
    uvs[:, 1] = 1 - uvs[:, 1]
    mesh.vertices = mesh.vertices[vmapping]
    mesh.faces = indices
    mesh.visual.uv = uvs

    return mesh

uv_inpaint

uv_inpaint(mesh: Trimesh, texture: ndarray, mask: ndarray) -> np.ndarray

Inpaints missing regions in the UV texture.

Parameters:

Name	Type	Description	Default
mesh	Trimesh	Mesh.	required
texture	ndarray	Texture image.	required
mask	ndarray	Mask image.	required

Returns:

Type	Description
ndarray	np.ndarray: Inpainted texture.

Source code in embodied_gen/data/backproject_v2.py

def uv_inpaint(
    self, mesh: trimesh.Trimesh, texture: np.ndarray, mask: np.ndarray
) -> np.ndarray:
    """Inpaints missing regions in the UV texture.

    Args:
        mesh (trimesh.Trimesh): Mesh.
        texture (np.ndarray): Texture image.
        mask (np.ndarray): Mask image.

    Returns:
        np.ndarray: Inpainted texture.
    """
    if self.inpaint_smooth:
        vertices, faces, uv_map = self.get_mesh_np_attrs(mesh)
        texture, mask = _texture_inpaint_smooth(
            texture, mask, vertices, faces, uv_map
        )

    texture = texture.clip(0, 1)
    texture = cv2.inpaint(
        (texture * 255).astype(np.uint8),
        255 - mask,
        3,
        cv2.INPAINT_NS,
    )

    return texture

entrypoint

entrypoint(delight_model: DelightingModel = None, imagesr_model: ImageRealESRGAN = None, **kwargs) -> trimesh.Trimesh

Entrypoint for texture backprojection from multi-view images.

Parameters:

Name	Type	Description	Default
delight_model	DelightingModel	Delighting model.	None
imagesr_model	ImageRealESRGAN	Super-resolution model.	None
**kwargs		Additional arguments to override CLI.	{}

Returns:

Type	Description
Trimesh	trimesh.Trimesh: Textured mesh.

Source code in embodied_gen/data/backproject_v2.py

def entrypoint(
    delight_model: DelightingModel = None,
    imagesr_model: ImageRealESRGAN = None,
    **kwargs,
) -> trimesh.Trimesh:
    """Entrypoint for texture backprojection from multi-view images.

    Args:
        delight_model (DelightingModel, optional): Delighting model.
        imagesr_model (ImageRealESRGAN, optional): Super-resolution model.
        **kwargs: Additional arguments to override CLI.

    Returns:
        trimesh.Trimesh: Textured mesh.
    """
    args = parse_args()
    for k, v in kwargs.items():
        if hasattr(args, k) and v is not None:
            setattr(args, k, v)

    # Setup camera parameters.
    camera_params = CameraSetting(
        num_images=args.num_images,
        elevation=args.elevation,
        distance=args.distance,
        resolution_hw=args.resolution_hw,
        fov=math.radians(args.fov),
        device=args.device,
    )

    args.color_path = as_list(args.color_path)
    if args.delight and delight_model is None:
        delight_model = DelightingModel()

    color_grid = [Image.open(color_path) for color_path in args.color_path]
    color_grid = vcat_pil_images(color_grid, image_mode="RGBA")
    if args.delight:
        color_grid = delight_model(color_grid)
        if not args.no_save_delight_img:
            save_dir = os.path.dirname(args.output_path)
            os.makedirs(save_dir, exist_ok=True)
            color_grid.save(f"{save_dir}/color_delight.png")

    multiviews = get_images_from_grid(color_grid, img_size=512)
    view_weights = [1, 0.1, 0.02, 0.1, 1, 0.02]
    view_weights += [0.01] * (len(multiviews) - len(view_weights))

    # Use RealESRGAN_x4plus for x4 (512->2048) image super resolution.
    if imagesr_model is None:
        imagesr_model = ImageRealESRGAN(outscale=4)
    multiviews = [imagesr_model(img) for img in multiviews]
    multiviews = [img.convert("RGB") for img in multiviews]
    mesh = trimesh.load(args.mesh_path)
    if isinstance(mesh, trimesh.Scene):
        mesh = mesh.dump(concatenate=True)

    if not args.skip_fix_mesh:
        mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)
        mesh_fixer = MeshFixer(mesh.vertices, mesh.faces, args.device)
        mesh.vertices, mesh.faces = mesh_fixer(
            filter_ratio=args.mesh_sipmlify_ratio,
            max_hole_size=0.04,
            resolution=1024,
            num_views=1000,
            norm_mesh_ratio=0.5,
        )
        if len(mesh.faces) > args.n_max_faces:
            mesh.vertices, mesh.faces = mesh_fixer(
                filter_ratio=0.8,
                max_hole_size=0.04,
                resolution=1024,
                num_views=1000,
                norm_mesh_ratio=0.5,
            )
        # Restore scale.
        mesh.vertices = mesh.vertices / scale
        mesh.vertices = mesh.vertices + center

    # Baking texture to mesh.
    texture_backer = TextureBacker(
        camera_params=camera_params,
        view_weights=view_weights,
        render_wh=args.resolution_hw,
        texture_wh=args.texture_wh,
        smooth_texture=not args.no_smooth_texture,
    )

    textured_mesh = texture_backer(multiviews, mesh, args.output_path)

    if args.save_glb_path is not None:
        os.makedirs(os.path.dirname(args.save_glb_path), exist_ok=True)
        textured_mesh.export(args.save_glb_path)

    return textured_mesh

parse_args

parse_args()

Parses command-line arguments for texture backprojection.

Returns:

Type	Description
	argparse.Namespace: Parsed arguments.

Source code in embodied_gen/data/backproject_v2.py

def parse_args():
    """Parses command-line arguments for texture backprojection.

    Returns:
        argparse.Namespace: Parsed arguments.
    """
    parser = argparse.ArgumentParser(description="Backproject texture")
    parser.add_argument(
        "--color_path",
        nargs="+",
        type=str,
        help="Multiview color image in grid file paths",
    )
    parser.add_argument(
        "--mesh_path",
        type=str,
        help="Mesh path, .obj, .glb or .ply",
    )
    parser.add_argument(
        "--output_path",
        type=str,
        help="Output mesh path with suffix",
    )
    parser.add_argument(
        "--num_images", type=int, default=6, help="Number of images to render."
    )
    parser.add_argument(
        "--elevation",
        nargs="+",
        type=float,
        default=[20.0, -10.0],
        help="Elevation angles for the camera (default: [20.0, -10.0])",
    )
    parser.add_argument(
        "--distance",
        type=float,
        default=5,
        help="Camera distance (default: 5)",
    )
    parser.add_argument(
        "--resolution_hw",
        type=int,
        nargs=2,
        default=(2048, 2048),
        help="Resolution of the output images (default: (2048, 2048))",
    )
    parser.add_argument(
        "--fov",
        type=float,
        default=30,
        help="Field of view in degrees (default: 30)",
    )
    parser.add_argument(
        "--device",
        type=str,
        choices=["cpu", "cuda"],
        default="cuda",
        help="Device to run on (default: `cuda`)",
    )
    parser.add_argument(
        "--skip_fix_mesh", action="store_true", help="Fix mesh geometry."
    )
    parser.add_argument(
        "--texture_wh",
        nargs=2,
        type=int,
        default=[2048, 2048],
        help="Texture resolution width and height",
    )
    parser.add_argument(
        "--mesh_sipmlify_ratio",
        type=float,
        default=0.9,
        help="Mesh simplification ratio (default: 0.9)",
    )
    parser.add_argument(
        "--delight", action="store_true", help="Use delighting model."
    )
    parser.add_argument(
        "--no_smooth_texture",
        action="store_true",
        help="Do not smooth the texture.",
    )
    parser.add_argument(
        "--save_glb_path", type=str, default=None, help="Save glb path."
    )
    parser.add_argument(
        "--no_save_delight_img",
        action="store_true",
        help="Disable saving delight image",
    )
    parser.add_argument("--n_max_faces", type=int, default=30000)
    args, unknown = parser.parse_known_args()

    return args

embodied_gen.data.convex_decomposer

decompose_convex_coacd

decompose_convex_coacd(filename: str, outfile: str, params: dict, verbose: bool = False, auto_scale: bool = True, scale_factor: float = 1.0) -> None

Decomposes a mesh using CoACD and saves the result.

This function loads a mesh from a file, runs the CoACD algorithm with the given parameters, optionally scales the resulting convex hulls to match the original mesh's bounding box, and exports the combined result to a file.

Parameters:

Name	Type	Description	Default
filename	str	Path to the input mesh file.	required
outfile	str	Path to save the decomposed output mesh.	required
params	dict	A dictionary of parameters for the CoACD algorithm.	required
verbose	bool	If True, sets the CoACD log level to 'info'.	False
auto_scale	bool	If True, automatically computes a scale factor to match the decomposed mesh's bounding box to the visual mesh's bounding box.	True
scale_factor	float	An additional scaling factor applied to the vertices of the decomposed mesh parts.	1.0

Source code in embodied_gen/data/convex_decomposer.py

def decompose_convex_coacd(
    filename: str,
    outfile: str,
    params: dict,
    verbose: bool = False,
    auto_scale: bool = True,
    scale_factor: float = 1.0,
) -> None:
    """Decomposes a mesh using CoACD and saves the result.

    This function loads a mesh from a file, runs the CoACD algorithm with the
    given parameters, optionally scales the resulting convex hulls to match the
    original mesh's bounding box, and exports the combined result to a file.

    Args:
        filename: Path to the input mesh file.
        outfile: Path to save the decomposed output mesh.
        params: A dictionary of parameters for the CoACD algorithm.
        verbose: If True, sets the CoACD log level to 'info'.
        auto_scale: If True, automatically computes a scale factor to match the
            decomposed mesh's bounding box to the visual mesh's bounding box.
        scale_factor: An additional scaling factor applied to the vertices of
            the decomposed mesh parts.
    """
    coacd.set_log_level("info" if verbose else "warn")

    mesh = trimesh.load(filename, force="mesh")
    mesh = coacd.Mesh(mesh.vertices, mesh.faces)

    result = coacd.run_coacd(mesh, **params)

    meshes = []
    for v, f in result:
        meshes.append(trimesh.Trimesh(v, f))

    # Compute collision_scale because convex decomposition usually makes the mesh larger.
    if auto_scale:
        all_mesh = sum([trimesh.Trimesh(*m) for m in result])
        convex_mesh_shape = np.ptp(all_mesh.vertices, axis=0)
        visual_mesh_shape = np.ptp(mesh.vertices, axis=0)
        scale_factor *= visual_mesh_shape / convex_mesh_shape

    combined = trimesh.Scene()
    for mesh_part in meshes:
        mesh_part.vertices *= scale_factor
        combined.add_geometry(mesh_part)

    combined.export(outfile)

decompose_convex_mesh

decompose_convex_mesh(filename: str, outfile: str, threshold: float = 0.05, max_convex_hull: int = -1, preprocess_mode: str = 'auto', preprocess_resolution: int = 30, resolution: int = 2000, mcts_nodes: int = 20, mcts_iterations: int = 150, mcts_max_depth: int = 3, pca: bool = False, merge: bool = True, seed: int = 0, auto_scale: bool = True, scale_factor: float = 1.005, verbose: bool = False) -> str

Decomposes a mesh into convex parts with retry logic.

This function serves as a wrapper for decompose_convex_coacd, providing explicit parameters for the CoACD algorithm and implementing a retry mechanism. If the initial decomposition fails, it attempts again with preprocess_mode set to 'on'.

Parameters:

Name	Type	Description	Default
filename	str	Path to the input mesh file.	required
outfile	str	Path to save the decomposed output mesh.	required
threshold	float	CoACD parameter. See CoACD documentation for details.	0.05
max_convex_hull	int	CoACD parameter. See CoACD documentation for details.	-1
preprocess_mode	str	CoACD parameter. See CoACD documentation for details.	'auto'
preprocess_resolution	int	CoACD parameter. See CoACD documentation for details.	30
resolution	int	CoACD parameter. See CoACD documentation for details.	2000
mcts_nodes	int	CoACD parameter. See CoACD documentation for details.	20
mcts_iterations	int	CoACD parameter. See CoACD documentation for details.	150
mcts_max_depth	int	CoACD parameter. See CoACD documentation for details.	3
pca	bool	CoACD parameter. See CoACD documentation for details.	False
merge	bool	CoACD parameter. See CoACD documentation for details.	True
seed	int	CoACD parameter. See CoACD documentation for details.	0
auto_scale	bool	If True, automatically scale the output to match the input bounding box.	True
scale_factor	float	Additional scaling factor to apply.	1.005
verbose	bool	If True, enables detailed logging.	False

Returns:

Type	Description
str	The path to the output file if decomposition is successful.

Raises:

Type	Description
RuntimeError	If convex decomposition fails after all attempts.

Source code in embodied_gen/data/convex_decomposer.py

def decompose_convex_mesh(
    filename: str,
    outfile: str,
    threshold: float = 0.05,
    max_convex_hull: int = -1,
    preprocess_mode: str = "auto",
    preprocess_resolution: int = 30,
    resolution: int = 2000,
    mcts_nodes: int = 20,
    mcts_iterations: int = 150,
    mcts_max_depth: int = 3,
    pca: bool = False,
    merge: bool = True,
    seed: int = 0,
    auto_scale: bool = True,
    scale_factor: float = 1.005,
    verbose: bool = False,
) -> str:
    """Decomposes a mesh into convex parts with retry logic.

    This function serves as a wrapper for `decompose_convex_coacd`, providing
    explicit parameters for the CoACD algorithm and implementing a retry
    mechanism. If the initial decomposition fails, it attempts again with
    `preprocess_mode` set to 'on'.

    Args:
        filename: Path to the input mesh file.
        outfile: Path to save the decomposed output mesh.
        threshold: CoACD parameter. See CoACD documentation for details.
        max_convex_hull: CoACD parameter. See CoACD documentation for details.
        preprocess_mode: CoACD parameter. See CoACD documentation for details.
        preprocess_resolution: CoACD parameter. See CoACD documentation for details.
        resolution: CoACD parameter. See CoACD documentation for details.
        mcts_nodes: CoACD parameter. See CoACD documentation for details.
        mcts_iterations: CoACD parameter. See CoACD documentation for details.
        mcts_max_depth: CoACD parameter. See CoACD documentation for details.
        pca: CoACD parameter. See CoACD documentation for details.
        merge: CoACD parameter. See CoACD documentation for details.
        seed: CoACD parameter. See CoACD documentation for details.
        auto_scale: If True, automatically scale the output to match the input
            bounding box.
        scale_factor: Additional scaling factor to apply.
        verbose: If True, enables detailed logging.

    Returns:
        The path to the output file if decomposition is successful.

    Raises:
        RuntimeError: If convex decomposition fails after all attempts.
    """
    coacd.set_log_level("info" if verbose else "warn")

    if os.path.exists(outfile):
        logger.warning(f"Output file {outfile} already exists, removing it.")
        os.remove(outfile)

    params = dict(
        threshold=threshold,
        max_convex_hull=max_convex_hull,
        preprocess_mode=preprocess_mode,
        preprocess_resolution=preprocess_resolution,
        resolution=resolution,
        mcts_nodes=mcts_nodes,
        mcts_iterations=mcts_iterations,
        mcts_max_depth=mcts_max_depth,
        pca=pca,
        merge=merge,
        seed=seed,
    )

    try:
        decompose_convex_coacd(
            filename, outfile, params, verbose, auto_scale, scale_factor
        )
        if os.path.exists(outfile):
            return outfile
    except Exception as e:
        if verbose:
            print(f"Decompose convex first attempt failed: {e}.")

    if preprocess_mode != "on":
        try:
            params["preprocess_mode"] = "on"
            decompose_convex_coacd(
                filename, outfile, params, verbose, auto_scale, scale_factor
            )
            if os.path.exists(outfile):
                return outfile
        except Exception as e:
            if verbose:
                print(
                    f"Decompose convex second attempt with preprocess_mode='on' failed: {e}"
                )

    raise RuntimeError(f"Convex decomposition failed on {filename}")

decompose_convex_mp

decompose_convex_mp(filename: str, outfile: str, threshold: float = 0.05, max_convex_hull: int = -1, preprocess_mode: str = 'auto', preprocess_resolution: int = 30, resolution: int = 2000, mcts_nodes: int = 20, mcts_iterations: int = 150, mcts_max_depth: int = 3, pca: bool = False, merge: bool = True, seed: int = 0, verbose: bool = False, auto_scale: bool = True) -> str

Decomposes a mesh into convex parts in a separate process.

This function uses the multiprocessing module to run the CoACD algorithm in a spawned subprocess. This is useful for isolating the decomposition process to prevent potential memory leaks or crashes in the main process. It includes a retry mechanism similar to decompose_convex_mesh.

See https://simulately.wiki/docs/toolkits/ConvexDecomp for details.

Parameters:

Name	Type	Description	Default
filename	str	Path to the input mesh file.	required
outfile	str	Path to save the decomposed output mesh.	required
threshold	float	CoACD parameter.	0.05
max_convex_hull	int	CoACD parameter.	-1
preprocess_mode	str	CoACD parameter.	'auto'
preprocess_resolution	int	CoACD parameter.	30
resolution	int	CoACD parameter.	2000
mcts_nodes	int	CoACD parameter.	20
mcts_iterations	int	CoACD parameter.	150
mcts_max_depth	int	CoACD parameter.	3
pca	bool	CoACD parameter.	False
merge	bool	CoACD parameter.	True
seed	int	CoACD parameter.	0
verbose	bool	If True, enables detailed logging in the subprocess.	False
auto_scale	bool	If True, automatically scale the output.	True

Returns:

Type	Description
str	The path to the output file if decomposition is successful.

Raises:

Type	Description
RuntimeError	If convex decomposition fails after all attempts.

Source code in embodied_gen/data/convex_decomposer.py

def decompose_convex_mp(
    filename: str,
    outfile: str,
    threshold: float = 0.05,
    max_convex_hull: int = -1,
    preprocess_mode: str = "auto",
    preprocess_resolution: int = 30,
    resolution: int = 2000,
    mcts_nodes: int = 20,
    mcts_iterations: int = 150,
    mcts_max_depth: int = 3,
    pca: bool = False,
    merge: bool = True,
    seed: int = 0,
    verbose: bool = False,
    auto_scale: bool = True,
) -> str:
    """Decomposes a mesh into convex parts in a separate process.

    This function uses the `multiprocessing` module to run the CoACD algorithm
    in a spawned subprocess. This is useful for isolating the decomposition
    process to prevent potential memory leaks or crashes in the main process.
    It includes a retry mechanism similar to `decompose_convex_mesh`.

    See https://simulately.wiki/docs/toolkits/ConvexDecomp for details.

    Args:
        filename: Path to the input mesh file.
        outfile: Path to save the decomposed output mesh.
        threshold: CoACD parameter.
        max_convex_hull: CoACD parameter.
        preprocess_mode: CoACD parameter.
        preprocess_resolution: CoACD parameter.
        resolution: CoACD parameter.
        mcts_nodes: CoACD parameter.
        mcts_iterations: CoACD parameter.
        mcts_max_depth: CoACD parameter.
        pca: CoACD parameter.
        merge: CoACD parameter.
        seed: CoACD parameter.
        verbose: If True, enables detailed logging in the subprocess.
        auto_scale: If True, automatically scale the output.

    Returns:
        The path to the output file if decomposition is successful.

    Raises:
        RuntimeError: If convex decomposition fails after all attempts.
    """
    params = dict(
        threshold=threshold,
        max_convex_hull=max_convex_hull,
        preprocess_mode=preprocess_mode,
        preprocess_resolution=preprocess_resolution,
        resolution=resolution,
        mcts_nodes=mcts_nodes,
        mcts_iterations=mcts_iterations,
        mcts_max_depth=mcts_max_depth,
        pca=pca,
        merge=merge,
        seed=seed,
    )

    ctx = mp.get_context("spawn")
    p = ctx.Process(
        target=decompose_convex_coacd,
        args=(filename, outfile, params, verbose, auto_scale),
    )
    p.start()
    p.join()
    if p.exitcode == 0 and os.path.exists(outfile):
        return outfile

    if preprocess_mode != "on":
        params["preprocess_mode"] = "on"
        p = ctx.Process(
            target=decompose_convex_coacd,
            args=(filename, outfile, params, verbose, auto_scale),
        )
        p.start()
        p.join()
        if p.exitcode == 0 and os.path.exists(outfile):
            return outfile

    raise RuntimeError(f"Convex decomposition failed on {filename}")