Bug fixes

Signed-off-by: Shameek Ganguly <shameek@intrinsic.ai>

Author: shameekganguly

sdformat_mjcf/src/sdformat_mjcf/sdformat_to_mjcf/converters/geometry.py

@@ -23,39 +23,16 @@
 COLLISION_GEOM_GROUP = 3
 VISUAL_GEOM_GROUP = 0
 
-def _add_geoms_after_conversion(body, name, pose, conversion_output,
-                                is_visual):
-    geom_list = []
-    for path, info in conversion_output.obj_files.items():
-        geom = body.add(
-            "geom",
-            name=su.find_unique_name(body, "geom", name),
-            pos=su.vec3d_to_list(pose.pos()),
-            euler=su.quat_to_euler_list(pose.rot()),
-        )
-        geom.type = "mesh"
-        file_without_extension = os.path.splitext(
-            os.path.basename(path))[0]
-        asset_loaded = geom.root.asset.find('mesh', file_without_extension)
-        if asset_loaded is None:
-            geom.mesh = geom.root.asset.add('mesh', file=path)
-            if is_visual:
-                # Some visual meshes are too thin and can cause inertia
-                # calculation in Mujoco to fail. Mark them as "shell" to prevent
-                # Mujoco from failing in compilation. The visual geometry will
-                # not be used for inertia computation anyway since we explicitly
-                # convert inertia from sdformat to mjcf.
-                # https://github.com/google-deepmind/mujoco/issues/2455
-                geom.mesh.inertia = "shell"
-                mat_asset = geom.root.asset.add("material",
-                          name="material_" + file_without_extension,
-                          specular=info.mat.specular,
-                          shininess=info.mat.shininess,
-                          rgba=info.mat.rgba)
-                geom.material = mat_asset
-        else:
-            geom.mesh = asset_loaded
-    return geom_list
+def _set_mesh_inertia(mjcf_mesh, is_visual):
+    if is_visual:
+        # Some visual meshes are too thin and can cause inertia
+        # calculation in Mujoco to fail. Mark them as "shell" to prevent
+        # Mujoco from failing in compilation. The visual geometry will
+        # not be used for inertia computation anyway since we explicitly
+        # convert inertia from sdformat to mjcf.
+        # https://github.com/google-deepmind/mujoco/issues/2455
+        mjcf_mesh.inertia = "shell"
+
 
 def _validate_uri(uri):
     if 'http://' in uri or 'https://' in uri:
@@ -66,7 +43,8 @@ def _validate_uri(uri):
             return uri.replace(prefix, '', 1)
     return uri
 
-def _is_unsupported_mesh_geo(root, sdf_geom):
+
+def _is_unsupported_mesh_geo(sdf_geom):
     if not sdf_geom.mesh_shape():
         return False
     mesh_shape = sdf_geom.mesh_shape()
@@ -79,6 +57,7 @@ def _is_unsupported_mesh_geo(root, sdf_geom):
     # for Mujoco.
     return extension not in [".stl", ".msh"]
 
+
 def add_geometry(body, name, pose, sdf_geom, is_visual=False):
     """
     Converts an SDFormat geometry to an MJCF geom and add it to the given body.
@@ -135,7 +114,7 @@ def add_geometry(body, name, pose, sdf_geom, is_visual=False):
         if _is_unsupported_mesh_geo(sdf_geom):
             raise RuntimeError(
                 f"Call `convert_and_add_mesh` for unsupported mesh geo {uri}")
-        extension_tokens = os.path.basename(mesh_shape.uri()).split(".")
+        extension_tokens = os.path.basename(uri).split(".")
         if (len(extension_tokens) == 1):
             raise RuntimeError("Unable to find the mesh extension {}"
                                .format(uri))
@@ -145,26 +124,9 @@ def add_geometry(body, name, pose, sdf_geom, is_visual=False):
         if asset_loaded is None:
             dirname = os.path.dirname(mesh_shape.file_path())
             mesh_file_path = os.path.join(dirname, uri)
-            try:
-                geom.mesh = geom.root.asset.add('mesh',
-                                                file=mesh_file_path)
-            except ValueError as e:
-                # Try converting mesh to .obj. This could result in multiple
-                # geos, one per mesh in the input file.
-                converted_path = os.path.join(dirname,
-                                              file_without_extension + ".obj")
-                print(f"Converting {mesh_file_path} to {converted_path}")
-                result = convert_mesh_to_obj_multimesh(mesh_file_path,
-                                                       converted_path)
-                if isinstance(result, str):
-                    geom.mesh = geom.root.asset.add('mesh',
-                                                file=result)
-                else:
-                    geom.remove()
-                    return _add_geoms_after_conversion(body, name, pose,
-                                                       result, is_visual)
-            if is_visual:
-                geom.mesh.inertia = "shell"
+            geom.mesh = geom.root.asset.add('mesh',
+                                            file=mesh_file_path)
+            _set_mesh_inertia(geom.mesh, is_visual)
         else:
             geom.mesh = asset_loaded
         geom.mesh.scale = su.vec3d_to_list(mesh_shape.scale())
@@ -174,41 +136,36 @@ def add_geometry(body, name, pose, sdf_geom, is_visual=False):
 
     return geom
 
+
 def _add_mesh_geom_with_assets(body, name, pose, mjcf_mesh_asset,
-                               is_visual=False, mjcf_material_asset=None):
+                               mjcf_material_asset=None):
     geom = body.add(
             "geom",
             name=su.find_unique_name(body, "geom", name),
             pos=su.vec3d_to_list(pose.pos()),
             euler=su.quat_to_euler_list(pose.rot()),
         )
     geom.type = "mesh"
-    if is_visual:
-        # Some visual meshes are too thin and can cause inertia
-        # calculation in Mujoco to fail. Mark them as "shell" to prevent
-        # Mujoco from failing in compilation. The visual geometry will
-        # not be used for inertia computation anyway since we explicitly
-        # convert inertia from sdformat to mjcf.
-        # https://github.com/google-deepmind/mujoco/issues/2455
-        geom.mesh.inertia = "shell"
     geom.mesh = mjcf_mesh_asset
     if mjcf_material_asset:
         geom.material = mjcf_material_asset
     return geom
 
+
 def convert_and_add_mesh(body, name, pose, sdf_mesh, is_visual=False):
     # Check if asset was loaded already with the uri key. If so, just add the
-    # asset:
+    # asset. This can happen if the converted mesh has a single sub-mesh,
+    # which was loaded already.
     uri = _validate_uri(sdf_mesh.uri())
     file_without_extension = os.path.splitext(os.path.basename(uri))[0]
     mesh_asset_name = file_without_extension
     mesh_loaded = body.root.asset.find('mesh', mesh_asset_name)
     material_asset_name = "material_" + file_without_extension
-    material_loaded = body.root.asset.find('mesh', material_asset_name)
+    material_loaded = body.root.asset.find('material', material_asset_name)
     if mesh_loaded:
-        return _add_mesh_geom_with_assets(body, name, pose, mesh_loaded,
-                                          is_visual,
+        geom = _add_mesh_geom_with_assets(body, name, pose, mesh_loaded,
                                           mjcf_material_asset=material_loaded)
+        return [geom]
 
     # Try converting mesh to sanitized .obj. This could result in multiple
     # geos, one per mesh in the input file.
@@ -224,50 +181,34 @@ def convert_and_add_mesh(body, name, pose, sdf_mesh, is_visual=False):
     for path, info in result.obj_files.items():
         output_file_without_extension = os.path.splitext(
             os.path.basename(path))[0]
-        sub_mesh_loaded = geom.root.asset.find('mesh',
+        sub_mesh_loaded = body.root.asset.find('mesh',
                                                output_file_without_extension)
         sub_mesh_material_asset_name = (
             "material_" + output_file_without_extension)
-        material_loaded = body.root.asset.find('mesh',
+        material_loaded = body.root.asset.find('material',
                                                sub_mesh_material_asset_name)
         if sub_mesh_loaded:
             geom_list.append(
                 _add_mesh_geom_with_assets(body, name, pose, mesh_loaded,
-                                          is_visual,
                                           mjcf_material_asset=material_loaded))
+            continue
 
+        # Add mesh and material assets
+        mesh = body.root.asset.add('mesh', file=path)
+        _set_mesh_inertia(mesh, is_visual)
+        mesh.scale = su.vec3d_to_list(sdf_mesh.scale())
+        material_asset = None
+        if is_visual:
+            material_asset = body.root.asset.add("material",
+                            name=sub_mesh_material_asset_name,
+                            specular=info.mat.specular,
+                            shininess=info.mat.shininess,
+                            rgba=info.mat.rgba)
+        geom_list.append(
+                _add_mesh_geom_with_assets(body, name, pose, mesh_loaded,
+                                          mjcf_material_asset=material_asset))
 
-        geom = body.add(
-            "geom",
-            name=su.find_unique_name(body, "geom", name),
-            pos=su.vec3d_to_list(pose.pos()),
-            euler=su.quat_to_euler_list(pose.rot()),
-        )
-        geom.type = "mesh"
-        file_without_extension = os.path.splitext(
-            os.path.basename(path))[0]
-        asset_loaded = geom.root.asset.find('mesh', file_without_extension)
-        if asset_loaded is None:
-            geom.mesh = geom.root.asset.add('mesh', file=path)
-            if is_visual:
-                # Some visual meshes are too thin and can cause inertia
-                # calculation in Mujoco to fail. Mark them as "shell" to prevent
-                # Mujoco from failing in compilation. The visual geometry will
-                # not be used for inertia computation anyway since we explicitly
-                # convert inertia from sdformat to mjcf.
-                # https://github.com/google-deepmind/mujoco/issues/2455
-                geom.mesh.inertia = "shell"
-                mat_asset = geom.root.asset.add("material",
-                          name="material_" + file_without_extension,
-                          specular=info.mat.specular,
-                          shininess=info.mat.shininess,
-                          rgba=info.mat.rgba)
-                geom.material = mat_asset
-        else:
-            geom.mesh = asset_loaded
     return geom_list
-    return _add_geoms_after_conversion(body, name, pose,
-                                            result, is_visual)
 
 
 def apply_surface_to_geometry(geom, sdf_surface):
@@ -297,10 +238,20 @@ def add_collision(body, col):
     """
     sem_pose = col.semantic_pose()
     pose = su.graph_resolver.resolve_pose(sem_pose)
-    geom = add_geometry(body, col.name(), pose, col.geometry())
-    geom.group = COLLISION_GEOM_GROUP
-    apply_surface_to_geometry(geom, col.surface())
-    return geom
+    if _is_unsupported_mesh_geo(col.geometry()):
+        sdf_mesh = col.geometry().mesh_shape()
+        geoms = convert_and_add_mesh(body, col.name(), pose, sdf_mesh,
+            is_visual=False)
+    else:
+        geom = add_geometry(body, col.name(), pose, col.geometry(),
+            is_visual=False)
+        geoms = [geom]
+    for geom in geoms:
+        geom.group = COLLISION_GEOM_GROUP
+        apply_surface_to_geometry(geom, col.surface())
+    if len(geoms) == 1:
+        return geoms[0]
+    return geoms
 
 
 def add_visual(body, vis):
@@ -316,17 +267,25 @@ def add_visual(body, vis):
     """
     sem_pose = vis.semantic_pose()
     pose = su.graph_resolver.resolve_pose(sem_pose)
-    geoms = add_geometry(body, vis.name(), pose, vis.geometry(), is_visual=True)
-    if not isinstance(geoms, list):
-        geoms = [geoms]
+    if _is_unsupported_mesh_geo(vis.geometry()):
+        sdf_mesh = vis.geometry().mesh_shape()
+        geoms = convert_and_add_mesh(body, vis.name(), pose, sdf_mesh,
+            is_visual=True)
+    else:
+        geom = add_geometry(body, vis.name(), pose, vis.geometry(),
+            is_visual=True)
+        geoms = [geom]
+
     for geom in geoms:
         geom.group = VISUAL_GEOM_GROUP
         # Visual geoms do not collide with any other geom, so we set their
         # contype and conaffinity to 0.
         geom.contype = 0
         geom.conaffinity = 0
     if vis.material() is not None:
-        mjcf_mat = add_material(geom, vis.material())
+        mjcf_mat = add_material(geom.root, vis.material())
         for geom in geoms:
             geom.material = mjcf_mat
+    if len(geoms) == 1:
+        return geoms[0]
     return geoms

sdformat_mjcf/src/sdformat_mjcf/sdformat_to_mjcf/mesh_io.py

@@ -4,63 +4,18 @@
 import sys
 import trimesh
 
-def convert_mesh_to_obj(input_filepath, output_filepath):
-    """
-    Converts an input mesh file (e.g. glb) to an STL file.
-    Supported file formats: .dae, .glb, .glTF and others that Assimp supports.
-
-    Args:
-        input_filepath (str): Path to the input mesh file.
-        output_filepath (str): Path where the output .stl file will be saved.
-    """
-    if not os.path.exists(input_filepath):
-        raise RuntimeError("Unable to find the input mesh file {}"
-                               .format(input_filepath))
-    try:
-        # trimesh.load handles scenes (multiple meshes) or single meshes
-        model = trimesh.load(input_filepath)
-    except Exception as e:
-        print(f"\nERROR: Failed to load GLB file. Check if the file is valid.")
-        print(f"Detail: {e}")
-        sys.exit(1)
-
-    base_path, _ = os.path.splitext(output_filepath)
-    mtl_filepath = base_path + ".mtl"
-    mtl_filename = os.path.basename(mtl_filepath)
-
-    try:
-        if isinstance(model, trimesh.Trimesh):
-            # If it's a single mesh, wrap it in a scene for consistent export
-            scene = trimesh.Scene([model])
-        elif isinstance(model, trimesh.Scene):
-            scene = model
-        else:
-            print(f"\nERROR: Loaded model object type is unexpected: "
-                  f"{type(model)}. Conversion failed.")
-            sys.exit(1)
-
-        # Export the scene to OBJ
-        scene.export(
-            output_filepath,
-            file_type='obj',
-            mtl_name=mtl_filename
-        )
-    except Exception as e:
-        print(f"\nERROR: Failed to export model to OBJ.")
-        print(f"Detail: {e}")
-        sys.exit(1)
-
-
 @dataclass
 class Material:
     specular: float
     shininess: float
     rgba: list[float]  # each channel has a value from 0 to 1
 
+
 @dataclass
 class MeshInfo:
     mat: Material
 
+
 @dataclass
 class ConversionOutput:
     # Dictionary with key = obj filename for mesh, value = assigned material
@@ -78,6 +33,7 @@ def _extract_material_info(mesh_material) -> Material:
     default_rgba = [0.5, 0.5, 0.5, 1.0] # Default grey
 
     # Simple defaults for common material properties
+    # TODO: Make this configurable
     specular = 0.5
     shininess = 30.0
     color = default_rgba
@@ -86,15 +42,12 @@ def _extract_material_info(mesh_material) -> Material:
     if isinstance(mesh_material, trimesh.visual.material.PBRMaterial):
         # PBRMaterial uses baseColorFactor
         color = mesh_material.baseColorFactor
-        print("PBRMaterial", color)
     elif isinstance(mesh_material, trimesh.visual.material.SimpleMaterial):
         # SimpleMaterial usually has a diffuse color
         color = mesh_material.diffuse
-        print("SimpleMaterial", color)
     elif hasattr(mesh_material, 'main_color'):
         # Fallback for older/custom materials
         color = mesh_material.main_color
-        print("main_color", color)
 
     # --- Color Standardization ---
     if isinstance(color, np.ndarray):
@@ -196,13 +149,11 @@ def convert_mesh_to_obj_multimesh(input_filepath: str, output_filepath: str) ->
 
             if extracted_material is None:
                 extracted_material = Material(specular=0.5, shininess=30.0, rgba=[0.5, 0.5, 0.5, 1.0])
-            else:
-                print(extracted_material.rgba)
 
             mesh_info = MeshInfo(mat=extracted_material)
 
             # Store result
-            conversion_results[os.path.basename(obj_filename)] = mesh_info
+            conversion_results[obj_filename] = mesh_info
 
         except Exception as e:
             print(f"\nERROR: Failed to export mesh '{G_name}' to OBJ.")