From e0483c2f8654dbb4c1620139be0c45efb0c3ee28 Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Fri, 2 Aug 2024 15:13:25 -0500 Subject: [PATCH] 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 = []