diff --git a/python/mujoco/usd/camera.py b/python/mujoco/usd/camera.py index b57f6108..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_component +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_component.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/demo.py b/python/mujoco/usd/demo.py index 442e33e3..6a2df9f3 100644 --- a/python/mujoco/usd/demo.py +++ b/python/mujoco/usd/demo.py @@ -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..5b95e47b 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -27,7 +27,6 @@ from PIL import Image as im from PIL import ImageOps import scipy import termcolor -import tqdm # TODO: b/288149332 - Remove once USD Python Binding works well with pytype. # pytype: disable=module-attr @@ -209,7 +208,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] @@ -251,13 +250,13 @@ class USDExporter: assert geom_name not in self.geom_names - texture_file = ( - self.texture_files[ - self.model.mat_texid[geom.matid][mujoco.mjTEXROLE_RGB] - ] - 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: @@ -268,7 +267,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 @@ -282,10 +281,11 @@ 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, - texture_file=texture_file, + geom_textures=geom_textures, ) # handling primitives in our scene else: @@ -297,10 +297,11 @@ 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, - 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 b63db73e..850e3645 100644 --- a/python/mujoco/usd/objects.py +++ b/python/mujoco/usd/objects.py @@ -16,13 +16,14 @@ import abc import collections -from typing import Optional, Dict, Any +from typing import Any, Dict, List, Optional, Tuple 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_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 @@ -50,16 +51,18 @@ 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, + geom_textures: List[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) @@ -117,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" @@ -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) @@ -222,11 +225,10 @@ class USDMesh(USDObject): obj_name: str, dataid: int, rgba: np.ndarray = np.array([1, 1, 1, 1]), - texture_file: Optional[str] = None, + geom_textures: List[Optional[Tuple[str, mujoco.mjtTexture]]] = None ): - super().__init__(stage, geom, obj_name, rgba, texture_file) + super().__init__(stage, model, geom, obj_name, rgba, geom_textures) - self.model = model self.dataid = dataid mesh_path = f"{self.xform_path}/Mesh_{obj_name}" @@ -241,15 +243,19 @@ 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 ( + 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( + "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) @@ -314,12 +320,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, + geom_textures: List[Optional[Tuple[str, mujoco.mjtTexture]]] = None ): - super().__init__(stage, 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() @@ -334,42 +341,66 @@ 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 ( + 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( + "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_component.mesh_generator(self.mesh_config) + tex_role = mujoco.mjtTextureRole + geom_rgb_texture = ( + self.geom_textures[tex_role.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 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. + tex_role = mujoco.mjtTextureRole + geom_rgb_texture = self.geom_textures[tex_role.mjTEXROLE_RGB][1] + + 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() 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) @@ -383,12 +414,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, + geom_textures: List[Optional[Tuple[str, mujoco.mjtTexture]]] = None ): - super().__init__(stage, 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() @@ -407,6 +439,7 @@ class USDTendon(USDObject): # setting mesh geometry properties for each of the parts in the tendon part_geometries = self._get_mesh_geometry() + part_geometry = None for name, part_geometry in part_geometries.items(): self.usd_refs[name]["usd_mesh"].GetPointsAttr().Set( part_geometry["mesh_vert"] @@ -418,33 +451,43 @@ 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: + tex_role = mujoco.mjtTextureRole + if geom.matid != -1 and self.geom_textures[tex_role.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_component.mesh_generator(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 @@ -452,6 +495,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] = { @@ -463,7 +507,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..ca29b276 100644 --- a/python/mujoco/usd/shapes.py +++ b/python/mujoco/usd/shapes.py @@ -14,33 +14,270 @@ # ============================================================================== """Built-in shapes for USD exporter.""" -from typing import Dict, Any +from typing import Any, Dict, Optional, Tuple, Union import mujoco import numpy as np -from open3d import open3d as o3d -def create_hemisphere( - radius: float, theta_steps: int = 50, phi_steps: int = 50 +def get_triangle_uvs( + vertices: np.ndarray, + triangles: np.ndarray, + texture_type: Optional[mujoco.mjtTexture] ): - """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]) + """Returns UV coordinates for a given mesh.""" + if not texture_type: + return None - pcd = o3d.geometry.PointCloud() - pcd.points = o3d.utility.Vector3dVector(points) + triangle_uvs = [] + if texture_type == mujoco.mjtTexture.mjTEXTURE_2D: + triangle_uvs = [ + [vertices[i][0], vertices[i][1]] for i in np.nditer(triangles) + ] - mesh = pcd.compute_convex_hull()[0] + 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 - return mesh + abs_x, abs_y, abs_z = abs(x), abs(y), abs(z) + u = 0 + v = 0 + + 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 primitive mesh. + + Attributes: + vertices: A numpy array of vertices. + triangles: A numpy array of triangles. + triangle_uvs: A numpy array of UV coordinates. + """ + + def __init__(self, + vertices: np.ndarray, + triangles: np.ndarray, + triangle_uvs: np.ndarray): + """Creates a TriangleMesh object. + + Args: + vertices: A numpy array of vertices. + triangles: A numpy array of triangles. + triangle_uvs: A numpy array of UV coordinates. + """ + self.vertices = vertices + self.triangles = triangles + self.triangle_uvs = triangle_uvs + + @classmethod + def create_box( + cls, + width: float, + height: float, + depth: float, + texture_type: Optional[mujoco.mjtTexture] + ) -> TriangleMesh: + """Creates a box.""" + 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, texture_type) + + return TriangleMesh(vertices, triangles, triangle_uvs) + + @classmethod + def create_sphere( + cls, + radius: float, + texture_type: Optional[mujoco.mjtTexture], + resolution: int + ) -> TriangleMesh: + """Creates a sphere.""" + 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, texture_type) + + return TriangleMesh(vertices, triangles, triangle_uvs) + + @classmethod + def create_hemisphere( + cls, + radius: float, + texture_type: Optional[mujoco.mjtTexture], + resolution: int, + ) -> TriangleMesh: + """Creates a hemisphere.""" + 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, texture_type) + + return TriangleMesh(vertices, triangles, triangle_uvs) + + @classmethod + def create_cylinder( + cls, + radius: float, + height: float, + texture_type: Optional[mujoco.mjtTexture], + resolution: int + ) -> TriangleMesh: + """Creates a cylinder.""" + 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, texture_type) + + 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 = 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)}") def decouple_config(config: Dict[str, Any]): @@ -60,7 +297,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, ): @@ -79,20 +316,22 @@ 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( @@ -131,8 +370,9 @@ def mesh_config_generator( return config -def mesh_generator( +def mesh_factory( mesh_config: Dict[str, Any], + texture_type: Optional[mujoco.mjtTexture], resolution: int = 100, ): """Generates a mesh given a config consisting of shapes.""" @@ -146,27 +386,30 @@ def mesh_generator( continue 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, + texture_type=texture_type ) 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, + 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 = 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, + texture_type=texture_type, + resolution=resolution ) else: raise ValueError("Shape not supported") @@ -182,10 +425,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"])