From 34060fcc4925c2502c014f28fbc19ff750e61730 Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Sat, 20 Jul 2024 19:33:11 -0500 Subject: [PATCH 1/8] Removing Open3D dependency --- python/mujoco/usd/camera.py | 5 +- python/mujoco/usd/demo.py | 8 +- python/mujoco/usd/exporter.py | 14 ++- python/mujoco/usd/objects.py | 27 ++-- python/mujoco/usd/shapes.py | 229 ++++++++++++++++++++++++++++------ 5 files changed, 225 insertions(+), 58 deletions(-) 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"]) From 5c79313b211fc68f4667fc2e1ee6be41b0bb2e9d Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Sat, 20 Jul 2024 21:44:57 -0500 Subject: [PATCH 2/8] Implementing texture scaling for 2D textures --- python/mujoco/usd/exporter.py | 2 ++ python/mujoco/usd/objects.py | 25 +++++++++++++++---------- python/mujoco/usd/shapes.py | 2 +- 3 files changed, 18 insertions(+), 11 deletions(-) diff --git a/python/mujoco/usd/exporter.py b/python/mujoco/usd/exporter.py index 6eb8e9ce..180fa0bf 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -286,6 +286,7 @@ class USDExporter: usd_geom = object_module.USDTendon( mesh_config=mesh_config, stage=self.stage, + model=self.model, geom=geom, obj_name=geom_name, rgba=geom.rgba, @@ -301,6 +302,7 @@ class USDExporter: usd_geom = object_module.USDPrimitiveMesh( mesh_config=mesh_config, stage=self.stage, + model=self.model, geom=geom, obj_name=geom_name, rgba=geom.rgba, diff --git a/python/mujoco/usd/objects.py b/python/mujoco/usd/objects.py index 2504495c..b8a19ef8 100644 --- a/python/mujoco/usd/objects.py +++ b/python/mujoco/usd/objects.py @@ -53,12 +53,14 @@ class USDObject(abc.ABC): def __init__( self, stage: Usd.Stage, + model: mujoco.MjModel, geom: mujoco.MjvGeom, obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): self.stage = stage + self.model = model self.geom = geom self.obj_name = obj_name self.rgba = rgba @@ -227,9 +229,8 @@ class USDMesh(USDObject): rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): - super().__init__(stage, geom, obj_name, rgba, texture_file) + super().__init__(stage, model, geom, obj_name, rgba, texture_file) - self.model = model self.dataid = dataid mesh_path = f"{self.xform_path}/Mesh_{obj_name}" @@ -317,12 +318,13 @@ class USDPrimitiveMesh(USDObject): self, mesh_config: Dict[Any, Any], stage: Usd.Stage, + model: mujoco.MjModel, geom: mujoco.MjvGeom, obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): - super().__init__(stage, geom, obj_name, rgba, texture_file) + super().__init__(stage, model, geom, obj_name, rgba, texture_file) self.mesh_config = mesh_config self.prim_mesh = self.generate_primitive_mesh() @@ -363,17 +365,19 @@ class USDPrimitiveMesh(USDObject): def _get_uv_geometry(self): assert self.prim_mesh - # x_scale, y_scale = self.geom.texrepeat + s_scale, t_scale = self.model.mat_texrepeat[self.geom.matid] mesh_texcoord = np.array(self.prim_mesh.triangle_uvs) mesh_facetexcoord = np.asarray(self.prim_mesh.triangles) - # x_multiplier, y_multiplier = 1, 1 - # if self.geom.texuniform: - # x_multiplier, y_multiplier = self.geom.size[:2] + if self.model.mat_texuniform[self.geom.matid]: + if self.geom.size[0] > 0: + s_scale *= self.geom.size[0] + if self.geom.size[1] > 0: + t_scale *= self.geom.size[1] - # mesh_texcoord[:, 0] *= x_scale * x_multiplier - # mesh_texcoord[:, 1] *= y_scale * y_multiplier + mesh_texcoord[:, 0] *= s_scale / (self.geom.size[0] * 2) + mesh_texcoord[:, 1] *= t_scale / (self.geom.size[1] * 2) return mesh_texcoord, mesh_facetexcoord.flatten() @@ -394,12 +398,13 @@ class USDTendon(USDObject): self, mesh_config: Dict[Any, Any], stage: Usd.Stage, + model: mujoco.MjModel, geom: mujoco.MjvGeom, obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): - super().__init__(stage, geom, obj_name, rgba, texture_file) + super().__init__(stage, model, geom, obj_name, rgba, texture_file) self.mesh_config = mesh_config self.tendon_parts = self.generate_primitive_mesh() diff --git a/python/mujoco/usd/shapes.py b/python/mujoco/usd/shapes.py index 5c86b694..06c05026 100644 --- a/python/mujoco/usd/shapes.py +++ b/python/mujoco/usd/shapes.py @@ -282,7 +282,7 @@ def mesh_config_generator( def mesh_factory( mesh_config: Dict[str, Any], - resolution: int = 20, + resolution: int = 100, ): """Generates a mesh given a config consisting of shapes.""" assert "name" in mesh_config From e79aeb63067fafdf3d934dd5c98893cf941e9d76 Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Sat, 20 Jul 2024 22:16:01 -0500 Subject: [PATCH 3/8] Using correct imports --- python/mujoco/usd/camera.py | 3 +-- python/mujoco/usd/demo.py | 3 +-- python/mujoco/usd/exporter.py | 12 ++++-------- python/mujoco/usd/objects.py | 6 ++---- 4 files changed, 8 insertions(+), 16 deletions(-) diff --git a/python/mujoco/usd/camera.py b/python/mujoco/usd/camera.py index e44ef585..d175bc7d 100644 --- a/python/mujoco/usd/camera.py +++ b/python/mujoco/usd/camera.py @@ -14,8 +14,7 @@ # ============================================================================== """Camera handling for USD exporter.""" -# import mujoco.usd.utils as utils_modules -import utils as utils_module +import mujoco.usd.utils as utils_modules import numpy as np diff --git a/python/mujoco/usd/demo.py b/python/mujoco/usd/demo.py index e5aa2d71..0e520098 100644 --- a/python/mujoco/usd/demo.py +++ b/python/mujoco/usd/demo.py @@ -18,8 +18,7 @@ import argparse import pathlib import mujoco -# from mujoco.usd import exporter -import exporter +from mujoco.usd import exporter def generate_usd_trajectory(local_args): """Generates a USD file given the user arguments.""" diff --git a/python/mujoco/usd/exporter.py b/python/mujoco/usd/exporter.py index 180fa0bf..73b6c3a2 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -18,14 +18,10 @@ 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 camera as camera_module -import lights as light_module -import objects as object_module -import 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 numpy as np from PIL import Image as im from PIL import ImageOps diff --git a/python/mujoco/usd/objects.py b/python/mujoco/usd/objects.py index b8a19ef8..62d9df8e 100644 --- a/python/mujoco/usd/objects.py +++ b/python/mujoco/usd/objects.py @@ -19,10 +19,8 @@ import collections from typing import Optional, Dict, Any import mujoco -# 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 mujoco.usd.shapes as shapes_module +import mujoco.usd.utils as utils_modules import numpy as np From 5abdeb8a71f5fcdeed713cf0bd100781fd95083a Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Sat, 20 Jul 2024 22:17:49 -0500 Subject: [PATCH 4/8] Correcting utils_module import name --- python/mujoco/usd/objects.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/mujoco/usd/objects.py b/python/mujoco/usd/objects.py index 62d9df8e..72eb836f 100644 --- a/python/mujoco/usd/objects.py +++ b/python/mujoco/usd/objects.py @@ -20,7 +20,7 @@ from typing import Optional, Dict, Any import mujoco import mujoco.usd.shapes as shapes_module -import mujoco.usd.utils as utils_modules +import mujoco.usd.utils as utils_module import numpy as np From d1ff008a4f0a24de45d852247ac8a7779ba18ff8 Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Sun, 21 Jul 2024 11:14:00 -0500 Subject: [PATCH 5/8] Adding UV maps only when texture present --- python/mujoco/usd/exporter.py | 16 +++-- python/mujoco/usd/objects.py | 112 ++++++++++++++++++---------------- python/mujoco/usd/shapes.py | 62 +++++++++++++------ 3 files changed, 109 insertions(+), 81 deletions(-) diff --git a/python/mujoco/usd/exporter.py b/python/mujoco/usd/exporter.py index 73b6c3a2..832bce32 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -251,12 +251,10 @@ class USDExporter: assert geom_name not in self.geom_names - texture_file = ( - self.texture_files[ - self.model.mat_texid[geom.matid][mujoco.mjtTextureRole.mjTEXROLE_RGB] - ] - if geom.matid != -1 - else None + geom_textures = ( + [(self.texture_files[i], self.model.tex_type[i]) if i != -1 else None for i in self.model.mat_texid[geom.matid]] + if geom.matid != -1 + else None ) # handling meshes in our scene @@ -268,7 +266,7 @@ class USDExporter: obj_name=geom_name, dataid=self.model.geom_dataid[geom.objid], rgba=geom.rgba, - texture_file=texture_file, + geom_textures=geom_textures, ) else: # handling tendons in our scene @@ -286,7 +284,7 @@ class USDExporter: geom=geom, obj_name=geom_name, rgba=geom.rgba, - texture_file=texture_file, + geom_textures=geom_textures, ) # handling primitives in our scene else: @@ -302,7 +300,7 @@ class USDExporter: geom=geom, obj_name=geom_name, rgba=geom.rgba, - texture_file=texture_file, + geom_textures=geom_textures, ) self.geom_names.add(geom_name) diff --git a/python/mujoco/usd/objects.py b/python/mujoco/usd/objects.py index 72eb836f..1d68ccbb 100644 --- a/python/mujoco/usd/objects.py +++ b/python/mujoco/usd/objects.py @@ -16,7 +16,7 @@ import abc import collections -from typing import Optional, Dict, Any +from typing import Optional, Dict, Any, Tuple import mujoco import mujoco.usd.shapes as shapes_module @@ -55,14 +55,14 @@ class USDObject(abc.ABC): geom: mujoco.MjvGeom, obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), - texture_file: Optional[str] = None, + geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None ): self.stage = stage self.model = model self.geom = geom self.obj_name = obj_name self.rgba = rgba - self.texture_file = texture_file + self.geom_textures = geom_textures self.xform_path = f"/World/Mesh_Xform_{obj_name}" self.usd_xform = UsdGeom.Xform.Define(stage, self.xform_path) @@ -120,7 +120,7 @@ class USDObject(abc.ABC): # setting the image texture attributes image_shader.CreateIdAttr("UsdUVTexture") image_shader.CreateInput("file", Sdf.ValueTypeNames.Asset).Set( - self.texture_file + self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB][0] ) image_shader.CreateInput("sourceColorSpace", Sdf.ValueTypeNames.Token).Set( "sRGB" @@ -225,9 +225,9 @@ class USDMesh(USDObject): obj_name: str, dataid: int, rgba: np.ndarray = np.array([1, 1, 1, 1]), - texture_file: Optional[str] = None, + geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None ): - super().__init__(stage, model, geom, obj_name, rgba, texture_file) + super().__init__(stage, model, geom, obj_name, rgba, geom_textures) self.dataid = dataid @@ -243,15 +243,14 @@ class USDMesh(USDObject): ) self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face) - # setting mesh uv properties - mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry() - self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar( - "UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying - ) - self.texcoords.Set(mesh_texcoord) - self.texcoords.SetIndices(Vt.IntArray(mesh_facetexcoord.tolist())) - - if self.texture_file: + if self.geom_textures and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]: + # setting mesh uv properties + mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry() + self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar( + "UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying + ) + self.texcoords.Set(mesh_texcoord) + self.texcoords.SetIndices(Vt.IntArray(mesh_facetexcoord.tolist())) self.attach_image_material(self.usd_mesh) else: self.attach_solid_material(self.usd_mesh) @@ -320,9 +319,9 @@ class USDPrimitiveMesh(USDObject): geom: mujoco.MjvGeom, obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), - texture_file: Optional[str] = None, + geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None ): - super().__init__(stage, model, geom, obj_name, rgba, texture_file) + super().__init__(stage, model, geom, obj_name, rgba, geom_textures) self.mesh_config = mesh_config self.prim_mesh = self.generate_primitive_mesh() @@ -337,31 +336,32 @@ class USDPrimitiveMesh(USDObject): [3 for _ in range(mesh_facenum)] ) self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face) - - # setting mesh uv properties - mesh_texcoord, _ = self._get_uv_geometry() - self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar( - "UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying - ) - - self.texcoords.Set(mesh_texcoord) - self.texcoords.SetIndices(Vt.IntArray(list(range(mesh_facenum * 3)))) - self._set_refinement_properties(self.usd_prim) - if self.texture_file: + if self.geom_textures and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]: + # setting mesh uv properties + mesh_texcoord, _ = self._get_uv_geometry() + self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar( + "UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying + ) + + self.texcoords.Set(mesh_texcoord) + self.texcoords.SetIndices(Vt.IntArray(list(range(mesh_facenum * 3)))) + self.attach_image_material(self.usd_mesh) else: self.attach_solid_material(self.usd_mesh) def generate_primitive_mesh(self): """Generates the mesh for the primitive USD object.""" - _, prim_mesh = shapes_module.mesh_factory(self.mesh_config) + geom_rgb_texture = self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB] if self.geom_textures else None + texture_type = geom_rgb_texture[1] if geom_rgb_texture else None + _, prim_mesh = shapes_module.mesh_factory(self.mesh_config, texture_type) prim_mesh.translate(-prim_mesh.get_center()) return prim_mesh def _get_uv_geometry(self): - assert self.prim_mesh + assert self.prim_mesh and self.prim_mesh.triangle_uvs is not None s_scale, t_scale = self.model.mat_texrepeat[self.geom.matid] @@ -400,9 +400,9 @@ class USDTendon(USDObject): geom: mujoco.MjvGeom, obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), - texture_file: Optional[str] = None, + geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None ): - super().__init__(stage, model, geom, obj_name, rgba, texture_file) + super().__init__(stage, model, geom, obj_name, rgba, geom_textures) self.mesh_config = mesh_config self.tendon_parts = self.generate_primitive_mesh() @@ -432,33 +432,36 @@ class USDTendon(USDObject): part_geometry["mesh_face"] ) - # setting uv properties for each of the parts in the tendon - part_uv_geometries = self._get_uv_geometry() - for name, part_uv_geometry in part_uv_geometries.items(): - self.texcoords = UsdGeom.PrimvarsAPI( - self.usd_refs[name]["usd_mesh"] - ).CreatePrimvar( - "UVMap", - Sdf.ValueTypeNames.TexCoord2fArray, - UsdGeom.Tokens.faceVarying, - ) - self.texcoords.Set(part_uv_geometry["mesh_texcoord"]) - self.texcoords.SetIndices( - Vt.IntArray(list(range(part_geometry["mesh_facenum"] * 3))) - ) - - for _, ref in self.usd_refs.items(): - self._set_refinement_properties(ref["usd_prim"]) - if self.texture_file: - self.attach_image_material(ref["usd_mesh"]) - else: - self.attach_solid_material(ref["usd_mesh"]) + if self.geom_textures and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]: + # setting uv properties for each of the parts in the tendon + part_uv_geometries = self._get_uv_geometry() + for name, part_uv_geometry in part_uv_geometries.items(): + self.texcoords = UsdGeom.PrimvarsAPI( + self.usd_refs[name]["usd_mesh"] + ).CreatePrimvar( + "UVMap", + Sdf.ValueTypeNames.TexCoord2fArray, + UsdGeom.Tokens.faceVarying, + ) + self.texcoords.Set(part_uv_geometry["mesh_texcoord"]) + self.texcoords.SetIndices( + Vt.IntArray(list(range(part_geometry["mesh_facenum"] * 3))) + ) + for _, ref in self.usd_refs.items(): + self._set_refinement_properties(ref["usd_prim"]) + self.attach_image_material(ref["usd_mesh"]) + else: + for _, ref in self.usd_refs.items(): + self._set_refinement_properties(ref["usd_prim"]) + self.attach_solid_material(ref["usd_mesh"]) def generate_primitive_mesh(self): """Generates the tendon mesh using primitives.""" mesh_parts = {} + geom_rgb_texture = self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB] if self.geom_textures else None + texture_type = geom_rgb_texture[1] if geom_rgb_texture else None for part_config in self.mesh_config: - mesh_name, prim_mesh = shapes_module.mesh_factory(part_config) + mesh_name, prim_mesh = shapes_module.mesh_factory(part_config, texture_type) prim_mesh.translate(-prim_mesh.get_center()) mesh_parts[mesh_name] = prim_mesh return mesh_parts @@ -466,6 +469,7 @@ class USDTendon(USDObject): def _get_uv_geometry(self): part_uv_geometries = collections.defaultdict(dict) for name, mesh in self.tendon_parts.items(): + assert mesh.triangle_uvs is not None mesh_texcoord = np.array(mesh.triangle_uvs) mesh_facetexcoord = np.asarray(mesh.triangles) part_uv_geometries[name] = { diff --git a/python/mujoco/usd/shapes.py b/python/mujoco/usd/shapes.py index 06c05026..bc82b9a3 100644 --- a/python/mujoco/usd/shapes.py +++ b/python/mujoco/usd/shapes.py @@ -14,17 +14,24 @@ # ============================================================================== """Built-in shapes for USD exporter.""" -from typing import Dict, Any, Tuple +from typing import Dict, Any, Tuple, Optional import mujoco import numpy as np -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 get_triangle_uvs( + vertices: np.array, + triangles: np.array, + texture_type: Optional[mujoco.mjtTexture] +): + if texture_type == None: + return None + + # (jabhi) assuming 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 class TriangleMesh(): """ Store UV and geometry information for a primitve mesh @@ -39,7 +46,11 @@ class TriangleMesh(): @classmethod def create_box( - cls, width: float, height: float, depth: float + cls, + width: float, + height: float, + depth: float, + texture_type: Optional[mujoco.mjtTexture] ): vertices = np.array([[0.0, 0.0, 0.0], [width, 0.0, 0.0], @@ -63,13 +74,16 @@ class TriangleMesh(): [0, 4, 1], [1, 4, 5]]) - triangle_uvs = get_triangle_uvs(vertices, triangles) + triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type) return TriangleMesh(vertices, triangles, triangle_uvs) @classmethod def create_sphere( - cls, radius: float, resolution: int + cls, + radius: float, + texture_type: Optional[mujoco.mjtTexture], + resolution: int ): vertices = [] triangles = [] @@ -93,13 +107,16 @@ class TriangleMesh(): vertices = np.array(vertices) triangles = np.array(triangles) - triangle_uvs = get_triangle_uvs(vertices, triangles) + triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type) return TriangleMesh(vertices, triangles, triangle_uvs) @classmethod def create_hemisphere( - cls, radius: float, resolution: int + cls, + radius: float, + texture_type: Optional[mujoco.mjtTexture], + resolution: int ): vertices = [] triangles = [] @@ -128,13 +145,17 @@ class TriangleMesh(): vertices = np.array(vertices) triangles = np.array(triangles) - triangle_uvs = get_triangle_uvs(vertices, triangles) + triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type) return TriangleMesh(vertices, triangles, triangle_uvs) @classmethod def create_cylinder( - cls, radius: float, height: float, resolution: int + cls, + radius: float, + height: float, + texture_type: Optional[mujoco.mjtTexture], + resolution: int ): vertices = [] triangles = [] @@ -166,7 +187,7 @@ class TriangleMesh(): vertices = np.array(vertices) triangles = np.array(triangles) - triangle_uvs = get_triangle_uvs(vertices, triangles) + triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type) return TriangleMesh(vertices, triangles, triangle_uvs) @@ -189,7 +210,7 @@ class 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) + new_triangle_uvs = np.vstack((self.triangle_uvs, other.triangle_uvs)) return TriangleMesh(new_vertices, new_triangles, new_triangle_uvs) raise TypeError(f"Cannot add TriangleMesh with {type(other)}") @@ -282,7 +303,8 @@ def mesh_config_generator( def mesh_factory( mesh_config: Dict[str, Any], - resolution: int = 100, + texture_type: Optional[mujoco.mjtTexture], + resolution: int = 10, ): """Generates a mesh given a config consisting of shapes.""" assert "name" in mesh_config @@ -300,22 +322,26 @@ def mesh_factory( 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"], + texture_type=texture_type ) elif "hemisphere" in shape: prim_mesh = TriangleMesh.create_hemisphere( radius=mesh_config[shape]["radius"], + texture_type=texture_type, resolution=resolution ) elif "sphere" in shape: prim_mesh = TriangleMesh.create_sphere( radius=mesh_config[shape]["radius"], + texture_type=texture_type, resolution=resolution ) elif "cylinder" in shape: prim_mesh = TriangleMesh.create_cylinder( radius=mesh_config[shape]["radius"], height=mesh_config[shape]["height"], + texture_type=texture_type, resolution=resolution ) else: From d7447c4a842dedf95b631ed53f3bbf448773c053 Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Sun, 21 Jul 2024 17:00:16 -0500 Subject: [PATCH 6/8] Partial cube mapping implementation in USD --- python/mujoco/usd/exporter.py | 2 +- python/mujoco/usd/objects.py | 21 ++++++++------ python/mujoco/usd/shapes.py | 53 ++++++++++++++++++++++++++++++----- 3 files changed, 59 insertions(+), 17 deletions(-) diff --git a/python/mujoco/usd/exporter.py b/python/mujoco/usd/exporter.py index 832bce32..daa1fae4 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -209,7 +209,7 @@ class USDExporter: """Load textures.""" data_adr = 0 self.texture_files = [] - for texture_id in tqdm.tqdm(range(self.model.ntex)): + for texture_id in range(self.model.ntex): texture_height = self.model.tex_height[texture_id] texture_width = self.model.tex_width[texture_id] texture_nchannel = self.model.tex_nchannel[texture_id] diff --git a/python/mujoco/usd/objects.py b/python/mujoco/usd/objects.py index 1d68ccbb..79d8ac3f 100644 --- a/python/mujoco/usd/objects.py +++ b/python/mujoco/usd/objects.py @@ -363,19 +363,22 @@ class USDPrimitiveMesh(USDObject): def _get_uv_geometry(self): assert self.prim_mesh and self.prim_mesh.triangle_uvs is not None - s_scale, t_scale = self.model.mat_texrepeat[self.geom.matid] - mesh_texcoord = np.array(self.prim_mesh.triangle_uvs) mesh_facetexcoord = np.asarray(self.prim_mesh.triangles) - if self.model.mat_texuniform[self.geom.matid]: - if self.geom.size[0] > 0: - s_scale *= self.geom.size[0] - if self.geom.size[1] > 0: - t_scale *= self.geom.size[1] + geom_rgb_texture = self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB][1] - mesh_texcoord[:, 0] *= s_scale / (self.geom.size[0] * 2) - mesh_texcoord[:, 1] *= t_scale / (self.geom.size[1] * 2) + if geom_rgb_texture == mujoco.mjtTexture.mjTEXTURE_2D: + s_scale, t_scale = self.model.mat_texrepeat[self.geom.matid] + + if self.model.mat_texuniform[self.geom.matid]: + if self.geom.size[0] > 0: + s_scale *= self.geom.size[0] + if self.geom.size[1] > 0: + t_scale *= self.geom.size[1] + + mesh_texcoord[:, 0] *= s_scale / (self.geom.size[0] * 2) + mesh_texcoord[:, 1] *= t_scale / (self.geom.size[1] * 2) return mesh_texcoord, mesh_facetexcoord.flatten() diff --git a/python/mujoco/usd/shapes.py b/python/mujoco/usd/shapes.py index bc82b9a3..28b340f8 100644 --- a/python/mujoco/usd/shapes.py +++ b/python/mujoco/usd/shapes.py @@ -27,11 +27,48 @@ def get_triangle_uvs( if texture_type == None: return None - # (jabhi) assuming 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 + triangle_uvs = [] + if texture_type == mujoco.mjtTexture.mjTEXTURE_2D: + triangle_uvs = [[vertices[i][0], vertices[i][1]] for i in np.nditer(triangles)] + + elif texture_type == mujoco.mjtTexture.mjTEXTURE_CUBE: + center = np.mean(vertices, axis=0) + for vertex_id in np.nditer(triangles): + x, y, z = vertices[vertex_id] - center + + abs_x, abs_y, abs_z = abs(x), abs(y), abs(z) + + if x > 0 and abs_x >= abs_y and abs_x >= abs_z: + u = -z / abs_x + v = y / abs_x + elif x <= 0 and abs_x >= abs_y and abs_x >= abs_z: + u = z / abs_x + v = y / abs_x + elif y > 0 and abs_y >= abs_x and abs_y >= abs_z: + u = x / abs_y + v = -z / abs_y + elif y <= 0 and abs_y >= abs_x and abs_y >= abs_z: + u = x / abs_y + v = z / abs_y + elif z > 0 and abs_z >= abs_x and abs_z >= abs_y: + u = x / abs_z + v = y / abs_z + elif z <= 0 and abs_z >= abs_x and abs_z >= abs_y: + u = -x / abs_z + v = y / abs_z + + u = (u + 1.0) / 2.0 + v = (v + 1.0) / 2.0 + v /= 6 + + assert 0 <= u and u <= 1 and 0 <= v and v <= 1 + + triangle_uvs.append([u, v]) + elif texture_type == mujoco.mjtTexture.mjTEXTURE_SKYBOX: + # defaults to 2D mapping temporarily + triangle_uvs = [[vertices[i][0], vertices[i][1]] for i in np.nditer(triangles)] + + return np.array(triangle_uvs) class TriangleMesh(): """ Store UV and geometry information for a primitve mesh @@ -210,7 +247,9 @@ class 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 = np.vstack((self.triangle_uvs, other.triangle_uvs)) + new_triangle_uvs = None + if self.triangle_uvs is not None: + new_triangle_uvs = np.vstack((self.triangle_uvs, other.triangle_uvs)) return TriangleMesh(new_vertices, new_triangles, new_triangle_uvs) raise TypeError(f"Cannot add TriangleMesh with {type(other)}") @@ -304,7 +343,7 @@ def mesh_config_generator( def mesh_factory( mesh_config: Dict[str, Any], texture_type: Optional[mujoco.mjtTexture], - resolution: int = 10, + resolution: int = 100, ): """Generates a mesh given a config consisting of shapes.""" assert "name" in mesh_config From a470f26a4ba8e78ab9f5bc581a7866c35a366551 Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Tue, 30 Jul 2024 15:30:16 -0500 Subject: [PATCH 7/8] Reverting to Union for python 3.9 --- python/mujoco/usd/shapes.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/mujoco/usd/shapes.py b/python/mujoco/usd/shapes.py index 28b340f8..63592c60 100644 --- a/python/mujoco/usd/shapes.py +++ b/python/mujoco/usd/shapes.py @@ -14,7 +14,7 @@ # ============================================================================== """Built-in shapes for USD exporter.""" -from typing import Dict, Any, Tuple, Optional +from typing import Dict, Any, Tuple, Optional, Union import mujoco import numpy as np @@ -270,7 +270,7 @@ def decouple_config(config: Dict[str, Any]): def mesh_config_generator( name: str, - geom_type: int | mujoco.mjtGeom, + geom_type: Union[int, mujoco.mjtGeom], size: np.ndarray, decouple: bool = False, ): From e0483c2f8654dbb4c1620139be0c45efb0c3ee28 Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Fri, 2 Aug 2024 15:13:25 -0500 Subject: [PATCH 8/8] Updating for PR comments --- python/mujoco/usd/camera.py | 4 ++-- python/mujoco/usd/exporter.py | 9 ++++----- python/mujoco/usd/objects.py | 16 ++++++++-------- python/mujoco/usd/shapes.py | 25 ++++++++++++------------- 4 files changed, 26 insertions(+), 28 deletions(-) diff --git a/python/mujoco/usd/camera.py b/python/mujoco/usd/camera.py index d175bc7d..a43cb3e2 100644 --- a/python/mujoco/usd/camera.py +++ b/python/mujoco/usd/camera.py @@ -14,7 +14,7 @@ # ============================================================================== """Camera handling for USD exporter.""" -import mujoco.usd.utils as utils_modules +import mujoco.usd.utils as utils_module import numpy as np @@ -49,7 +49,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_modules.create_transform_matrix( + transformation_mat = utils_module.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/exporter.py b/python/mujoco/usd/exporter.py index daa1fae4..5a94f47d 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -251,11 +251,10 @@ class USDExporter: assert geom_name not in self.geom_names - geom_textures = ( - [(self.texture_files[i], self.model.tex_type[i]) if i != -1 else None for i in self.model.mat_texid[geom.matid]] - if geom.matid != -1 - else None - ) + if geom.matid == -1: + geom_textures = [] + else: + geom_textures = [(self.texture_files[i], self.model.tex_type[i]) if i != -1 else None for i in self.model.mat_texid[geom.matid]] # handling meshes in our scene if geom.type == mujoco.mjtGeom.mjGEOM_MESH: diff --git a/python/mujoco/usd/objects.py b/python/mujoco/usd/objects.py index 79d8ac3f..13f19a59 100644 --- a/python/mujoco/usd/objects.py +++ b/python/mujoco/usd/objects.py @@ -16,7 +16,7 @@ import abc import collections -from typing import Optional, Dict, Any, Tuple +from typing import Any, Dict, List, Optional, Tuple import mujoco import mujoco.usd.shapes as shapes_module @@ -55,7 +55,7 @@ class USDObject(abc.ABC): geom: mujoco.MjvGeom, obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), - geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None + geom_textures: List[Optional[Tuple[str, mujoco.mjtTexture]]] = None ): self.stage = stage self.model = model @@ -225,7 +225,7 @@ class USDMesh(USDObject): obj_name: str, dataid: int, rgba: np.ndarray = np.array([1, 1, 1, 1]), - geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None + geom_textures: List[Optional[Tuple[str, mujoco.mjtTexture]]] = None ): super().__init__(stage, model, geom, obj_name, rgba, geom_textures) @@ -243,7 +243,7 @@ class USDMesh(USDObject): ) self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face) - if self.geom_textures and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]: + if geom.matid != -1 and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]: # setting mesh uv properties mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry() self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar( @@ -319,7 +319,7 @@ class USDPrimitiveMesh(USDObject): geom: mujoco.MjvGeom, obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), - geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None + geom_textures: List[Optional[Tuple[str, mujoco.mjtTexture]]] = None ): super().__init__(stage, model, geom, obj_name, rgba, geom_textures) @@ -338,7 +338,7 @@ class USDPrimitiveMesh(USDObject): self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face) self._set_refinement_properties(self.usd_prim) - if self.geom_textures and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]: + if geom.matid != -1 and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]: # setting mesh uv properties mesh_texcoord, _ = self._get_uv_geometry() self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar( @@ -403,7 +403,7 @@ class USDTendon(USDObject): geom: mujoco.MjvGeom, obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), - geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None + geom_textures: List[Optional[Tuple[str, mujoco.mjtTexture]]] = None ): super().__init__(stage, model, geom, obj_name, rgba, geom_textures) @@ -435,7 +435,7 @@ class USDTendon(USDObject): part_geometry["mesh_face"] ) - if self.geom_textures and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]: + if geom.matid != -1 and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]: # setting uv properties for each of the parts in the tendon part_uv_geometries = self._get_uv_geometry() for name, part_uv_geometry in part_uv_geometries.items(): diff --git a/python/mujoco/usd/shapes.py b/python/mujoco/usd/shapes.py index 63592c60..4f135790 100644 --- a/python/mujoco/usd/shapes.py +++ b/python/mujoco/usd/shapes.py @@ -14,14 +14,14 @@ # ============================================================================== """Built-in shapes for USD exporter.""" -from typing import Dict, Any, Tuple, Optional, Union +from typing import Any, Dict, Optional, Tuple, Union import mujoco import numpy as np def get_triangle_uvs( - vertices: np.array, - triangles: np.array, + vertices: np.ndarray, + triangles: np.ndarray, texture_type: Optional[mujoco.mjtTexture] ): if texture_type == None: @@ -70,13 +70,12 @@ def get_triangle_uvs( return np.array(triangle_uvs) -class TriangleMesh(): - """ Store UV and geometry information for a primitve mesh - """ +class TriangleMesh: + """Store UV and geometry information for a primitive mesh.""" def __init__(self, - vertices: np.array, - triangles: np.array, - triangle_uvs: np.array): + vertices: np.ndarray, + triangles: np.ndarray, + triangle_uvs: np.ndarray): self.vertices = vertices self.triangles = triangles self.triangle_uvs = triangle_uvs @@ -88,7 +87,7 @@ class TriangleMesh(): height: float, depth: float, texture_type: Optional[mujoco.mjtTexture] - ): + ) -> 'TriangleMesh': vertices = np.array([[0.0, 0.0, 0.0], [width, 0.0, 0.0], [0.0, 0.0, depth], @@ -121,7 +120,7 @@ class TriangleMesh(): radius: float, texture_type: Optional[mujoco.mjtTexture], resolution: int - ): + ) -> 'TriangleMesh': vertices = [] triangles = [] for i in range(2*resolution + 1): @@ -154,7 +153,7 @@ class TriangleMesh(): radius: float, texture_type: Optional[mujoco.mjtTexture], resolution: int - ): + ) -> 'TriangleMesh': vertices = [] triangles = [] for i in range(resolution + 1): @@ -193,7 +192,7 @@ class TriangleMesh(): height: float, texture_type: Optional[mujoco.mjtTexture], resolution: int - ): + ) -> 'TriangleMesh': vertices = [] triangles = []