Adding logic to include meshes using MjModel

python/mujoco/usd_component.py

@@ -1,7 +1,17 @@
 from enum import Enum
 
 import mujoco
-from pxr import Usd, UsdGeom, Vt
+from pxr import Usd, UsdGeom, UsdLux, Vt, Gf
+from scipy.spatial.transform import Rotation as R
+
+def create_usd_geom_primitive(geom, stage):
+  geom_type = geom.type
+  if geom_type==USDGeomType.Plane.value:
+    return USDPlane(geom, stage)
+  elif geom_type==USDGeomType.Sphere.value:
+    return USDSphere(geom, stage)
+  elif geom_type==USDGeomType.Cube.value:
+    return USDCube(geom, stage)
 
 class USDGeomType(Enum):
   """
@@ -12,6 +22,7 @@ class USDGeomType(Enum):
   Plane = 0
   Sphere = 2
   Cube = 6
+  Mesh = 7
 
 class USDGeom(object):
   """
@@ -28,12 +39,22 @@ class USDGeom(object):
     self.ref = None
 
   def update_geom(self, new_geom):
-    raise NotImplementedError
+    self.update_pos(new_geom.pos)
+    self.update_rotation(new_geom.mat)
+    self.update_size(new_geom.size)
+    self.update_color(new_geom.rgba)
   
   def update_pos(self, new_pos):
     pos = tuple([float(x) for x in new_pos])
     self.xform.AddTranslateOp().Set(pos)
 
+  def update_rotation(self, new_mat):
+    r = R.from_matrix(new_mat)
+    euler_rotation = r.as_euler('xyz', degrees=True)
+    rotation = Gf.Vec3f(float(euler_rotation[0]), float(euler_rotation[1]), float(euler_rotation[2]))
+    self.xform.AddRotateXYZOp().Set(rotation)
+
+  # TODO: check to make sure scale and size are the same thing
   def update_size(self, new_size):
     size = tuple([float(x) for x in new_size])
     self.xform.AddScaleOp().Set(value=size)
@@ -41,8 +62,7 @@ class USDGeom(object):
   def update_color(self, new_color):
     # new_color is the rgba (we extract first three)
     rgba = [(float(x) for x in new_color[:3])]
-    color = self.prim.GetDisplayColorAttr()
+    color = self.prim.GetDisplayColorAttr().Set(rgba)
-    color.Set(rgba)
 
   # TODO: create another method for transparency
   def update_transparency(self, new_transparency):
@@ -70,10 +90,10 @@ class USDPlane(USDGeom):
     self.prim = UsdGeom.Plane.Define(stage, plane_path)
     self.ref = stage.GetPrimAtPath(plane_path)
 
-  def update_geom(self, new_geom):
+  def update_size(self, new_size):
-    self.update_pos(new_geom.pos)
+    self.prim.GetAxisAttr().Set("Z")
-    self.update_size(new_geom.size)
+    self.prim.GetWidthAttr().Set(float(new_size[0]))
-    self.update_color(new_geom.rgba)
+    self.prim.GetLengthAttr().Set(float(new_size[1]))
 
 class USDSphere(USDGeom):
   """
@@ -83,8 +103,8 @@ class USDSphere(USDGeom):
   sphere_count = 0
 
   def __init__(self, 
-               geom=None,
+               geom,
-               stage=None):
+               stage):
     super().__init__(geom, stage)
     self.type = 2
     USDSphere.sphere_count += 1
@@ -94,11 +114,6 @@ class USDSphere(USDGeom):
     self.prim = UsdGeom.Sphere.Define(stage, sphere_path)
     self.ref = stage.GetPrimAtPath(sphere_path)
 
-  def update_geom(self, new_geom):
-    self.update_pos(new_geom.pos)
-    self.update_size(new_geom.size)
-    self.update_color(new_geom.rgba)
-
 class USDCube(USDGeom):
   """
   Stores information regarding a cube geom in USD
@@ -107,8 +122,8 @@ class USDCube(USDGeom):
   cube_count = 0
 
   def __init__(self, 
-               geom=None,
+               geom,
-               stage=None):
+               stage):
     super().__init__(geom, stage)
     self.type = 6
     USDCube.cube_count += 1
@@ -118,24 +133,76 @@ class USDCube(USDGeom):
     self.prim = UsdGeom.Cube.Define(stage, cube_path)
     self.ref = stage.GetPrimAtPath(cube_path)
 
+class USDMesh(USDGeom):
+  """
+  Stores information regarding a mesh geom in USD
+  """
+
+  mesh_count = 0
+
+  def __init__(self,
+               mesh_idx,
+               geom,
+               stage,
+               model,
+               mesh_vertex_ranges,
+               mesh_face_ranges):
+    super().__init__(geom, stage)
+    self.type = 7
+    USDMesh.mesh_count += 1
+    xform_path = f'/Mesh_Xform_{USDMesh.mesh_count}'
+    mesh_path= f'{xform_path}/Mesh_{USDMesh.mesh_count}'
+    self.xform = UsdGeom.Xform.Define(stage, xform_path)
+    self.prim = UsdGeom.Mesh.Define(stage, mesh_path)
+    self.ref = stage.GetPrimAtPath(mesh_path)
+    
+    self.vertices = model.mesh_vert[mesh_vertex_ranges[mesh_idx]:mesh_vertex_ranges[mesh_idx+1]]
+    self.prim.GetPointsAttr().Set(self.vertices)
+
+    self.prim.GetFaceVertexCountsAttr().Set([3 for _ in range(model.mesh_facenum[mesh_idx])])
+
+    self.faces = model.mesh_face[mesh_face_ranges[mesh_idx]:mesh_face_ranges[mesh_idx+1]]
+    self.prim.GetFaceVertexIndicesAttr().Set(self.faces)
+
   def update_geom(self, new_geom):
     self.update_pos(new_geom.pos)
-    self.update_size(new_geom.size)
+    self.update_rotation(new_geom.mat)
-    self.update_color(new_geom.rgba)
 
-def create_usd_geom(geom, stage):
+    # TODO: remove this, temporary
-  geom_type = geom.type
+    self.xform.AddScaleOp().Set(value=(10.0, 10.0, 10.0))
-  if geom_type==USDGeomType.Plane.value:
-    return USDPlane(geom, stage)
-  elif geom_type==USDGeomType.Sphere.value:
-    return USDSphere(geom, stage)
-  elif geom_type==USDGeomType.Cube.value:
-    return USDCube(geom, stage)
   
 class USDLight(object):
-  def __init__(self):
+  """
-    pass
+  Class for the created lights
+  """
+
+  light_count = 0
+
+  def __init__(self,
+               stage):
+    self.stage = stage
+    USDLight.light_count += 1
+    xform_path = f'/Light_Xform_{USDLight.light_count}'
+    light_path = f'{xform_path}/Light_{USDLight.light_count}'
+    self.xform = UsdGeom.Xform.Define(stage, xform_path)
+    self.prim = UsdLux.SphereLight.Define(stage, light_path)
+    self.ref = stage.GetPrimAtPath(light_path)
 
   def update_light(self, new_light):
-    pass
+    pos = tuple([float(x) for x in new_light.pos])
-  
+    self.xform.AddTranslateOp().Set(pos)
+
+    # TODO attributes:
+    #   - direction
+    #   - intensity
+    #   - exposure
+    #   - radius
+    #   - specular
+
+    self.prim.GetIntensityAttr().Set(5000);
+
+
+
+
+
+    

python/mujoco/usd_renderer.py

@@ -2,6 +2,7 @@
 import mujoco
 import mujoco.viewer as viewer
 from mujoco.usd_component import *
+from mujoco.usd_utilities import *
 from pxr import Usd, UsdGeom
 
 class USDRenderer(object):
@@ -47,13 +48,29 @@ class USDRenderer(object):
     """
     Loads and initializes the necessary objects to render the scene
     """
+
+    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)
+
+    current_mesh_idx = 0
+
     # initializes an array to store all the geoms in the scene
     # populates with "empty" USDGeom objects
     self.usd_geoms = []
     geoms = self.scene.geoms
     self.ngeom = self.scene.ngeom
     for i in range(self.ngeom):
-      self.usd_geoms.append(create_usd_geom(geoms[i], self.stage))
+      if geoms[i].type == USDGeomType.Mesh.value:
+        self.usd_geoms.append(USDMesh(current_mesh_idx,
+                                      geoms[i], 
+                                      self.stage, 
+                                      self.model,
+                                      mesh_vertex_ranges,
+                                      mesh_face_ranges))
+        current_mesh_idx += 1
+      else:
+        self.usd_geoms.append(create_usd_geom_primitive(geoms[i], self.stage))
 
     # initializes an array to store all the lights in the scene
     # populates with "empty" USDLight objects
@@ -61,7 +78,7 @@ class USDRenderer(object):
     lights = self.scene.lights
     self.nlight = self.scene.nlight
     for i in range(self.nlight):
-      self.usd_lights.append(USDLight())
+      self.usd_lights.append(USDLight(self.stage))
 
   def _update(self):
     self._update_geoms()
@@ -83,9 +100,12 @@ class USDRenderer(object):
     lights = self.scene.lights
     nlight = self.scene.nlight
     for i in range(nlight):
-      print(self.usd_lights[i])
+      self.usd_lights[i].update_light(lights[i])
       
   def _update_camera(self):
+    """
+    Updates the camera to match the current scene
+    """
     pass
 
   def start_viewer(self):
@@ -94,11 +114,11 @@ class USDRenderer(object):
 
   def render(self):
     # should render the usd file given a particular renderer that
-    # works with USD files?
+    # works with USD files
     # TODO: determine if this is valid functionality
     pass
 
-  # TODO: remove later, this is only for debuggin purposes
+  # TODO: remove later, this is only for debugging purposes
   def print_geom_information(self):
     for i in range(self.ngeom):
       print(self.usd_geoms[i])

python/mujoco/usd_utilities.py

@@ -0,0 +1,8 @@
+
+def get_mesh_ranges(nmesh, arr):
+    mesh_ranges = [0]
+    running_sum = 0
+    for i in range(nmesh):
+        running_sum += arr[i]
+        mesh_ranges.append(running_sum)
+    return mesh_ranges