Adding UV maps only when texture present

python/mujoco/usd/exporter.py

@@ -251,10 +251,8 @@ class USDExporter:
 
     assert geom_name not in self.geom_names
 
-    texture_file = (
+    geom_textures = (
-        self.texture_files[
+      [(self.texture_files[i], self.model.tex_type[i]) if i != -1 else None for i in self.model.mat_texid[geom.matid]]
-            self.model.mat_texid[geom.matid][mujoco.mjtTextureRole.mjTEXROLE_RGB]
-        ]
       if geom.matid != -1
       else None
     )
@@ -268,7 +266,7 @@ class USDExporter:
           obj_name=geom_name,
           dataid=self.model.geom_dataid[geom.objid],
           rgba=geom.rgba,
-          texture_file=texture_file,
+          geom_textures=geom_textures,
       )
     else:
       # handling tendons in our scene
@@ -286,7 +284,7 @@ class USDExporter:
             geom=geom,
             obj_name=geom_name,
             rgba=geom.rgba,
-            texture_file=texture_file,
+            geom_textures=geom_textures,
         )
       # handling primitives in our scene
       else:
@@ -302,7 +300,7 @@ class USDExporter:
             geom=geom,
             obj_name=geom_name,
             rgba=geom.rgba,
-            texture_file=texture_file,
+            geom_textures=geom_textures,
         )
 
     self.geom_names.add(geom_name)

python/mujoco/usd/objects.py

@@ -16,7 +16,7 @@
 
 import abc
 import collections
-from typing import Optional, Dict, Any
+from typing import Optional, Dict, Any, Tuple
 
 import mujoco
 import mujoco.usd.shapes as shapes_module
@@ -55,14 +55,14 @@ class USDObject(abc.ABC):
       geom: mujoco.MjvGeom,
       obj_name: str,
       rgba: np.ndarray = np.array([1, 1, 1, 1]),
-      texture_file: Optional[str] = None,
+      geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None
   ):
     self.stage = stage
     self.model = model
     self.geom = geom
     self.obj_name = obj_name
     self.rgba = rgba
-    self.texture_file = texture_file
+    self.geom_textures = geom_textures
 
     self.xform_path = f"/World/Mesh_Xform_{obj_name}"
     self.usd_xform = UsdGeom.Xform.Define(stage, self.xform_path)
@@ -120,7 +120,7 @@ class USDObject(abc.ABC):
     # setting the image texture attributes
     image_shader.CreateIdAttr("UsdUVTexture")
     image_shader.CreateInput("file", Sdf.ValueTypeNames.Asset).Set(
-        self.texture_file
+        self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB][0]
     )
     image_shader.CreateInput("sourceColorSpace", Sdf.ValueTypeNames.Token).Set(
         "sRGB"
@@ -225,9 +225,9 @@ class USDMesh(USDObject):
       obj_name: str,
       dataid: int,
       rgba: np.ndarray = np.array([1, 1, 1, 1]),
-      texture_file: Optional[str] = None,
+      geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None
   ):
-    super().__init__(stage, model, geom, obj_name, rgba, texture_file)
+    super().__init__(stage, model, geom, obj_name, rgba, geom_textures)
 
     self.dataid = dataid
 
@@ -243,6 +243,7 @@ class USDMesh(USDObject):
     )
     self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
 
+    if self.geom_textures and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]:
       # setting mesh uv properties
       mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry()
       self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar(
@@ -250,8 +251,6 @@ class USDMesh(USDObject):
       )
       self.texcoords.Set(mesh_texcoord)
       self.texcoords.SetIndices(Vt.IntArray(mesh_facetexcoord.tolist()))
-
-    if self.texture_file:
       self.attach_image_material(self.usd_mesh)
     else:
       self.attach_solid_material(self.usd_mesh)
@@ -320,9 +319,9 @@ class USDPrimitiveMesh(USDObject):
       geom: mujoco.MjvGeom,
       obj_name: str,
       rgba: np.ndarray = np.array([1, 1, 1, 1]),
-      texture_file: Optional[str] = None,
+      geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None
   ):
-    super().__init__(stage, model, geom, obj_name, rgba, texture_file)
+    super().__init__(stage, model, geom, obj_name, rgba, geom_textures)
 
     self.mesh_config = mesh_config
     self.prim_mesh = self.generate_primitive_mesh()
@@ -337,7 +336,9 @@ class USDPrimitiveMesh(USDObject):
         [3 for _ in range(mesh_facenum)]
     )
     self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
+    self._set_refinement_properties(self.usd_prim)
 
+    if self.geom_textures and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]:
       # setting mesh uv properties
       mesh_texcoord, _ = self._get_uv_geometry()
       self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar(
@@ -347,21 +348,20 @@ class USDPrimitiveMesh(USDObject):
       self.texcoords.Set(mesh_texcoord)
       self.texcoords.SetIndices(Vt.IntArray(list(range(mesh_facenum * 3))))
 
-    self._set_refinement_properties(self.usd_prim)
-
-    if self.texture_file:
       self.attach_image_material(self.usd_mesh)
     else:
       self.attach_solid_material(self.usd_mesh)
 
   def generate_primitive_mesh(self):
     """Generates the mesh for the primitive USD object."""
-    _, prim_mesh = shapes_module.mesh_factory(self.mesh_config)
+    geom_rgb_texture = self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB] if self.geom_textures else None
+    texture_type = geom_rgb_texture[1] if geom_rgb_texture else None
+    _, prim_mesh = shapes_module.mesh_factory(self.mesh_config, texture_type)
     prim_mesh.translate(-prim_mesh.get_center())
     return prim_mesh
 
   def _get_uv_geometry(self):
-    assert self.prim_mesh
+    assert self.prim_mesh and self.prim_mesh.triangle_uvs is not None
 
     s_scale, t_scale = self.model.mat_texrepeat[self.geom.matid]
 
@@ -400,9 +400,9 @@ class USDTendon(USDObject):
       geom: mujoco.MjvGeom,
       obj_name: str,
       rgba: np.ndarray = np.array([1, 1, 1, 1]),
-      texture_file: Optional[str] = None,
+      geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None
   ):
-    super().__init__(stage, model, geom, obj_name, rgba, texture_file)
+    super().__init__(stage, model, geom, obj_name, rgba, geom_textures)
 
     self.mesh_config = mesh_config
     self.tendon_parts = self.generate_primitive_mesh()
@@ -432,6 +432,7 @@ class USDTendon(USDObject):
           part_geometry["mesh_face"]
       )
 
+    if self.geom_textures and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]:
       # setting uv properties for each of the parts in the tendon
       part_uv_geometries = self._get_uv_geometry()
       for name, part_uv_geometry in part_uv_geometries.items():
@@ -446,19 +447,21 @@ class USDTendon(USDObject):
         self.texcoords.SetIndices(
             Vt.IntArray(list(range(part_geometry["mesh_facenum"] * 3)))
         )
-
         for _, ref in self.usd_refs.items():
           self._set_refinement_properties(ref["usd_prim"])
-      if self.texture_file:
           self.attach_image_material(ref["usd_mesh"])
     else:
+      for _, ref in self.usd_refs.items():
+          self._set_refinement_properties(ref["usd_prim"])
           self.attach_solid_material(ref["usd_mesh"])
 
   def generate_primitive_mesh(self):
     """Generates the tendon mesh using primitives."""
     mesh_parts = {}
+    geom_rgb_texture = self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB] if self.geom_textures else None
+    texture_type = geom_rgb_texture[1] if geom_rgb_texture else None
     for part_config in self.mesh_config:
-      mesh_name, prim_mesh = shapes_module.mesh_factory(part_config)
+      mesh_name, prim_mesh = shapes_module.mesh_factory(part_config, texture_type)
       prim_mesh.translate(-prim_mesh.get_center())
       mesh_parts[mesh_name] = prim_mesh
     return mesh_parts
@@ -466,6 +469,7 @@ class USDTendon(USDObject):
   def _get_uv_geometry(self):
     part_uv_geometries = collections.defaultdict(dict)
     for name, mesh in self.tendon_parts.items():
+      assert mesh.triangle_uvs is not None
       mesh_texcoord = np.array(mesh.triangle_uvs)
       mesh_facetexcoord = np.asarray(mesh.triangles)
       part_uv_geometries[name] = {

python/mujoco/usd/shapes.py

@@ -14,13 +14,20 @@
 # ==============================================================================
 """Built-in shapes for USD exporter."""
 
-from typing import Dict, Any, Tuple
+from typing import Dict, Any, Tuple, Optional
 
 import mujoco
 import numpy as np
 
-def get_triangle_uvs(vertices: np.array, triangles: np.array):
+def get_triangle_uvs(
-    # (jabhi) asusming all mappings are 2d mapping temporarily
+    vertices: np.array,
+    triangles: np.array,
+    texture_type: Optional[mujoco.mjtTexture]
+):
+  if texture_type == None:
+    return None
+
+  # (jabhi) assuming all mappings are 2d mapping temporarily
   triangle_uvs = np.array([
     [vertices[i][0], vertices[i][1]] for i in np.nditer(triangles)]
   )
@@ -39,7 +46,11 @@ class TriangleMesh():
 
   @classmethod
   def create_box(
-      cls, width: float, height: float, depth: float
+      cls,
+      width: float,
+      height: float,
+      depth: float,
+      texture_type: Optional[mujoco.mjtTexture]
   ):
     vertices = np.array([[0.0, 0.0, 0.0],
                          [width, 0.0, 0.0],
@@ -63,13 +74,16 @@ class TriangleMesh():
                           [0, 4, 1],
                           [1, 4, 5]])
     
-    triangle_uvs = get_triangle_uvs(vertices, triangles)
+    triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type)
     
     return TriangleMesh(vertices, triangles, triangle_uvs)
 
   @classmethod
   def create_sphere(
-      cls, radius: float, resolution: int
+      cls,
+      radius: float,
+      texture_type: Optional[mujoco.mjtTexture],
+      resolution: int
   ):
     vertices = []
     triangles = []
@@ -93,13 +107,16 @@ class TriangleMesh():
     vertices = np.array(vertices)
     triangles = np.array(triangles)
 
-    triangle_uvs = get_triangle_uvs(vertices, triangles)
+    triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type)
     
     return TriangleMesh(vertices, triangles, triangle_uvs)
 
   @classmethod
   def create_hemisphere(
-    cls, radius: float, resolution: int
+    cls,
+    radius: float,
+    texture_type: Optional[mujoco.mjtTexture],
+    resolution: int
   ):
     vertices = []
     triangles = []
@@ -128,13 +145,17 @@ class TriangleMesh():
     vertices = np.array(vertices)
     triangles = np.array(triangles)
 
-    triangle_uvs = get_triangle_uvs(vertices, triangles)
+    triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type)
     
     return TriangleMesh(vertices, triangles, triangle_uvs)
 
   @classmethod
   def create_cylinder(
-      cls, radius: float, height: float, resolution: int
+      cls,
+      radius: float,
+      height: float,
+      texture_type: Optional[mujoco.mjtTexture],
+      resolution: int
   ):
     vertices = []
     triangles = []
@@ -166,7 +187,7 @@ class TriangleMesh():
     vertices = np.array(vertices)
     triangles = np.array(triangles)
 
-    triangle_uvs = get_triangle_uvs(vertices, triangles)
+    triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type)
     
     return TriangleMesh(vertices, triangles, triangle_uvs)
 
@@ -189,7 +210,7 @@ class TriangleMesh():
       new_vertices = np.vstack((self.vertices, other.vertices))
       other_triangles = other.triangles + len(self.vertices)
       new_triangles = np.vstack((self.triangles, other_triangles))
-      new_triangle_uvs = get_triangle_uvs(new_vertices, new_triangles)
+      new_triangle_uvs = np.vstack((self.triangle_uvs, other.triangle_uvs))
       return TriangleMesh(new_vertices, new_triangles, new_triangle_uvs)
     raise TypeError(f"Cannot add TriangleMesh with {type(other)}")
 
@@ -282,7 +303,8 @@ def mesh_config_generator(
 
 def mesh_factory(
     mesh_config: Dict[str, Any],
-    resolution: int = 100,
+    texture_type: Optional[mujoco.mjtTexture],
+    resolution: int = 10,
 ):
   """Generates a mesh given a config consisting of shapes."""
   assert "name" in mesh_config
@@ -300,22 +322,26 @@ def mesh_factory(
       prim_mesh = TriangleMesh.create_box(
           width=mesh_config[shape]["width"],
           height=mesh_config[shape]["height"],
-          depth=mesh_config[shape]["depth"]
+          depth=mesh_config[shape]["depth"],
+          texture_type=texture_type
       )
     elif "hemisphere" in shape:
       prim_mesh = TriangleMesh.create_hemisphere(
           radius=mesh_config[shape]["radius"],
+          texture_type=texture_type,
           resolution=resolution
       )
     elif "sphere" in shape:
       prim_mesh = TriangleMesh.create_sphere(
           radius=mesh_config[shape]["radius"],
+          texture_type=texture_type,
           resolution=resolution
       )
     elif "cylinder" in shape:
       prim_mesh = TriangleMesh.create_cylinder(
           radius=mesh_config[shape]["radius"],
           height=mesh_config[shape]["height"],
+          texture_type=texture_type,
           resolution=resolution
       )
     else: