Merge pull request #288 from Genesis-Embodied-AI/zhenjia/load_mjcf_bo…

…x_texture

load box texture when parsing mjcf

genesis/utils/mjcf.py

@@ -45,7 +45,6 @@ def parse_mjcf(path):
 
 
 def parse_link(mj, i_l, q_offset, dof_offset, scale):
-
     # mj.body
     l_info = dict()
 
@@ -286,6 +285,32 @@ def parse_geom(mj, i_g, scale, convexify, surface, xml_path):
         tmesh = trimesh.creation.box(extents=mj.geom_size[i_g, :3] * 2)
         gs_type = gs.GEOM_TYPE.BOX
 
+        # TODO: not sure if it is the right way to load texture for box
+        mat_id = mj.geom_matid[i_g]
+        if mat_id >= 0:
+            mat_id = mj.geom_matid[i_g]
+            tex_id = next((x for x in mj.mat_texid[mat_id] if x != -1), None)
+
+            if tex_id is not None:
+                tex_path = mj.paths[mj.tex_pathadr[tex_id] :].decode("utf-8").split("\x00")[0]
+                texturedir = extract_compiler_attributes(xml_path)["texturedir"]
+                assets_dir = os.path.join(get_assets_dir(), os.path.join(os.path.dirname(xml_path), texturedir))
+
+                uv_coordinates = tmesh.vertices[:, :2].copy()
+                uv_coordinates -= uv_coordinates.min(axis=0)
+                uv_coordinates /= uv_coordinates.max(axis=0)
+                image = Image.open(os.path.join(assets_dir, tex_path))
+                image_array = np.array(image)
+                if image_array.ndim == 2:  # convert gray image to RGBA
+                    rgba = np.zeros((image_array.shape[0], image_array.shape[1], 4), dtype=np.uint8)
+                    rgba[:, :, :3] = image_array[:, :, None]
+                    rgba[:, :, 3] = 255
+                    image_array = rgba
+                tex_repeat = mj.mat_texrepeat[mat_id].astype(int)
+                image_array = np.tile(image_array, (tex_repeat[0], tex_repeat[1], 1))
+                visual = TextureVisuals(uv=uv_coordinates, image=Image.fromarray(image_array))
+                tmesh.visual = visual
+
     elif mj_type == mujoco.mjtGeom.mjGEOM_MESH:
         i = mj.geom_dataid[i_g]
         last = (i + 1) >= mj.nmesh