combining shapes into one class

2024-04-04 21:10:15 -04:00
parent 09bc73e117
commit f1fa9c0b8b
5 changed files with 205 additions and 340 deletions
@@ -0,0 +1,129 @@
+import copy
+import mujoco
+import pprint
+import numpy as np
+import open3d as o3d
+
+def mesh_config_generator(
+    name: str,
+    geom_type: mujoco.mjtGeom,
+    size: np.ndarray
+):
+
+    if geom_type == mujoco.mjtGeom.mjGEOM_PLANE:
+      return {
+        "name": name,
+        "box": {
+          "width": size[0] * 2 if size[0] > 0 else 100,
+          "height": size[1] * 2 if size[1] > 0 else 100,
+          "depth": 0.001,
+          "map_texture_to_each_face": True,
+        }
+      }
+    elif geom_type == mujoco.mjtGeom.mjGEOM_SPHERE:
+      return {
+        "name": name,
+        "sphere": {
+          "radius": float(size[0])
+        }
+      }
+    elif geom_type == mujoco.mjtGeom.mjGEOM_CAPSULE:
+      cylinder = mesh_config_generator(name, mujoco.mjtGeom.mjGEOM_CYLINDER, size)
+      left_sphere = mesh_config_generator(name, mujoco.mjtGeom.mjGEOM_SPHERE, size)
+      right_sphere = copy.deepcopy(left_sphere)
+      left_sphere["sphere"]["transform"] = {
+        "translate": (0, 0, -size[2])
+      }
+      right_sphere["sphere"]["transform"] = {
+        "translate": (0, 0, size[2])
+      } 
+      return {
+        "name": name,
+        "cylinder": cylinder["cylinder"],
+        "left_sphere": left_sphere["sphere"],
+        "right_sphere": right_sphere["sphere"],
+      }
+    elif geom_type == mujoco.mjtGeom.mjGEOM_ELLIPSOID:
+      sphere = mesh_config_generator(name, mujoco.mjtGeom.mjGEOM_SPHERE, [1.0])
+      sphere["sphere"]["transform"] = {
+        "scale": tuple(size)
+      }
+      return {
+        "name": name,
+        "sphere": sphere["sphere"],
+      }
+    elif geom_type == mujoco.mjtGeom.mjGEOM_CYLINDER:
+      return {
+        "name": name,
+        "cylinder": {
+          "radius": size[0],
+          "height": size[2] * 2,
+        }
+      }
+    elif geom_type == mujoco.mjtGeom.mjGEOM_BOX:
+      return {
+        "name": name,
+        "box": {
+          "width": size[0] * 2,
+          "height": size[1] * 2,
+          "depth": size[2] * 2,
+        }
+      }
+    else:
+      raise NotImplemented(f"{geom_type} primitive geom type not implemented with USD integration")
+
+
+def mesh_generator(
+    mesh_config: dict
+):
+
+    assert "name" in mesh_config
+
+    mesh = None
+
+    for shape, config in mesh_config.items():
+
+      if shape == "name":
+        continue
+
+      if "box" in shape:
+        prim_mesh = o3d.geometry.TriangleMesh.create_box(
+            width=mesh_config[shape]["width"],
+            height=mesh_config[shape]["height"],
+            depth=mesh_config[shape]["depth"],
+            create_uv_map=True,
+            map_texture_to_each_face=True,
+        )
+      elif "sphere" in shape:
+        prim_mesh = o3d.geometry.TriangleMesh.create_sphere(
+            radius=mesh_config[shape]["radius"],
+            create_uv_map=True
+        )
+      elif "cylinder" in shape:
+        prim_mesh = o3d.geometry.TriangleMesh.create_cylinder(
+            radius=mesh_config[shape]["radius"],
+            height=mesh_config[shape]["height"],
+            create_uv_map=True,
+        )
+
+      if "transform" in config:
+        for transform, val in config["transform"].items():
+          if transform == "translate":
+            prim_mesh.translate(val)
+          if transform == "scale":
+            prim_mesh.vertices = o3d.utility.Vector3dVector(
+              np.asarray(prim_mesh.vertices) * np.array(val))
+
+      if not mesh:
+        mesh = prim_mesh
+      else:
+        mesh += prim_mesh
+
+    return mesh_config["name"], mesh
+      
+
+
+
+
+    
+