Use trimesh instead of bpy (#35)

---------

Co-authored-by: Nathan Dunkelberger <nathan.b.dunkelberger@nasa.gov>

Author: saikishor

package.xml

@@ -26,6 +26,9 @@
   <test_depend>urdfdom_py</test_depend>
   <test_depend>xacro</test_depend>
 
+  <exec_depend>python3-collada-pip</exec_depend>
+  <exec_depend>python3-scipy</exec_depend>
+
   <export>
     <build_type>ament_cmake</build_type>
   </export>

scripts/make_mjcf_from_robot_description.py

@@ -18,7 +18,6 @@
 # under the License.
 
 import argparse
-import bpy
 import mujoco
 import os
 import pathlib
@@ -32,6 +31,7 @@
 import math
 
 import numpy as np
+import trimesh  # Added trimesh
 
 from urdf_parser_py.urdf import URDF
 
@@ -52,11 +52,11 @@ def add_mujoco_info(raw_xml, output_filepath, publish_topic):
 
     # Use relative path for fixed directory otherwise use absolute path
     if not publish_topic:
-        meshdir = "assets"
+        asset_dir = "assets"
     else:
-        meshdir = os.path.join(output_filepath, "assets")
+        asset_dir = os.path.join(output_filepath, "assets")
 
-    compiler_element.setAttribute("meshdir", meshdir)
+    compiler_element.setAttribute("assetdir", asset_dir)
     compiler_element.setAttribute("balanceinertia", "true")
     compiler_element.setAttribute("discardvisual", "false")
     compiler_element.setAttribute("strippath", "false")
@@ -191,9 +191,79 @@ def replace_package_names(xml_data):
     return xml_data
 
 
+# get required images from dae so that we can copy them to the temporary filepath
+def get_images_from_dae(dae_path):
+
+    dae_dir = os.path.dirname(dae_path)
+
+    doc = minidom.parse(dae_path)
+
+    image_paths = []
+    seen = set()
+
+    # access data from dae files with this structure to access image_filepath
+    # <library_images>
+    #     <image id="id" name="name">
+    #         <init_from>image_filepath</init_from>
+    #     </image>
+    # </library_images>
+    for image in doc.getElementsByTagName("image"):
+        init_from_elems = image.getElementsByTagName("init_from")
+        if not init_from_elems:
+            continue
+
+        path = init_from_elems[0].firstChild
+        if path is None:
+            continue
+
+        image_path = path.nodeValue.strip()
+
+        # Resolve relative paths
+        if not os.path.isabs(image_path):
+            image_path = os.path.normpath(os.path.join(dae_dir, image_path))
+
+        # make sure it is an image
+        if image_path.lower().endswith((".png", ".jpg", ".jpeg")):
+            # ignore duplucates
+            if image_path not in seen:
+                seen.add(image_path)
+                image_paths.append(image_path)
+
+    return image_paths
+
+
+# Change all files that match "material_{some_int}.{png, jpg, jpeg}"
+# in a specified directory to be "material_{modifier}_{some_int}.{png, jpg, jpeg}"
+# This is important because trimesh puts out materials that look like
+# material_{some_int}.{png, jpg, jpeg}, and they need to be indexed per item
+def rename_material_textures(dir_path, modifier):
+    dir_path = pathlib.Path(dir_path)
+
+    # regex to match files we want to modify
+    pattern = re.compile(r"^material_(\d+)\.(png|jpg|jpeg)$", re.IGNORECASE)
+
+    for path in dir_path.iterdir():
+        if not path.is_file():
+            continue
+
+        m = pattern.match(path.name)
+        if not m:
+            continue
+
+        # extract important components and reorder
+        index, ext = m.groups()
+        new_name = f"material_{modifier}_{index}.{ext}"
+        new_path = path.with_name(new_name)
+
+        print(f"{path.name} -> {new_name}")
+        path.rename(new_path)
+
+
 def convert_to_objs(mesh_info_dict, directory, xml_data, convert_stl_to_obj, decompose_dict):
     # keep track of files we have already processed so we don't do it again
     converted_filenames = []
+    # keep track of what material number we are on to get unique materials
+    mtl_num = 0
 
     # clean assets directory and remake required paths
     if os.path.exists(f"{directory}assets/"):
@@ -212,10 +282,6 @@ def convert_to_objs(mesh_info_dict, directory, xml_data, convert_stl_to_obj, dec
 
         print(f"processing {full_filepath}")
 
-        # Clear any existing objects in the scene
-        bpy.ops.object.select_all(action="SELECT")
-        bpy.ops.object.delete(use_global=False)
-
         # if we want to decompose the mesh, put it in decomposed filepath, otherwise put it in full
         if filename_no_ext in decompose_dict:
             assets_relative_filepath = f"{DECOMPOSED_PATH_NAME}/{filename_no_ext}/"
@@ -224,22 +290,31 @@ def convert_to_objs(mesh_info_dict, directory, xml_data, convert_stl_to_obj, dec
             assets_relative_filepath = f"{COMPOSED_PATH_NAME}/"
         assets_relative_filepath += filename_no_ext
 
+        output_path = f"{directory}assets/{assets_relative_filepath}.obj"
+
         # Import the .stl or .dae file
         if filename_ext.lower() == ".stl":
             if filename_no_ext in decompose_dict and not convert_stl_to_obj:
                 raise ValueError("The --convert_stl_to_obj argument must be specified to decompose .stl mesh")
             if convert_stl_to_obj:
-                bpy.ops.wm.stl_import(filepath=full_filepath)
+                # Load mesh using trimesh
+                mesh = trimesh.load(full_filepath)
 
                 # bring in file color from urdf
-                new_mat = bpy.data.materials.new(name=f"material_{filename_no_ext}")
-                new_mat.diffuse_color = mesh_item["color"]
-                o = bpy.context.selected_objects[0]
-                o.active_material = new_mat
-
-                bpy.ops.wm.obj_export(
-                    filepath=f"{directory}assets/{assets_relative_filepath}.obj", forward_axis="Y", up_axis="Z"
-                )
+                if "color" in mesh_item:
+                    # trimesh expects 0-255 uint8 for colors
+                    rgba = mesh_item["color"]
+                    # make a material for the rgba values and export it
+                    mtl_modifier = f"m{mtl_num}"
+                    mtl_name = "mtl_" + mtl_modifier
+                    material = trimesh.visual.material.SimpleMaterial(name=mtl_name, diffuse=rgba)  # RGBA
+                    mesh.visual = trimesh.visual.TextureVisuals(material=material)
+
+                    # increment material number
+                    mtl_num = mtl_num + 1
+
+                # Export to OBJ
+                mesh.export(output_path, include_color=True, mtl_name=mtl_name)
                 xml_data = xml_data.replace(full_filepath, f"{assets_relative_filepath}.obj")
 
             else:
@@ -258,24 +333,66 @@ def convert_to_objs(mesh_info_dict, directory, xml_data, convert_stl_to_obj, dec
             pass
             # objs are ok as is
         elif filename_ext.lower() == ".dae":
+            # keep track of the image files that we need to copy from the dae
+            image_files = get_images_from_dae(full_filepath)
+            # keep track of the files that were copied to tmp directory to delete
+            copied_image_files = []
+
             # set z axis to up in the dae file because that is how mujoco expects it
             z_up_dae_txt = set_up_axis_to_z_up(full_filepath)
 
             # make a temporary file rather than overwriting the old one
-            temp_file = tempfile.NamedTemporaryFile(suffix=".dae", delete=False)
+            temp_file = tempfile.NamedTemporaryFile(suffix=".dae", mode="w+", delete=False)
             temp_filepath = temp_file.name
+            temp_folder = os.path.dirname(temp_filepath)
             try:
-                with open(temp_filepath, "w") as f:
-                    f.write(z_up_dae_txt)
-                    # import into blender
-                    bpy.ops.wm.collada_import(filepath=temp_filepath)
-            finally:
+                temp_file.write(z_up_dae_txt)
                 temp_file.close()
-                os.remove(temp_filepath)
 
-            bpy.ops.wm.obj_export(
-                filepath=f"{directory}assets/{assets_relative_filepath}.obj", forward_axis="Y", up_axis="Z"
-            )
+                # copy relevant images into the temporary directory
+                for image_file in image_files:
+                    shutil.copy2(image_file, temp_folder)
+                    copied_image_files.append(f"{temp_folder}/{os.path.basename(image_file)}")
+
+                # Load scene using trimesh
+                scene = trimesh.load(temp_filepath, force=trimesh.scene)
+
+                # give the material a unique name so that it can be properly referenced
+                mtl_modifier = f"m{mtl_num}"
+                mtl_name = "mtl_" + mtl_modifier
+                mtl_filepath = os.path.dirname(output_path) + f"/{mtl_name}"
+
+                scene.export(output_path, include_color=True, mtl_name=mtl_name)
+                # rename textures
+                rename_material_textures(dir_path=os.path.dirname(output_path), modifier=mtl_modifier)
+
+                # we need to modify the material names to not all be material_X so they don't conflict
+                # all of the objs will have a line that looks like
+                #   usemtl material_0
+                # and all of the mtl files will have a line that looks like
+                #   newmtl material_0
+                # We will modify both of them to be numberd by mtl_num like so
+                #   usemtl material_{mtl_num}_0
+                #   newmtl material_{mtl_num}_0
+
+                if os.path.exists(mtl_filepath):
+                    for filepath in [mtl_filepath, output_path]:
+                        with open(filepath) as f:
+                            data = f.read()
+                        data = data.replace("material_", f"material_{mtl_modifier}_")
+                        with open(filepath, "w") as f:
+                            f.write(data)
+                # increment
+                mtl_num = mtl_num + 1
+            finally:
+                if os.path.exists(temp_filepath):
+                    os.remove(temp_filepath)
+
+            # get rid of copied image files in the temporary file directory
+            for copied_image_file in copied_image_files:
+                if os.path.exists(copied_image_file):
+                    os.remove(copied_image_file)
+
             xml_data = xml_data.replace(full_filepath, f"{assets_relative_filepath}.obj")
         else:
             print(f"Can't convert {full_filepath} \n\tOnly stl and dae file extensions are supported at the moment")
@@ -446,6 +563,17 @@ def update_obj_assets(dom, output_filepath, mesh_info_dict):
             for sub_material in sub_materials:
                 asset_element.appendChild(sub_material)
 
+            # bring in the textures, and modify filepath to properly reference filepaths
+            sub_textures = sub_asset_element.getElementsByTagName("texture")
+            for sub_texture in sub_textures:
+                if sub_texture.hasAttribute("file"):
+                    sub_texture_file = sub_texture.getAttribute("file")
+                    if composed_type == DECOMPOSED_PATH_NAME:
+                        sub_texture.setAttribute("file", f"{composed_type}/{mesh_name}/{mesh_name}/{sub_texture_file}")
+                    else:
+                        sub_texture.setAttribute("file", f"{composed_type}/{mesh_name}/{sub_texture_file}")
+                    asset_element.appendChild(sub_texture)
+
             sub_body = sub_dom.getElementsByTagName("body")
             sub_body = sub_body[0]