Visualizing specific geom groups

python/mujoco/usd_component.py

@@ -6,12 +6,19 @@ from mujoco.usd_utilities import *
 from pxr import Usd, UsdGeom, UsdLux, UsdShade, Vt, Gf, Sdf
 from scipy.spatial.transform import Rotation as R
 
+# TODO: clean this up and remove the if statements
 def create_usd_geom_primitive(geom, stage, texture_file):
   geom_type = geom.type
   if geom_type==USDGeomType.Plane.value:
     return USDPlane(geom, stage, texture_file)
   elif geom_type==USDGeomType.Sphere.value:
     return USDSphere(geom, stage, texture_file)
+  elif geom_type==USDGeomType.Capsule.value:
+    return USDCapsule(geom, stage, texture_file)
+  elif geom_type==USDGeomType.Cylinder.value:
+    return USDCylinder(geom, stage, texture_file)
+  elif geom_type==USDGeomType.Cube.value:
+    return USDCube(geom, stage, texture_file)
   elif geom_type==USDGeomType.Cube.value:
     return USDCube(geom, stage, texture_file)
 
@@ -22,7 +29,11 @@ class USDGeomType(Enum):
   https://mujoco.readthedocs.io/en/latest/APIreference/APItypes.html#mjtgeom
   """
   Plane = 0
+  # Hfield = 1
   Sphere = 2
+  Capsule = 3
+  # Ellipsoid = 4
+  Cylinder = 5
   Cube = 6
   Mesh = 7
 
@@ -35,18 +46,40 @@ class USDGeom(object):
                stage=None,
                texture_file=None):
     self.geom = geom
-    self.stage = stage # TODO: remove, not being used
+    self.stage = stage
     self.texture_file = texture_file
     self.type = None
     self.xform = None
     self.prim = None
     self.ref = None
+
+  def set_texture(self):
+    if self.texture_file:
+      mtl_path = Sdf.Path(f"/World/Looks/Material_{os.path.splitext(os.path.basename(self.texture_file))[0]}")
+      mtl = UsdShade.Material.Define(self.stage, mtl_path)
+      shader = UsdShade.Shader.Define(self.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(self.stage, mtl_path.AppendPath("DiffuseColorTx"))
+      diffuse_tx.CreateIdAttr('UsdUVTexture') 
+
+      diffuse_tx.CreateInput('file', Sdf.ValueTypeNames.Asset).Set(self.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)
     self.update_size(new_geom.size)
     self.update_color(new_geom.rgba)
+    self.update_transparency(new_geom.rgba[3])
   
   def update_pos(self, new_pos):
     pos = tuple([float(x) for x in new_pos])
@@ -66,11 +99,10 @@ 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().Set(rgba)
+    self.prim.GetDisplayColorAttr().Set(rgba)
 
-  # TODO: create another method for transparency
   def update_transparency(self, new_transparency):
-    pass
+    self.prim.GetDisplayOpacityAttr().Set(new_transparency)
 
   def __str__(self):
     return f'type = {self.type} \ngeom = {self.geom}'
@@ -89,12 +121,14 @@ class USDPlane(USDGeom):
     super().__init__(geom, stage, texture_file)
     self.type = 0
     USDPlane.plane_count += 1
-    xform_path = f'/Plane_Xform_{USDPlane.plane_count}'
+    xform_path = f'/World/Plane_Xform_{USDPlane.plane_count}'
     plane_path = f'{xform_path}/Plane_{USDPlane.plane_count}'
     self.xform = UsdGeom.Xform.Define(stage, xform_path)
     self.prim = UsdGeom.Plane.Define(stage, plane_path)
     self.ref = stage.GetPrimAtPath(plane_path)
 
+    self.set_texture()
+
   def update_size(self, new_size):
     self.prim.GetAxisAttr().Set("Z")
     self.prim.GetWidthAttr().Set(float(new_size[0]))
@@ -114,12 +148,80 @@ class USDSphere(USDGeom):
     super().__init__(geom, stage, texture_file)
     self.type = 2
     USDSphere.sphere_count += 1
-    xform_path = f'/Sphere_Xform_{USDSphere.sphere_count}'
+    xform_path = f'/World/Sphere_Xform_{USDSphere.sphere_count}'
     sphere_path = f'{xform_path}/Sphere_{USDSphere.sphere_count}'
     self.xform = UsdGeom.Xform.Define(stage, xform_path)
     self.prim = UsdGeom.Sphere.Define(stage, sphere_path)
     self.ref = stage.GetPrimAtPath(sphere_path)
 
+    self.set_texture()
+
+class USDCapsule(USDGeom):
+  """
+  Stores information regarding a capsule geom in USD
+  """
+
+  capsule_count = 0
+
+  def __init__(self, 
+               geom=None,
+               stage=None,
+               texture_file=None):
+    super().__init__(geom, stage, texture_file)
+    self.type = 3
+    USDCapsule.capsule_count += 1
+    xform_path = f'/World/Capsule_Xform_{USDCapsule.capsule_count}'
+    capsule_path = f'{xform_path}/Capsule_{USDCapsule.capsule_count}'
+    self.xform = UsdGeom.Xform.Define(stage, xform_path)
+    self.prim = UsdGeom.Capsule.Define(stage, sphere_path)
+    self.ref = stage.GetPrimAtPath(capsule_path)
+
+    self.set_texture()
+
+class USDCylinder(USDGeom):
+  """
+  Stores information regarding a capsule geom in USD
+  """
+
+  cylinder_count = 0
+
+  def __init__(self, 
+               geom=None,
+               stage=None,
+               texture_file=None):
+    super().__init__(geom, stage, texture_file)
+    self.type = 5
+    USDCylinder.cylinder_count += 1
+    xform_path = f'/World/Cylinder_Xform_{USDCylinder.cylinder_count}'
+    cylinder_path = f'{xform_path}/Cylinder_{USDCylinder.cylinder_count}'
+    self.xform = UsdGeom.Xform.Define(stage, xform_path)
+    self.prim = UsdGeom.Cylinder.Define(stage, cylinder_path)
+    self.ref = stage.GetPrimAtPath(cylinder_path)
+
+    self.set_texture()
+
+class USDSphere(USDGeom):
+  """
+  Stores information regarding a sphere geom in USD
+  """
+
+  sphere_count = 0
+
+  def __init__(self, 
+               geom=None,
+               stage=None,
+               texture_file=None):
+    super().__init__(geom, stage, texture_file)
+    self.type = 2
+    USDSphere.sphere_count += 1
+    xform_path = f'/World/Sphere_Xform_{USDSphere.sphere_count}'
+    sphere_path = f'{xform_path}/Sphere_{USDSphere.sphere_count}'
+    self.xform = UsdGeom.Xform.Define(stage, xform_path)
+    self.prim = UsdGeom.Sphere.Define(stage, sphere_path)
+    self.ref = stage.GetPrimAtPath(sphere_path)
+
+    self.set_texture()
+
 class USDCube(USDGeom):
   """
   Stores information regarding a cube geom in USD
@@ -134,31 +236,13 @@ class USDCube(USDGeom):
     super().__init__(geom, stage, texture_file)
     self.type = 6
     USDCube.cube_count += 1
-    xform_path = f'/Cube_Xform_{USDCube.cube_count}'
+    xform_path = f'/World/Cube_Xform_{USDCube.cube_count}'
     cube_path = f'{xform_path}/Cube_{USDCube.cube_count}'
     self.xform = UsdGeom.Xform.Define(stage, xform_path)
     self.prim = UsdGeom.Cube.Define(stage, cube_path)
     self.ref = stage.GetPrimAtPath(cube_path)
 
-    if texture_file:
-      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') 
-
-      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)
+    self.set_texture()
 
 class USDMesh(USDGeom):
   """
@@ -186,7 +270,7 @@ class USDMesh(USDGeom):
     mesh_face = model.mesh_face[mesh_face_adr_from:mesh_face_adr_to]
 
     self.type = 7
-    xform_path = f'/Mesh_Xform_{USDMesh.mesh_count}'
+    xform_path = f'/World/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)
@@ -243,7 +327,6 @@ class USDMesh(USDGeom):
 
     USDMesh.mesh_count += 1
 
-
   def update_geom(self, new_geom):
     self.update_pos(new_geom.pos)
     self.update_rotation(new_geom.mat)
@@ -264,7 +347,7 @@ class USDLight(object):
                stage):
     self.stage = stage
     USDLight.light_count += 1
-    xform_path = f'/Light_Xform_{USDLight.light_count}'
+    xform_path = f'/World/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)
@@ -281,7 +364,36 @@ class USDLight(object):
     #   - radius
     #   - specular
 
-    self.prim.GetIntensityAttr().Set(30000);
+    self.prim.GetIntensityAttr().Set(15000);
+
+
+class USDCamera(object):
+  """
+  Class for created cameras
+  """
+
+  camera_count = 0
+
+  def __init__(self,
+               stage):
+    self.stage = stage
+    USDCamera.camera_count += 1
+    camera_path = f'/World/Camera_{USDCamera.camera_count}'
+    self.prim = UsdGeom.Camera.Define(stage, camera_path)
+    self.ref = stage.GetPrimAtPath(camera_path)
+
+  def update_camera(self, new_camera):
+    # print("---- Updating camera in USD ----")
+    xformAPI = UsdGeom.XformCommonAPI(self.prim)
+
+    # hardcoded values for Robosuite testing and prototype 
+    # TODO: use actual camera values
+    xformAPI.SetTranslate((2.25, 0.0, 2.1))
+    xformAPI.SetRotate((65.0, 0.0, 90.0))
+    xformAPI.SetScale((1, 1, 1))
+
+    self.prim.CreateFocalLengthAttr().Set(24)
+    self.prim.CreateFocusDistanceAttr().Set(400)
 
 
 

python/mujoco/usd_renderer.py

@@ -5,6 +5,8 @@ from mujoco.usd_component import *
 from mujoco.usd_utilities import *
 from pxr import Usd, UsdGeom
 
+from mujoco import _structs
+
 from PIL import Image as im
 from PIL import ImageOps
 
@@ -15,7 +17,8 @@ class USDRenderer(object):
   def __init__(self,
                model,
                height=480,
-               width=480):
+               width=480,
+               geom_groups=[1,1,1,1,1,1]):
     self.model = model
     self.data = None
     self.renderer = mujoco.Renderer(model, height, width)
@@ -25,6 +28,10 @@ class USDRenderer(object):
     self.stage = Usd.Stage.CreateNew('usd_stage.usda')
     UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.z)
 
+    # TODO: maybe change where we initialize this?
+    self.scene_option = _structs.MjvOption()
+    self.scene_option.geomgroup = geom_groups
+
   @property
   def usd(self):
     return self.stage.GetRootLayer().ExportToString()
@@ -37,7 +44,7 @@ class USDRenderer(object):
     self.stage.GetRootLayer().Save()
 
   def update_scene(self, data):
-    self.renderer.update_scene(data)
+    self.renderer.update_scene(data, scene_option=self.scene_option)
     self.data = data
 
     if not self.loaded_scene_info:
@@ -53,15 +60,9 @@ class USDRenderer(object):
     Loads and initializes the necessary objects to render the scene
     """
 
-    # TODO: remove these and replace by reading directly from model
-    # 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
+    # TODO: remove once added to mujoco
     data_adr = 0
     texture_files = []
     for texid in range(self.model.ntex):
@@ -85,20 +86,17 @@ class USDRenderer(object):
     self.ngeom = self.scene.ngeom
     for i in range(self.ngeom):
       geom = geoms[i]
-      if geom.category == 1:
-        self.usd_geoms.append(None)
-        continue
       if geom.texid == -1:
         texture_file = None
       else:
         texture_file = texture_files[geom.texid]
 
       if geom.type == USDGeomType.Mesh.value:
-        self.usd_geoms.append(USDMesh(self.model.geom_dataid[i],
-                                      geom, 
-                                      self.stage, 
-                                      self.model,
-                                      texture_file))
+        self.usd_geoms.append(USDMesh(self.model.geom_dataid[geom.objid],
+                                        geom, 
+                                        self.stage, 
+                                        self.model,
+                                        texture_file))
       else:
         self.usd_geoms.append(create_usd_geom_primitive(geom, 
                                                         self.stage,
@@ -112,10 +110,17 @@ class USDRenderer(object):
     for i in range(self.nlight):
       self.usd_lights.append(USDLight(self.stage))
 
+    # initializes an array to store all the cameras in the scene
+    # populates with "empty" USDCamera objects
+    self.usd_cameras = []
+    cameras = self.scene.camera
+    for i in range(2):
+      self.usd_cameras.append(USDCamera(self.stage))
+
   def _update(self):
     self._update_geoms()
     self._update_lights()
-    self._update_camera()
+    self._update_cameras()
 
   def _update_geoms(self):
     """
@@ -135,11 +140,13 @@ class USDRenderer(object):
     for i in range(nlight):
       self.usd_lights[i].update_light(lights[i])
       
-  def _update_camera(self):
+  def _update_cameras(self):
     """
     Updates the camera to match the current scene
     """
-    pass
+    cameras = self.scene.camera
+    for i in range(2):
+      self.usd_cameras[i].update_camera(cameras[i])
 
   def start_viewer(self):
     if self.data: