combining shapes into one class

python/mujoco/usd/exporter.py

@@ -15,7 +15,12 @@
 import os
 
 import mujoco
-import mujoco.usd.component as component_module
+
+# import mujoco.usd.shapes as shapes_module
+# import mujoco.usd.component as component_module
+import shapes as shapes_module
+import component as component_module
+
 import numpy as np
 import scipy
 import termcolor
@@ -234,7 +239,6 @@ class USDExporter:
 
     texture_file = self.texture_files[geom.texid] if geom.texid != -1 else None
 
-    # handles meshes in scene
     if geom.type == mujoco.mjtGeom.mjGEOM_MESH:
       usd_geom = component_module.USDMesh(
           stage=self.stage,
@@ -245,56 +249,22 @@ class USDExporter:
           rgba=geom.rgba,
           texture_file=texture_file,
       )
-    elif geom.type == mujoco.mjtGeom.mjGEOM_PLANE:
-      usd_geom = component_module.USDPlaneMesh(
-          stage=self.stage,
-          geom=geom,
-          obj_name=geom_name,
-          rgba=geom.rgba,
-          texture_file=texture_file,
-      )
-    elif geom.type == mujoco.mjtGeom.mjGEOM_SPHERE:
-      usd_geom = component_module.USDSphereMesh(
-          stage=self.stage,
-          geom=geom,
-          obj_name=geom_name,
-          rgba=geom.rgba,
-          texture_file=texture_file,
-      )
-    elif geom.type == mujoco.mjtGeom.mjGEOM_CAPSULE:
-      usd_geom = component_module.USDCapsule(
-          stage=self.stage,
-          geom=geom,
-          obj_name=geom_name,
-          rgba=geom.rgba,
-          texture_file=texture_file,
-      )
-    elif geom.type == mujoco.mjtGeom.mjGEOM_ELLIPSOID:
-      usd_geom = component_module.USDEllipsoid(
-          stage=self.stage,
-          geom=geom,
-          obj_name=geom_name,
-          rgba=geom.rgba,
-          texture_file=texture_file,
-      )
-    elif geom.type == mujoco.mjtGeom.mjGEOM_CYLINDER:
-      usd_geom = component_module.USDCylinderMesh(
-          stage=self.stage,
-          geom=geom,
-          obj_name=geom_name,
-          rgba=geom.rgba,
-          texture_file=texture_file,
-      )
-    elif geom.type == mujoco.mjtGeom.mjGEOM_BOX:
-      usd_geom = component_module.USDCubeMesh(
-          stage=self.stage,
-          geom=geom,
-          obj_name=geom_name,
-          rgba=geom.rgba,
-          texture_file=texture_file,
-      )
     else:
-      usd_geom = None
+
+      mesh_config = shapes_module.mesh_config_generator(
+          name=geom_name,
+          geom_type=geom.type,
+          size=geom.size   
+      )
+
+      usd_geom = component_module.USDPrimitiveMesh(
+          mesh_config=mesh_config,
+          stage=self.stage,
+          geom=geom,
+          obj_name=geom_name,
+          rgba=geom.rgba,
+          texture_file=texture_file,
+      )
 
     self.geom_name2usd[geom_name] = usd_geom