diff --git a/python/mujoco/usd/camera.py b/python/mujoco/usd/camera.py index b57f6108..e44ef585 100644 --- a/python/mujoco/usd/camera.py +++ b/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) diff --git a/python/mujoco/usd/demo.py b/python/mujoco/usd/demo.py index 442e33e3..e5aa2d71 100644 --- a/python/mujoco/usd/demo.py +++ b/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) diff --git a/python/mujoco/usd/exporter.py b/python/mujoco/usd/exporter.py index db4a0fc5..6eb8e9ce 100644 --- a/python/mujoco/usd/exporter.py +++ b/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.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 diff --git a/python/mujoco/usd/objects.py b/python/mujoco/usd/objects.py index b63db73e..2504495c 100644 --- a/python/mujoco/usd/objects.py +++ b/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.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] = { diff --git a/python/mujoco/usd/shapes.py b/python/mujoco/usd/shapes.py index bd0973b9..5c86b694 100644 --- a/python/mujoco/usd/shapes.py +++ b/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( - 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"])