Removing Open3D dependency

python/mujoco/usd/camera.py

@@ -14,7 +14,8 @@
 # ==============================================================================
 """Camera handling for USD exporter."""
 
-import mujoco.usd.utils as utils_component
+# import mujoco.usd.utils as utils_modules
+import utils as utils_module
 
 import numpy as np
 
@@ -49,7 +50,7 @@ class USDCamera:
 
   def update(self, cam_pos: np.ndarray, cam_mat: np.ndarray, frame: int):
     """Updates the position and orientation of the camera in the scene."""
-    transformation_mat = utils_component.create_transform_matrix(
+    transformation_mat = utils_modules.create_transform_matrix(
         rotation_matrix=cam_mat, translation_vector=cam_pos
     ).T
     self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)

python/mujoco/usd/demo.py

@@ -18,8 +18,8 @@ import argparse
 import pathlib
 
 import mujoco
-from mujoco.usd import exporter
+# from mujoco.usd import exporter
-
+import exporter
 
 def generate_usd_trajectory(local_args):
   """Generates a USD file given the user arguments."""
@@ -42,6 +42,10 @@ def generate_usd_trajectory(local_args):
     if exp.frame_count < d.time * local_args.framerate:
       exp.update_scene(data=d)
 
+  exp.add_light(pos=(0, 0, 0),
+                intensity=2000,
+                light_type="dome")
+
   exp.save_scene(filetype=local_args.export_extension)
 
 

python/mujoco/usd/exporter.py

@@ -18,10 +18,14 @@ import os
 from typing import List, Optional
 
 import mujoco
-import mujoco.usd.camera as camera_module
+# import mujoco.usd.camera as camera_module
-import mujoco.usd.lights as light_module
+# import mujoco.usd.lights as light_module
-import mujoco.usd.objects as object_module
+# import mujoco.usd.objects as object_module
-import mujoco.usd.shapes as shapes_module
+# import mujoco.usd.shapes as shapes_module
+import camera as camera_module
+import lights as light_module
+import objects as object_module
+import shapes as shapes_module
 import numpy as np
 from PIL import Image as im
 from PIL import ImageOps
@@ -253,7 +257,7 @@ class USDExporter:
 
     texture_file = (
         self.texture_files[
-            self.model.mat_texid[geom.matid][mujoco.mjTEXROLE_RGB]
+            self.model.mat_texid[geom.matid][mujoco.mjtTextureRole.mjTEXROLE_RGB]
         ]
         if geom.matid != -1
         else None

python/mujoco/usd/objects.py

@@ -19,10 +19,13 @@ import collections
 from typing import Optional, Dict, Any
 
 import mujoco
-import mujoco.usd.shapes as shapes_component
+# import mujoco.usd.shapes as shapes_module
-import mujoco.usd.utils as utils_component
+# import mujoco.usd.utils as utils_modules
+import shapes as shapes_module
+import utils as utils_module
 import numpy as np
 
+
 # TODO: b/288149332 - Remove once USD Python Binding works well with pytype.
 # pytype: disable=module-attr
 from pxr import Gf
@@ -185,7 +188,7 @@ class USDObject(abc.ABC):
       scale: Optional[np.ndarray] = None,
   ):
     """Updates the position and orientation of an object."""
-    transformation_mat = utils_component.create_transform_matrix(
+    transformation_mat = utils_module.create_transform_matrix(
         rotation_matrix=mat, translation_vector=pos
     ).T
     self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
@@ -353,23 +356,31 @@ class USDPrimitiveMesh(USDObject):
 
   def generate_primitive_mesh(self):
     """Generates the mesh for the primitive USD object."""
-    _, prim_mesh = shapes_component.mesh_generator(self.mesh_config)
+    _, prim_mesh = shapes_module.mesh_factory(self.mesh_config)
     prim_mesh.translate(-prim_mesh.get_center())
     return prim_mesh
 
   def _get_uv_geometry(self):
     assert self.prim_mesh
 
+    # x_scale, y_scale = self.geom.texrepeat
+
     mesh_texcoord = np.array(self.prim_mesh.triangle_uvs)
     mesh_facetexcoord = np.asarray(self.prim_mesh.triangles)
-    # TODO(etom): bring back support for rescaling the texture coordinates.
+    
+    # x_multiplier, y_multiplier = 1, 1
+    # if self.geom.texuniform:
+    #   x_multiplier, y_multiplier = self.geom.size[:2]
+
+    # mesh_texcoord[:, 0] *= x_scale * x_multiplier
+    # mesh_texcoord[:, 1] *= y_scale * y_multiplier
 
     return mesh_texcoord, mesh_facetexcoord.flatten()
 
   def _get_mesh_geometry(self):
     assert self.prim_mesh
 
-    # get mesh geometry from the open3d mesh model
+    # get mesh geometry
     mesh_vert = np.asarray(self.prim_mesh.vertices)
     mesh_face = np.asarray(self.prim_mesh.triangles)
 
@@ -444,7 +455,7 @@ class USDTendon(USDObject):
     """Generates the tendon mesh using primitives."""
     mesh_parts = {}
     for part_config in self.mesh_config:
-      mesh_name, prim_mesh = shapes_component.mesh_generator(part_config)
+      mesh_name, prim_mesh = shapes_module.mesh_factory(part_config)
       prim_mesh.translate(-prim_mesh.get_center())
       mesh_parts[mesh_name] = prim_mesh
     return mesh_parts
@@ -463,7 +474,7 @@ class USDTendon(USDObject):
   def _get_mesh_geometry(self):
     part_geometries = collections.defaultdict(dict)
     for name, mesh in self.tendon_parts.items():
-      # get mesh geometry from the open3d mesh model
+      # get mesh geometry
       mesh_vert = np.asarray(mesh.vertices)
       mesh_face = np.asarray(mesh.triangles)
       part_geometries[name] = {

python/mujoco/usd/shapes.py

@@ -14,34 +14,184 @@
 # ==============================================================================
 """Built-in shapes for USD exporter."""
 
-from typing import Dict, Any
+from typing import Dict, Any, Tuple
 
 import mujoco
 import numpy as np
-from open3d import open3d as o3d
 
+def get_triangle_uvs(vertices: np.array, triangles: np.array):
+    # (jabhi) asusming all mappings are 2d mapping temporarily
+    triangle_uvs = np.array([
+      [vertices[i][0], vertices[i][1]] for i in np.nditer(triangles)]
+    )
+    return triangle_uvs
 
-def create_hemisphere(
+class TriangleMesh():
-    radius: float, theta_steps: int = 50, phi_steps: int = 50
+  """ Store UV and geometry information for a primitve mesh
-):
+  """
-  """Creates a hemisphere mesh from a point cloud."""
+  def __init__(self,
-  points = []
+               vertices: np.array,
-  for i in range(phi_steps + 1):
+               triangles: np.array,
-    phi = np.pi / 2 * i / phi_steps
+               triangle_uvs: np.array):
-    for j in range(theta_steps + 1):
+    self.vertices = vertices
-      theta = 2 * np.pi * j / theta_steps
+    self.triangles = triangles
-      x = radius * np.sin(phi) * np.cos(theta)
+    self.triangle_uvs = triangle_uvs
-      y = radius * np.sin(phi) * np.sin(theta)
-      z = radius * np.cos(phi)
-      points.append([x, y, z])
 
-  pcd = o3d.geometry.PointCloud()
+  @classmethod
-  pcd.points = o3d.utility.Vector3dVector(points)
+  def create_box(
+      cls, width: float, height: float, depth: float
+  ):
+    vertices = np.array([[0.0, 0.0, 0.0],
+                         [width, 0.0, 0.0],
+                         [0.0, 0.0, depth],
+                         [width, 0.0, depth],
+                         [0.0, height, 0.0],
+                         [width, height, 0.0],
+                         [0.0, height, depth],
+                         [width, height, depth]])
     
-  mesh = pcd.compute_convex_hull()[0]
+    triangles = np.array([[4, 7, 5],
+                          [4, 6, 7],
+                          [0, 2, 4],
+                          [2, 6, 4],
+                          [0, 1, 2],
+                          [1, 3, 2],
+                          [1, 5, 7],
+                          [1, 7, 3],
+                          [2, 3, 7],
+                          [2, 7, 6],
+                          [0, 4, 1],
+                          [1, 4, 5]])
     
-  return mesh
+    triangle_uvs = get_triangle_uvs(vertices, triangles)
     
+    return TriangleMesh(vertices, triangles, triangle_uvs)
+
+  @classmethod
+  def create_sphere(
+      cls, radius: float, resolution: int
+  ):
+    vertices = []
+    triangles = []
+    for i in range(2*resolution + 1):
+      phi = np.pi * i / (2*resolution)
+      for j in range(resolution + 1):
+        theta = 2 * np.pi * j / resolution
+        x = radius * np.sin(phi) * np.cos(theta)
+        y = radius * np.sin(phi) * np.sin(theta)
+        z = radius * np.cos(phi)
+        vertices.append([x, y, z])
+
+    for i in range(2*resolution):
+      for j in range(resolution):
+        first = i * (resolution + 1) + j
+        second = first + resolution + 1
+
+        triangles.append([first, second, first + 1])
+        triangles.append([second, second + 1, first + 1])
+
+    vertices = np.array(vertices)
+    triangles = np.array(triangles)
+
+    triangle_uvs = get_triangle_uvs(vertices, triangles)
+    
+    return TriangleMesh(vertices, triangles, triangle_uvs)
+
+  @classmethod
+  def create_hemisphere(
+    cls, radius: float, resolution: int
+  ):
+    vertices = []
+    triangles = []
+    for i in range(resolution + 1):
+      phi = np.pi / 2 * i / (resolution)
+      for j in range(resolution + 1):
+        theta = 2 * np.pi * j / resolution
+        x = radius * np.sin(phi) * np.cos(theta)
+        y = radius * np.sin(phi) * np.sin(theta)
+        z = radius * np.cos(phi)
+        vertices.append([x, y, z])
+    vertices.append([0, 0, 0])
+
+    for i in range(resolution):
+      for j in range(resolution):
+        first = i * (resolution + 1) + j
+        second = first + resolution + 1
+
+        triangles.append([first, second, first + 1])
+        triangles.append([second, second + 1, first + 1])
+
+    for i in range(resolution):
+      first = resolution * (resolution + 1) + i
+      triangles.append([first, first + 1, len(vertices) - 1])
+
+    vertices = np.array(vertices)
+    triangles = np.array(triangles)
+
+    triangle_uvs = get_triangle_uvs(vertices, triangles)
+    
+    return TriangleMesh(vertices, triangles, triangle_uvs)
+
+  @classmethod
+  def create_cylinder(
+      cls, radius: float, height: float, resolution: int
+  ):
+    vertices = []
+    triangles = []
+
+    # adding all the vertices for the cylinder including
+    # two center vertices at ends
+    for i in range(2):
+      z = 0 if i == 0 else height
+      for j in range(resolution + 1):
+        theta = 2 * np.pi * j / resolution
+        x = radius * np.cos(theta)
+        y = radius * np.sin(theta)
+        vertices.append([x, y, z])
+    vertices.append([0, 0, 0])
+    vertices.append([0, 0, height])
+
+    # constructing the end faces for the cylinder
+    for i in range(2):
+      for j in range(resolution):
+        first = (resolution + 1) * i + j
+        triangles.append([first, first + 1, len(vertices) - (2 - i)])
+
+    # constructing side of cylinder
+    for i in range(resolution):
+      second = resolution + 1 + i
+      triangles.append([i, second, second + 1])
+      triangles.append([i, i + 1, second + 1])
+
+    vertices = np.array(vertices)
+    triangles = np.array(triangles)
+
+    triangle_uvs = get_triangle_uvs(vertices, triangles)
+    
+    return TriangleMesh(vertices, triangles, triangle_uvs)
+
+  def translate(self, translation: np.array):
+    self.vertices = self.vertices + translation
+
+  def rotate(self, rotation: np.array, center: Tuple[float]):
+    translated_point = self.vertices - center
+    self.vertices = np.dot(translated_point, rotation) + center
+  
+  def scale(self, scale: np.array):
+    self.vertices = self.vertices * scale
+
+  def get_center(self):
+    center = np.mean(self.vertices, axis=0)
+    return center
+  
+  def __add__(self, other):
+    if isinstance(other, TriangleMesh):
+      new_vertices = np.vstack((self.vertices, other.vertices))
+      other_triangles = other.triangles + len(self.vertices)
+      new_triangles = np.vstack((self.triangles, other_triangles))
+      new_triangle_uvs = get_triangle_uvs(new_vertices, new_triangles)
+      return TriangleMesh(new_vertices, new_triangles, new_triangle_uvs)
+    raise TypeError(f"Cannot add TriangleMesh with {type(other)}")
 
 def decouple_config(config: Dict[str, Any]):
   """Breaks a shape config into is subcomponent shapes."""
@@ -79,20 +229,20 @@ def mesh_config_generator(
     config = {"name": name, "sphere": {"radius": float(size[0])}}
   elif geom_type == mujoco.mjtGeom.mjGEOM_CAPSULE:
     cylinder = mesh_config_generator(name, mujoco.mjtGeom.mjGEOM_CYLINDER, size)
+    cylinder["cylinder"]["transform"] = {"transform": {"translate": (0, 0, size[2])}}
     config = {
         "name": name,
         "cylinder": cylinder["cylinder"],
         "left_hemisphere": {
             "radius": size[0],
             "transform": {
-                "translate": (0, 0, -size[2]),
                 "rotate": (np.pi, 0, 0),
             },
         },
         "right_hemisphere": {
             "radius": size[0],
-            "transform": {"translate": (0, 0, size[2])},
+            "transform": {"translate": (0, 0, 2*size[2])},
-        },
+        }
     }
   elif geom_type == mujoco.mjtGeom.mjGEOM_ELLIPSOID:
     sphere = mesh_config_generator(
@@ -130,10 +280,9 @@ def mesh_config_generator(
 
   return config
 
-
+def mesh_factory(
-def mesh_generator(
     mesh_config: Dict[str, Any],
-    resolution: int = 100,
+    resolution: int = 20,
 ):
   """Generates a mesh given a config consisting of shapes."""
   assert "name" in mesh_config
@@ -145,28 +294,29 @@ def mesh_generator(
     if "name" in shape:
       continue
 
+    prim_mesh = None
+
     if "box" in shape:
-      prim_mesh = o3d.geometry.TriangleMesh.create_box(
+      prim_mesh = TriangleMesh.create_box(
           width=mesh_config[shape]["width"],
           height=mesh_config[shape]["height"],
-          depth=mesh_config[shape]["depth"],
+          depth=mesh_config[shape]["depth"]
-          create_uv_map=True,
-          map_texture_to_each_face=True,
       )
     elif "hemisphere" in shape:
-      prim_mesh = create_hemisphere(radius=mesh_config[shape]["radius"])
+      prim_mesh = TriangleMesh.create_hemisphere(
-    elif "sphere" in shape:
-      prim_mesh = o3d.geometry.TriangleMesh.create_sphere(
           radius=mesh_config[shape]["radius"],
-          resolution=resolution,
+          resolution=resolution
-          create_uv_map=True,
+      )
+    elif "sphere" in shape:
+      prim_mesh = TriangleMesh.create_sphere(
+          radius=mesh_config[shape]["radius"],
+          resolution=resolution
       )
     elif "cylinder" in shape:
-      prim_mesh = o3d.geometry.TriangleMesh.create_cylinder(
+      prim_mesh = TriangleMesh.create_cylinder(
           radius=mesh_config[shape]["radius"],
           height=mesh_config[shape]["height"],
-          resolution=resolution,
+          resolution=resolution
-          create_uv_map=True,
       )
     else:
       raise ValueError("Shape not supported")
@@ -182,10 +332,7 @@ def mesh_generator(
         rotation = rotation.reshape((3, 3))
         prim_mesh.rotate(rotation, center=(0, 0, 0))
       if "scale" in config["transform"]:
-        prim_mesh.vertices = o3d.utility.Vector3dVector(
+        prim_mesh.scale(config["transform"]["scale"])
-            np.asarray(prim_mesh.vertices)
-            * np.array(config["transform"]["scale"])
-        )
       if "translate" in config["transform"]:
         prim_mesh.translate(config["transform"]["translate"])