chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Removing Open3D dependency

Browse Source
This commit is contained in:
Abhishek Joshi
2024-07-20 19:33:11 -05:00
parent d49d4bf64c
commit 34060fcc49
5 changed files with 225 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/usd/camera.py
+3 -2
View File
@@ -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
+6 -2
View File
@@ -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
+9 -5
View File
@@ -18,10 +18,14 @@ import os
from typing import List, Optional
import mujoco
import mujoco.usd.camera as camera_module
import mujoco.usd.lights as light_module
import mujoco.usd.objects as object_module
import mujoco.usd.shapes as shapes_module
# import mujoco.usd.camera as camera_module
# import mujoco.usd.lights as light_module
# import mujoco.usd.objects as object_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 -8
View File
@@ -19,10 +19,13 @@ import collections
from typing import Optional, Dict, Any
import mujoco
import mujoco.usd.shapes as shapes_component
import mujoco.usd.utils as utils_component
# import mujoco.usd.shapes as shapes_module
# 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
+188 -41
View File
@@ -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(
radius: float, theta_steps: int = 50, phi_steps: int = 50
):
"""Creates a hemisphere mesh from a point cloud."""
points = []
for i in range(phi_steps + 1):
phi = np.pi / 2 * i / phi_steps
for j in range(theta_steps + 1):
theta = 2 * np.pi * j / theta_steps
x = radius * np.sin(phi) * np.cos(theta)
y = radius * np.sin(phi) * np.sin(theta)
z = radius * np.cos(phi)
points.append([x, y, z])
class TriangleMesh():
""" Store UV and geometry information for a primitve mesh
"""
def __init__(self,
vertices: np.array,
triangles: np.array,
triangle_uvs: np.array):
self.vertices = vertices
self.triangles = triangles
self.triangle_uvs = triangle_uvs
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(points)
@classmethod
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]])
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]])
triangle_uvs = get_triangle_uvs(vertices, triangles)
return TriangleMesh(vertices, triangles, triangle_uvs)
mesh = pcd.compute_convex_hull()[0]
@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])
return mesh
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_generator(
def mesh_factory(
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"],
create_uv_map=True,
map_texture_to_each_face=True,
depth=mesh_config[shape]["depth"]
)
elif "hemisphere" in shape:
prim_mesh = create_hemisphere(radius=mesh_config[shape]["radius"])
elif "sphere" in shape:
prim_mesh = o3d.geometry.TriangleMesh.create_sphere(
prim_mesh = TriangleMesh.create_hemisphere(
radius=mesh_config[shape]["radius"],
resolution=resolution,
create_uv_map=True,
resolution=resolution
)
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,
create_uv_map=True,
resolution=resolution
)
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(
np.asarray(prim_mesh.vertices)
* np.array(config["transform"]["scale"])
)
prim_mesh.scale(config["transform"]["scale"])
if "translate" in config["transform"]:
prim_mesh.translate(config["transform"]["translate"])