Adding meshes and materials

python/mujoco/usd_component.py

@@ -1,7 +1,8 @@
+import os
 from enum import Enum
 
 import mujoco
-from pxr import Usd, UsdGeom, UsdLux, Vt, Gf
+from pxr import Usd, UsdGeom, UsdLux, UsdShade, Vt, Gf, Sdf
 from scipy.spatial.transform import Rotation as R
 
 def create_usd_geom_primitive(geom, stage):
@@ -146,7 +147,10 @@ class USDMesh(USDGeom):
                stage,
                model,
                mesh_vertex_ranges,
-               mesh_face_ranges):
+               mesh_face_ranges,
+               mesh_texcoord_ranges,
+               mesh_facetexcoord_ranges,
+               texture_file):
     super().__init__(geom, stage)
     self.type = 7
     USDMesh.mesh_count += 1
@@ -164,12 +168,44 @@ class USDMesh(USDGeom):
     self.faces = model.mesh_face[mesh_face_ranges[mesh_idx]:mesh_face_ranges[mesh_idx+1]]
     self.prim.GetFaceVertexIndicesAttr().Set(self.faces)
 
+    texid = geom.texid
+    texcoords = model.mesh_texcoord[mesh_texcoord_ranges[texid]:mesh_texcoord_ranges[texid+1]]
+
+    facetexcoords = model.mesh_facetexcoord.flatten()
+    facetexcoords = facetexcoords[mesh_facetexcoord_ranges[mesh_idx]:mesh_facetexcoord_ranges[mesh_idx+1]]
+    self.texcoords = UsdGeom.PrimvarsAPI(self.prim).CreatePrimvar("st",
+                                    Sdf.ValueTypeNames.TexCoord2fArray,
+                                    UsdGeom.Tokens.faceVarying)
+
+    self.texcoords.Set(texcoords)
+    self.texcoords.SetIndices(Vt.IntArray(facetexcoords.tolist()));
+
+    mtl_path = Sdf.Path(f"/World/Looks/Material_{os.path.splitext(os.path.basename(texture_file))[0]}")
+    mtl = UsdShade.Material.Define(stage, mtl_path)
+    shader = UsdShade.Shader.Define(stage, mtl_path.AppendPath("Shader"))
+    shader.CreateIdAttr("UsdPreviewSurface")
+    shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set((1.0, 0.0, 0.0))
+    shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
+    shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
+
+    diffuse_tx = UsdShade.Shader.Define(stage,mtl_path.AppendPath("DiffuseColorTx"))
+    diffuse_tx.CreateIdAttr('UsdUVTexture') 
+
+    # TODO: don't hardcode the image file
+    diffuse_tx.CreateInput('file', Sdf.ValueTypeNames.Asset).Set(texture_file)
+    diffuse_tx.CreateOutput('rgb', Sdf.ValueTypeNames.Float3)
+    shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).ConnectToSource(diffuse_tx.ConnectableAPI(), 'rgb')
+    mtl.CreateSurfaceOutput().ConnectToSource(shader.ConnectableAPI(), "surface")
+
+    self.prim.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
+    UsdShade.MaterialBindingAPI(self.prim).Bind(mtl)
+
   def update_geom(self, new_geom):
     self.update_pos(new_geom.pos)
     self.update_rotation(new_geom.mat)
 
     # TODO: remove this, temporary
-    self.xform.AddScaleOp().Set(value=(10.0, 10.0, 10.0))
+    # self.xform.AddScaleOp().Set(value=(10.0, 10.0, 10.0))
   
 class USDLight(object):
   """

python/mujoco/usd_renderer.py

@@ -1,10 +1,12 @@
-
+import os
 import mujoco
 import mujoco.viewer as viewer
 from mujoco.usd_component import *
 from mujoco.usd_utilities import *
 from pxr import Usd, UsdGeom
 
+from PIL import Image as im
+
 class USDRenderer(object):
   """
   Renderer class the creates USD representations for mujoco scenes
@@ -20,6 +22,7 @@ class USDRenderer(object):
     self.loaded_scene_info = False
 
     self.stage = Usd.Stage.CreateNew('usd_stage.usda')
+    UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.z)
 
   @property
   def usd(self):
@@ -52,9 +55,27 @@ class USDRenderer(object):
     if self.model.nmesh > 0:
       mesh_vertex_ranges = get_mesh_ranges(self.model.nmesh, self.model.mesh_vertnum)
       mesh_face_ranges = get_mesh_ranges(self.model.nmesh, self.model.mesh_facenum)
+      mesh_texcoord_ranges = get_mesh_ranges(self.model.nmesh, self.model.mesh_texcoordnum)
+      mesh_facetexcoord_ranges = get_facetexcoord_ranges(self.model.nmesh, self.model.mesh_facenum)
+
+
+    # create and load the texture files
+    # iterate through all the textures and build list of tex_rgb ranges
+    data_adr = 0
+    texture_files = []
+    for texid in range(self.model.ntex):
+      height = self.model.tex_height[texid]
+      width = self.model.tex_width[texid]
+      pixels = 3*height*width
+      rgb = self.model.tex_rgb[data_adr:data_adr+pixels]
+      img = rgb.reshape(height, width, 3)
+      file_name = f'{texid}.png'
+      im.fromarray(img).save(file_name)
+      texture_file = os.path.abspath(file_name)
+      texture_files.append(texture_file)
+      data_adr += pixels
 
     current_mesh_idx = 0
-
     # initializes an array to store all the geoms in the scene
     # populates with "empty" USDGeom objects
     self.usd_geoms = []
@@ -67,7 +88,10 @@ class USDRenderer(object):
                                       self.stage, 
                                       self.model,
                                       mesh_vertex_ranges,
-                                      mesh_face_ranges))
+                                      mesh_face_ranges,
+                                      mesh_texcoord_ranges,
+                                      mesh_facetexcoord_ranges,
+                                      texture_files[geoms[i].texid]))
         current_mesh_idx += 1
       else:
         self.usd_geoms.append(create_usd_geom_primitive(geoms[i], self.stage))
@@ -110,7 +134,7 @@ class USDRenderer(object):
 
   def start_viewer(self):
     if self.data:
-      viewer.launch(self.model, self.data)
+      viewer.launch(self.model)
 
   def render(self):
     # should render the usd file given a particular renderer that

python/mujoco/usd_utilities.py

@@ -5,4 +5,12 @@ def get_mesh_ranges(nmesh, arr):
     for i in range(nmesh):
         running_sum += arr[i]
         mesh_ranges.append(running_sum)
-    return mesh_ranges
+    return mesh_ranges
+
+def get_facetexcoord_ranges(nmesh, arr):
+    facetexcoords_ranges = [0]
+    running_sum = 0
+    for i in range(nmesh):
+        running_sum += arr[i] * 3
+        facetexcoords_ranges.append(running_sum)
+    return facetexcoords_ranges