Add fallback conversion to obj for unsupported mesh files

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

Author: shameekganguly

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

@@ -16,13 +16,70 @@
 
 import os
 
+from sdformat_mjcf.sdformat_to_mjcf.mesh_io import convert_mesh_to_obj_multimesh
+from sdformat_mjcf.sdformat_to_mjcf.converters.material import add_material
 import sdformat_mjcf.utils.sdf_utils as su
 
 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 add_geometry(body, name, pose, sdf_geom):
+def _validate_uri(uri):
+    if 'http://' in uri or 'https://' in uri:
+        raise RuntimeError("Fuel meshes are not yet supported")
+    if 'model://' in uri:
+            prefix = 'model://'
+            # TODO: Support sdf::ParserConfig::AddURIPath to resolve URIs.
+            return uri.replace(prefix, '', 1)
+    return uri
+
+def _is_unsupported_mesh_geo(root, sdf_geom):
+    if not sdf_geom.mesh_shape():
+        return False
+    mesh_shape = sdf_geom.mesh_shape()
+    uri = _validate_uri(mesh_shape.uri())
+    _, basename = os.path.split(uri)
+    filename, extension = os.path.splitext(basename)
+    # Mujoco supports .obj, .stl and .msh
+    # Out of these, .obj is only supported if it has a single mesh in it.
+    # So treat it as unsupported for this check so these files are sanitized
+    # 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.
 
@@ -74,22 +131,40 @@ def add_geometry(body, name, pose, sdf_geom):
         geom.size = [sphere_shape.radius()]
     elif sdf_geom.mesh_shape():
         mesh_shape = sdf_geom.mesh_shape()
-        uri = mesh_shape.uri()
+        uri = _validate_uri(mesh_shape.uri())
+        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(".")
         if (len(extension_tokens) == 1):
             raise RuntimeError("Unable to find the mesh extension {}"
                                .format(uri))
-        file_without_extension = os.path.splitext(
-            os.path.basename(mesh_shape.uri()))[0]
-        if 'http://' in uri or 'https://' in uri:
-            raise RuntimeError("Fuel meshes are not yet supported")
+        file_without_extension = os.path.splitext(os.path.basename(uri))[0]
         geom.type = "mesh"
         asset_loaded = geom.root.asset.find('mesh', file_without_extension)
-        dirname = os.path.dirname(mesh_shape.file_path())
-        mesh_file_path = os.path.join(dirname, uri)
         if asset_loaded is None:
-            geom.mesh = geom.root.asset.add('mesh',
-                                            file=mesh_file_path)
+            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"
         else:
             geom.mesh = asset_loaded
         geom.mesh.scale = su.vec3d_to_list(mesh_shape.scale())
@@ -99,6 +174,101 @@ def add_geometry(body, name, pose, sdf_geom):
 
     return geom
 
+def _add_mesh_geom_with_assets(body, name, pose, mjcf_mesh_asset,
+                               is_visual=False, 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:
+    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)
+    if mesh_loaded:
+        return _add_mesh_geom_with_assets(body, name, pose, mesh_loaded,
+                                          is_visual,
+                                          mjcf_material_asset=material_loaded)
+
+    # Try converting mesh to sanitized .obj. This could result in multiple
+    # geos, one per mesh in the input file.
+    dirname = os.path.dirname(sdf_mesh.file_path())
+    mesh_file_path = os.path.join(dirname, uri)
+    # Pass a nominal path to `convert_mesh_to_obj_multimesh`. If multiple
+    # sub-meshes are present, only the file name without extension from this
+    # nominal path will be used as a prefix for the output files.
+    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)
+    geom_list = []
+    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',
+                                               output_file_without_extension)
+        sub_mesh_material_asset_name = (
+            "material_" + output_file_without_extension)
+        material_loaded = body.root.asset.find('mesh',
+                                               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))
+
+
+        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):
     """
@@ -142,14 +312,21 @@ def add_visual(body, vis):
     :param mjcf.Element body: The MJCF body to which the geom is added.
     :param sdformat.Visual vis: Visual object to be converted.
     :return: The newly created MJCF geom.
-    :rtype: mjcf.Element
+    :rtype: mjcf.Element or list of mjcf
     """
     sem_pose = vis.semantic_pose()
     pose = su.graph_resolver.resolve_pose(sem_pose)
-    geom = add_geometry(body, vis.name(), pose, vis.geometry())
-    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
-    return geom
+    geoms = add_geometry(body, vis.name(), pose, vis.geometry(), is_visual=True)
+    if not isinstance(geoms, list):
+        geoms = [geoms]
+    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())
+        for geom in geoms:
+            geom.material = mjcf_mat
+    return geoms

sdformat_mjcf/src/sdformat_mjcf/sdformat_to_mjcf/converters/link.py

@@ -17,7 +17,6 @@
     add_visual,
 )
 from sdformat_mjcf.sdformat_to_mjcf.converters.light import add_light
-from sdformat_mjcf.sdformat_to_mjcf.converters.material import add_material
 from sdformat_mjcf.sdformat_to_mjcf.converters.sensor import add_sensor
 import sdformat_mjcf.utils.sdf_utils as su
 
@@ -91,9 +90,7 @@ def add_link(body, link, parent_name="world", link_pose=None):
     for vi in range(link.visual_count()):
         vis = link.visual_by_index(vi)
         if vis.geometry() is not None:
-            visual_geom = add_visual(body, vis)
-            if vis.material() is not None:
-                add_material(visual_geom, vis.material())
+            add_visual(body, vis)
 
     for li in range(link.light_count()):
         light = link.light_by_index(li)

sdformat_mjcf/src/sdformat_mjcf/sdformat_to_mjcf/converters/material.py

@@ -24,11 +24,12 @@
 MATERIAL_NUMBER = 0
 
 
-def add_material(geom, material):
+def add_material(root, material):
     """
     Converts an SDFormat material to an MJCF material.
 
-    :param mjcf.Element geom: The MJCF geom to add the material
+    :param mjcf.RootElement root: The MJCF root in which the material will be
+    created.
     :param sdf.Material material: The SDF material to convert
     :return: The newly created MJCF material.
     :rtype: mjcf.Element
@@ -38,7 +39,7 @@ def add_material(geom, material):
     em_color = material.emissive()
     specular = (sp_color.r() + sp_color.g() + sp_color.b()) / 3.0
     emissive = (em_color.r() + em_color.g() + em_color.b()) / 3.0
-    asset = geom.root.asset
+    asset = root.asset
     r_mat = None
     if pbr is not None:
         workflow = pbr.workflow(sdf.PbrWorkflowType.METAL)
@@ -81,5 +82,4 @@ def add_material(geom, material):
                                 clamp(diff.g() * 0.8 + amb.g() * 0.4, 0, 1),
                                 clamp(diff.b() * 0.8 + amb.b() * 0.4, 0, 1),
                                 clamp(diff.a() * 0.8 + amb.a() * 0.4, 0, 1)])
-    geom.material = r_mat
     return r_mat