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: