From 58dc62c4ecfdbe500a23c83e639a22e8f0f078ab Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Tue, 13 Aug 2024 11:09:02 -0700 Subject: [PATCH 01/12] Fixing USD pipeline after copybara changes --- python/mujoco/usd/objects.py | 10 ++++++++-- python/mujoco/usd/shapes.py | 10 +++++----- python/mujoco/usd/utils.py | 1 + 3 files changed, 14 insertions(+), 7 deletions(-) diff --git a/python/mujoco/usd/objects.py b/python/mujoco/usd/objects.py index 850e3645..f74eae6b 100644 --- a/python/mujoco/usd/objects.py +++ b/python/mujoco/usd/objects.py @@ -392,8 +392,14 @@ class USDPrimitiveMesh(USDObject): 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) + s_size, t_size = self.geom.size[:2] + if self.geom.type == mujoco.mjtGeom.mjGEOM_PLANE: + s_size = s_size if s_size > 0 else 1 + t_size = t_size if t_size > 0 else 1 + + if self.model.mat_texuniform[self.geom.matid]: + mesh_texcoord[:, 0] *= s_scale / (s_size * 2) + mesh_texcoord[:, 1] *= t_scale / (t_size * 2) return mesh_texcoord, mesh_facetexcoord.flatten() diff --git a/python/mujoco/usd/shapes.py b/python/mujoco/usd/shapes.py index ca29b276..ceaea43b 100644 --- a/python/mujoco/usd/shapes.py +++ b/python/mujoco/usd/shapes.py @@ -26,7 +26,7 @@ def get_triangle_uvs( texture_type: Optional[mujoco.mjtTexture] ): """Returns UV coordinates for a given mesh.""" - if not texture_type: + if texture_type is None: return None triangle_uvs = [] @@ -110,7 +110,7 @@ class TriangleMesh: height: float, depth: float, texture_type: Optional[mujoco.mjtTexture] - ) -> TriangleMesh: + ) -> "TriangleMesh": """Creates a box.""" vertices = np.array([[0.0, 0.0, 0.0], [width, 0.0, 0.0], @@ -144,7 +144,7 @@ class TriangleMesh: radius: float, texture_type: Optional[mujoco.mjtTexture], resolution: int - ) -> TriangleMesh: + ) -> "TriangleMesh": """Creates a sphere.""" vertices = [] triangles = [] @@ -178,7 +178,7 @@ class TriangleMesh: radius: float, texture_type: Optional[mujoco.mjtTexture], resolution: int, - ) -> TriangleMesh: + ) -> "TriangleMesh": """Creates a hemisphere.""" vertices = [] triangles = [] @@ -218,7 +218,7 @@ class TriangleMesh: height: float, texture_type: Optional[mujoco.mjtTexture], resolution: int - ) -> TriangleMesh: + ) -> "TriangleMesh": """Creates a cylinder.""" vertices = [] triangles = [] diff --git a/python/mujoco/usd/utils.py b/python/mujoco/usd/utils.py index 1e42afa2..ce3402ff 100644 --- a/python/mujoco/usd/utils.py +++ b/python/mujoco/usd/utils.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== +"""Utility functions for USD exporter.""" import numpy as np From 39cbe5d5b0ad3f53ae140a505c86ee13a37fee2a Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Fri, 16 Aug 2024 12:06:25 -0500 Subject: [PATCH 02/12] Updating USD golden file and exporter_test.py --- python/mujoco/testdata/usd_golden.usda | 48 ++++++++++++++++++-------- python/mujoco/usd/demo.py | 5 ++- python/mujoco/usd/exporter.py | 1 - python/mujoco/usd/exporter_test.py | 34 +++++++++++------- python/mujoco/usd/shapes.py | 6 ++-- python/pyproject.toml | 1 + 6 files changed, 62 insertions(+), 33 deletions(-) diff --git a/python/mujoco/testdata/usd_golden.usda b/python/mujoco/testdata/usd_golden.usda index 1a154486..0c1c08b1 100644 --- a/python/mujoco/testdata/usd_golden.usda +++ b/python/mujoco/testdata/usd_golden.usda @@ -5,7 +5,7 @@ """ endTimeCode = 1 startTimeCode = 0 - timeCodesPerSecond = 24 + timeCodesPerSecond = 60 upAxis = "Z" ) @@ -13,48 +13,66 @@ def Xform "World" { def Xform "Light_Xform_0" { - double3 xformOp:translate + double3 xformOp:translate.timeSamples = { + 0: (0, -3.674234628677368, 3.674234628677368), + } uniform token[] xformOpOrder = ["xformOp:translate"] def SphereLight "Light_0" { + color3f inputs:color = (0.4, 0.4, 0.4) + float inputs:intensity = 10000 bool inputs:normalize = 1 float inputs:radius = 0.3 bool treatAsPoint = 0 } } - def Xform "CubeMesh_Xform_white_box" + def Xform "Camera_Xform_closeup" { matrix4d xformOp:transform.timeSamples = { - 0: ( (0, 0, 0, 0), (0, 0, 0, 0), (0, 0, 0, 0), (0, 0, 0, 1) ), + 0: ( (1, 0, -0, 0), (0, 0.009999499656260014, 0.9999499917030334, 0), (0, -0.9999499917030334, 0.009999499656260014, 0), (0, -6, 0, 1) ), } uniform token[] xformOpOrder = ["xformOp:transform"] - def Mesh "CubeMesh_white_box" ( + def Camera "Camera_closeup" + { + float2 clippingRange = (0.0001, 1000000) + float focalLength = 12 + float focusDistance = 400 + float horizontalAperture = 12 + } + } + + def Xform "Mesh_Xform_white_box_id0_geom" + { + token visibility.timeSamples = { + 0: "inherited", + 1: "invisible", + } + float3 xformOp:scale + matrix4d xformOp:transform.timeSamples = { + 0: ( (1, 0, 0, 0), (0, 1, 0, 0), (0, 0, 1, 0), (0, 0, 0, 1) ), + } + uniform token[] xformOpOrder = ["xformOp:transform", "xformOp:scale"] + + def Mesh "Mesh_white_box_id0_geom" ( apiSchemas = ["MaterialBindingAPI"] ) { int[] faceVertexCounts = [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3] int[] faceVertexIndices = [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] - rel material:binding = + rel material:binding = point3f[] points = [(-1, -1, -1), (1, -1, -1), (-1, -1, 1), (1, -1, 1), (-1, 1, -1), (1, 1, -1), (-1, 1, 1), (1, 1, 1)] - texCoord2f[] primvars:UVMap = [] ( - interpolation = "faceVarying" - ) - int[] primvars:UVMap:indices = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35] uniform token subdivisionScheme = "none" - token visibility.timeSamples = { - 0: "inherited", - } } } def "_materials" { - def Material "Material_white_box" + def Material "Material_white_box_id0_geom" { - token outputs:surface.connect = + token outputs:surface.connect = def Shader "Principled_BSDF" { diff --git a/python/mujoco/usd/demo.py b/python/mujoco/usd/demo.py index 6a2df9f3..d378df4e 100644 --- a/python/mujoco/usd/demo.py +++ b/python/mujoco/usd/demo.py @@ -79,7 +79,10 @@ if __name__ == '__main__': ) parser.add_argument( - '--camera_names', type=str, nargs='+', help='cameras to include in usd' + '--camera_names', + type=str, + nargs='+', + help='cameras to include in usd' ) parser.add_argument( diff --git a/python/mujoco/usd/exporter.py b/python/mujoco/usd/exporter.py index 5b95e47b..7879b997 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -194,7 +194,6 @@ class USDExporter: if self.updates == 0: self._initialize_usd_stage() - self._load_lights() self._load_cameras() diff --git a/python/mujoco/usd/exporter_test.py b/python/mujoco/usd/exporter_test.py index 01424bbf..5b1ba255 100644 --- a/python/mujoco/usd/exporter_test.py +++ b/python/mujoco/usd/exporter_test.py @@ -22,12 +22,12 @@ from etils import epath import mujoco -# Open3D and USD are not fully supported on all MuJoCo architectures. +# USD is not fully supported on all MuJoCo architectures. execute_test = True try: from mujoco.usd import exporter as exporter_module # pylint: disable=g-import-not-at-top except ImportError: - logging.warning('Skipping test due to missing import') + logging.warning("Skipping test due to missing import") execute_test = False @@ -35,9 +35,10 @@ class ExporterTest(absltest.TestCase): def test_usd_export(self): if not execute_test: - return + self.fail("Test test_usd_export failed to execute.") - output_dir = os.getenv('TEST_UNDECLARED_OUTPUTS_DIR') + output_dir_root = os.getenv("TEST_UNDECLARED_OUTPUTS_DIR") + output_dir_name = "usd_test" xml = """ @@ -50,19 +51,26 @@ class ExporterTest(absltest.TestCase): data = mujoco.MjData(model) exporter = exporter_module.USDExporter( model, - output_directory_name='mujoco_usdpkg', - output_directory_root=output_dir, + output_directory_name=output_dir_name, + output_directory_root=output_dir_root, + camera_names=["closeup"] ) + mujoco.mj_step(model, data) exporter.update_scene(data) - exporter.save_scene('export.usda') + exporter.save_scene("usda") with open(os.path.join( - output_dir, 'mujoco_usdpkg/frames', 'frame_1.export.usda'), 'r') as f: + output_dir_root, + f"{output_dir_name}/frames", + "frame_1.usda"), "r", encoding="utf-8" + ) as f: golden_path = os.path.join( - epath.resource_path('mujoco'), 'testdata', 'usd_golden.usda') - with open(golden_path, 'r') as golden_file: - self.assertEqual(f.read(), golden_file.read()) + epath.resource_path("mujoco"), + "testdata", + "usd_golden.usda" + ) + with open(golden_path, "r", encoding="utf-8") as golden_file: + self.assertEqual(f.readlines(), golden_file.readlines()) - -if __name__ == '__main__': +if __name__ == "__main__": absltest.main() diff --git a/python/mujoco/usd/shapes.py b/python/mujoco/usd/shapes.py index ceaea43b..48b91771 100644 --- a/python/mujoco/usd/shapes.py +++ b/python/mujoco/usd/shapes.py @@ -254,14 +254,14 @@ class TriangleMesh: return TriangleMesh(vertices, triangles, triangle_uvs) - def translate(self, translation: np.array): + def translate(self, translation: np.ndarray) -> None: self.vertices = self.vertices + translation - def rotate(self, rotation: np.array, center: Tuple[float, ...]): + def rotate(self, rotation: np.ndarray, center: Tuple[float, ...]) -> None: translated_point = self.vertices - center self.vertices = np.dot(translated_point, rotation) + center - def scale(self, scale: np.array): + def scale(self, scale: np.ndarray) -> None: self.vertices = self.vertices * scale def get_center(self): diff --git a/python/pyproject.toml b/python/pyproject.toml index 02f9721a..46d6e734 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -58,4 +58,5 @@ mujoco = [ "include/mujoco/*.h", "testdata/*.xml", "testdata/*.msh", + "testdata/*.usda", ] From 492cb7d07896c88a48fb2e14f0ae7eaaac81d863 Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Mon, 19 Aug 2024 11:19:18 -0500 Subject: [PATCH 03/12] Remove try catch in exporter test --- python/mujoco/usd/exporter_test.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/python/mujoco/usd/exporter_test.py b/python/mujoco/usd/exporter_test.py index 5b1ba255..9c5f8043 100644 --- a/python/mujoco/usd/exporter_test.py +++ b/python/mujoco/usd/exporter_test.py @@ -20,22 +20,11 @@ import os from absl.testing import absltest from etils import epath import mujoco - - -# USD is not fully supported on all MuJoCo architectures. -execute_test = True -try: - from mujoco.usd import exporter as exporter_module # pylint: disable=g-import-not-at-top -except ImportError: - logging.warning("Skipping test due to missing import") - execute_test = False - +from mujoco.usd import exporter as exporter_module # pylint: disable=g-import-not-at-top class ExporterTest(absltest.TestCase): def test_usd_export(self): - if not execute_test: - self.fail("Test test_usd_export failed to execute.") output_dir_root = os.getenv("TEST_UNDECLARED_OUTPUTS_DIR") output_dir_name = "usd_test" From f358cf664b08a51d8552b54aa46434b5a7541815 Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Mon, 19 Aug 2024 11:23:43 -0500 Subject: [PATCH 04/12] Removing unused logging import --- python/mujoco/usd/exporter_test.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/mujoco/usd/exporter_test.py b/python/mujoco/usd/exporter_test.py index 9c5f8043..6a61628f 100644 --- a/python/mujoco/usd/exporter_test.py +++ b/python/mujoco/usd/exporter_test.py @@ -14,7 +14,6 @@ # ============================================================================== """Tests for the MuJoCo USD Exporter.""" -import logging import os from absl.testing import absltest From 85d19e189bffd2c53621689c24f2edac657fad24 Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Tue, 20 Aug 2024 21:39:24 -0500 Subject: [PATCH 05/12] Adding usd-core optional dependency --- python/build_requirements.txt | 4 ++++ python/pyproject.toml | 3 +++ 2 files changed, 7 insertions(+) diff --git a/python/build_requirements.txt b/python/build_requirements.txt index 46a39e32..8261c954 100644 --- a/python/build_requirements.txt +++ b/python/build_requirements.txt @@ -91,3 +91,7 @@ pluggy==1.3.0 \ --hash=sha256:d89c696a773f8bd377d18e5ecda92b7a3793cbe66c87060a6fb58c7b6e1061f7 pyparsing==3.1.1 \ --hash=sha256:32c7c0b711493c72ff18a981d24f28aaf9c1fb7ed5e9667c9e84e3db623bdbfb + +# Dependencies for USD +usd-core==24.8 \ + --hash=sha256:39fe8e266875e1105886cab870df4bcbe2d40a84696177bc574c22b96d843bf6 \ No newline at end of file diff --git a/python/pyproject.toml b/python/pyproject.toml index 46d6e734..7d0cb8bb 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -60,3 +60,6 @@ mujoco = [ "testdata/*.msh", "testdata/*.usda", ] + +[project.optional-dependencies] +usd = ["usd-core"] From 83c88141497ca518119d2906ddddbdb5e74fa14e Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Fri, 30 Aug 2024 12:14:26 -0400 Subject: [PATCH 06/12] Adding new usd-core hashes --- python/build_requirements.txt | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/build_requirements.txt b/python/build_requirements.txt index 1d02e327..30ae9eee 100644 --- a/python/build_requirements.txt +++ b/python/build_requirements.txt @@ -95,4 +95,7 @@ pyparsing==3.1.1 \ # Dependencies for USD usd-core==24.8 \ - --hash=sha256:39fe8e266875e1105886cab870df4bcbe2d40a84696177bc574c22b96d843bf6 \ No newline at end of file + --hash=sha256:39fe8e266875e1105886cab870df4bcbe2d40a84696177bc574c22b96d843bf6 + --hash=sha256:8b38b347dce9336d00dacd4c0e5813bcb43c61e165c653d754e3ee3dc4b2b715 + --hash=sha256:d1dfe295ccbc57cac39e6dee2d7ce831d7a6a504b8a17dc974e631e3521f83a6 + --hash=sha256:d3e06fd8c953c4de3d24591cdb9e8e65ca28a876519378d1e4a5d9bca419d7c1 \ No newline at end of file From 4a8b54c7ecf4ee6b57a13a57b847bbde3f45ae02 Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Fri, 30 Aug 2024 12:16:16 -0400 Subject: [PATCH 07/12] Adding backslash char --- python/build_requirements.txt | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/build_requirements.txt b/python/build_requirements.txt index 30ae9eee..2e5e1ad9 100644 --- a/python/build_requirements.txt +++ b/python/build_requirements.txt @@ -95,7 +95,7 @@ pyparsing==3.1.1 \ # Dependencies for USD usd-core==24.8 \ - --hash=sha256:39fe8e266875e1105886cab870df4bcbe2d40a84696177bc574c22b96d843bf6 - --hash=sha256:8b38b347dce9336d00dacd4c0e5813bcb43c61e165c653d754e3ee3dc4b2b715 - --hash=sha256:d1dfe295ccbc57cac39e6dee2d7ce831d7a6a504b8a17dc974e631e3521f83a6 + --hash=sha256:39fe8e266875e1105886cab870df4bcbe2d40a84696177bc574c22b96d843bf6 \ + --hash=sha256:8b38b347dce9336d00dacd4c0e5813bcb43c61e165c653d754e3ee3dc4b2b715 \ + --hash=sha256:d1dfe295ccbc57cac39e6dee2d7ce831d7a6a504b8a17dc974e631e3521f83a6 \ --hash=sha256:d3e06fd8c953c4de3d24591cdb9e8e65ca28a876519378d1e4a5d9bca419d7c1 \ No newline at end of file From f54380206cf7926e596d3a66ba12e94aa5d71ccd Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Fri, 30 Aug 2024 12:30:15 -0400 Subject: [PATCH 08/12] Adding pillow dependency --- python/build_requirements.txt | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/build_requirements.txt b/python/build_requirements.txt index 2e5e1ad9..8cd71090 100644 --- a/python/build_requirements.txt +++ b/python/build_requirements.txt @@ -98,4 +98,7 @@ usd-core==24.8 \ --hash=sha256:39fe8e266875e1105886cab870df4bcbe2d40a84696177bc574c22b96d843bf6 \ --hash=sha256:8b38b347dce9336d00dacd4c0e5813bcb43c61e165c653d754e3ee3dc4b2b715 \ --hash=sha256:d1dfe295ccbc57cac39e6dee2d7ce831d7a6a504b8a17dc974e631e3521f83a6 \ + --hash=sha256:d3e06fd8c953c4de3d24591cdb9e8e65ca28a876519378d1e4a5d9bca419d7c1 + +pillow==10.2.0 \ --hash=sha256:d3e06fd8c953c4de3d24591cdb9e8e65ca28a876519378d1e4a5d9bca419d7c1 \ No newline at end of file From 1d2ee40c9bb4957bd43cebdf9387e0b32c7c2d00 Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Fri, 30 Aug 2024 12:32:45 -0400 Subject: [PATCH 09/12] Updating pillow hashes --- python/build_requirements.txt | 4 +++- python/pyproject.toml | 5 ++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/python/build_requirements.txt b/python/build_requirements.txt index 8cd71090..ec2f0693 100644 --- a/python/build_requirements.txt +++ b/python/build_requirements.txt @@ -101,4 +101,6 @@ usd-core==24.8 \ --hash=sha256:d3e06fd8c953c4de3d24591cdb9e8e65ca28a876519378d1e4a5d9bca419d7c1 pillow==10.2.0 \ - --hash=sha256:d3e06fd8c953c4de3d24591cdb9e8e65ca28a876519378d1e4a5d9bca419d7c1 \ No newline at end of file + --hash=sha256:a086c2af425c5f62a65e12fbf385f7c9fcb8f107d0849dba5839461a129cf311 \ + --hash=sha256:35bb52c37f256f662abdfa49d2dfa6ce5d93281d323a9af377a120e89a9eafb5 \ + --hash=sha256:1da3b2703afd040cf65ec97efea81cfba59cdbed9c11d8efc5ab09df9509fc56 \ No newline at end of file diff --git a/python/pyproject.toml b/python/pyproject.toml index 7d0cb8bb..3bcbc52c 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -62,4 +62,7 @@ mujoco = [ ] [project.optional-dependencies] -usd = ["usd-core"] +usd = [ + "usd-core", + "pillow" +] From 5dc9c07fa662a68af4e6b6fb1bcf34425e0f9a3b Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Fri, 30 Aug 2024 12:56:25 -0400 Subject: [PATCH 10/12] Removing scipy dependency --- python/mujoco/usd/exporter.py | 32 ++++++++++++-------------------- 1 file changed, 12 insertions(+), 20 deletions(-) diff --git a/python/mujoco/usd/exporter.py b/python/mujoco/usd/exporter.py index 7879b997..29055a32 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -25,8 +25,6 @@ import mujoco.usd.shapes as shapes_module import numpy as np from PIL import Image as im from PIL import ImageOps -import scipy -import termcolor # TODO: b/288149332 - Remove once USD Python Binding works well with pytype. # pytype: disable=module-attr @@ -165,12 +163,9 @@ class USDExporter: os.makedirs(self.assets_directory) if self.verbose: - print( - termcolor.colored( - "Writing output frames and assets to" - f" {self.output_directory_path}", - "green", - ) + print("Writing output frames and assets to" + f" {self.output_directory_path}", + "green" ) def update_scene( @@ -235,12 +230,9 @@ class USDExporter: data_adr += pixels if self.verbose: - print( - termcolor.colored( - f"Completed writing {self.model.ntex} textures to" - f" {self.assets_directory}", - "green", - ) + print(f"Completed writing {self.model.ntex} textures to" + f" {self.assets_directory}", + "green", ) def _load_geom(self, geom: mujoco.MjvGeom): @@ -449,9 +441,11 @@ class USDExporter: new_camera = camera_module.USDCamera( stage=self.stage, obj_name=obj_name) - r = scipy.spatial.transform.Rotation.from_euler( - "xyz", rotation_xyz, degrees=True) - new_camera.update(cam_pos=np.array(pos), cam_mat=r.as_matrix(), frame=0) + R = np.zeros(9) + quat = np.zeros(4) + mujoco.mju_euler2Quat(quat, rotation_xyz, "xyz") + mujoco.mju_quat2Mat(R, quat) + new_camera.update(cam_pos=np.array(pos), cam_mat=R, frame=0) def save_scene(self, filetype: str = "usd"): """Saves the scene to a USD file.""" @@ -468,9 +462,7 @@ class USDExporter: ) if self.verbose: print( - termcolor.colored( - f"Completed writing frame_{self.frame_count}.{filetype}", "green" - ) + f"Completed writing frame_{self.frame_count}.{filetype}", "green" ) def _get_geom_name(self, geom): From c0dc787cee807d2f16e206de5e92031e57f6c78e Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Fri, 30 Aug 2024 18:23:13 -0400 Subject: [PATCH 11/12] Removing mujoco renderer dependency --- python/mujoco/usd/exporter.py | 64 ++++++++++++++++++++++++++++------- 1 file changed, 51 insertions(+), 13 deletions(-) diff --git a/python/mujoco/usd/exporter.py b/python/mujoco/usd/exporter.py index 29055a32..ea9bbd2d 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -15,9 +15,10 @@ """USD exporter.""" import os -from typing import List, Optional +from typing import List, Optional, Union import mujoco +from mujoco import _structs, _enums, _functions import mujoco.usd.camera as camera_module import mujoco.usd.lights as light_module import mujoco.usd.objects as object_module @@ -113,7 +114,7 @@ class USDExporter: self.usd_cameras = [] # initializing rendering requirements - self.renderer = mujoco.Renderer(model, height, width, max_geom) + self._scene = _structs.MjvScene(model=model, maxgeom=max_geom) self._initialize_usd_stage() self._scene_option = mujoco.MjvOption() # using default scene option @@ -131,7 +132,7 @@ class USDExporter: @property def scene(self): """Returns the scene.""" - return self.renderer.scene + return self._scene def _initialize_usd_stage(self): """Initializes a USD stage to represent the mujoco scene.""" @@ -164,10 +165,48 @@ class USDExporter: if self.verbose: print("Writing output frames and assets to" - f" {self.output_directory_path}", - "green" + f" {self.output_directory_path}" ) + def _update_scene( + self, + data: _structs.MjData, + camera: Union[int, str] = -1, + scene_option: Optional[_structs.MjvOption] = None + ): + camera_id = camera + if isinstance(camera_id, str): + camera_id = _functions.mj_name2id( + self.model, _enums.mjtObj.mjOBJ_CAMERA.value, camera_id + ) + if camera_id == -1: + raise ValueError(f"The camera '{camera}' does not exist.") + if camera_id < -1 or camera_id >= self.model.ncam: + raise ValueError(f"The camera id {camera_id} is out of" + f" range [-1, {self.model.ncam}).") + + # Render camera. + camera = _structs.MjvCamera() + camera.fixedcamid = camera_id + + # Defaults to mjCAMERA_FREE, otherwise mjCAMERA_FIXED refers to a + # camera explicitly defined in the model. + if camera_id == -1: + camera.type = _enums.mjtCamera.mjCAMERA_FREE + _functions.mjv_defaultFreeCamera(self.model, camera) + else: + camera.type = _enums.mjtCamera.mjCAMERA_FIXED + + scene_option = scene_option or self._scene_option + _functions.mjv_updateScene( + self.model, + data, + scene_option, + None, + camera, _enums.mjtCatBit.mjCAT_ALL.value, + self._scene, + ) + def update_scene( self, data: mujoco.MjData, @@ -185,7 +224,7 @@ class USDExporter: scene_option = scene_option or self._scene_option # update the mujoco renderer - self.renderer.update_scene(data, scene_option=scene_option) + self._update_scene(data, scene_option=scene_option) if self.updates == 0: self._initialize_usd_stage() @@ -232,7 +271,6 @@ class USDExporter: if self.verbose: print(f"Completed writing {self.model.ntex} textures to" f" {self.assets_directory}", - "green", ) def _load_geom(self, geom: mujoco.MjvGeom): @@ -374,8 +412,8 @@ class USDExporter: camera = self.usd_cameras[i] camera_name = self.camera_names[i] - self.renderer.update_scene( - data, scene_option=scene_option, camera=camera_name + self._update_scene( + data, camera=camera_name, scene_option=scene_option, ) avg_camera = mujoco.mjv_averageCamera( @@ -441,11 +479,11 @@ class USDExporter: new_camera = camera_module.USDCamera( stage=self.stage, obj_name=obj_name) - R = np.zeros(9) + rotation = np.zeros(9) quat = np.zeros(4) mujoco.mju_euler2Quat(quat, rotation_xyz, "xyz") - mujoco.mju_quat2Mat(R, quat) - new_camera.update(cam_pos=np.array(pos), cam_mat=R, frame=0) + mujoco.mju_quat2Mat(rotation, quat) + new_camera.update(cam_pos=np.array(pos), cam_mat=rotation, frame=0) def save_scene(self, filetype: str = "usd"): """Saves the scene to a USD file.""" @@ -462,7 +500,7 @@ class USDExporter: ) if self.verbose: print( - f"Completed writing frame_{self.frame_count}.{filetype}", "green" + f"Completed writing frame_{self.frame_count}.{filetype}" ) def _get_geom_name(self, geom): From fe1d6dad0232f6d42e76776f941a207bdac3af40 Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Wed, 4 Sep 2024 18:18:37 -0400 Subject: [PATCH 12/12] Using temp directory for usd exporter test --- python/mujoco/usd/exporter_test.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/mujoco/usd/exporter_test.py b/python/mujoco/usd/exporter_test.py index 6a61628f..c6c74aee 100644 --- a/python/mujoco/usd/exporter_test.py +++ b/python/mujoco/usd/exporter_test.py @@ -18,6 +18,7 @@ import os from absl.testing import absltest from etils import epath +import tempfile import mujoco from mujoco.usd import exporter as exporter_module # pylint: disable=g-import-not-at-top @@ -25,7 +26,7 @@ class ExporterTest(absltest.TestCase): def test_usd_export(self): - output_dir_root = os.getenv("TEST_UNDECLARED_OUTPUTS_DIR") + output_dir_root = os.getenv('TEST_UNDECLARED_OUTPUTS_DIR', tempfile.gettempdir()) output_dir_name = "usd_test" xml = """