From 508d13cab82058a9420d10223b484218ffa97086 Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Tue, 11 Feb 2025 17:38:16 -0500 Subject: [PATCH] formatting --- python/mujoco/usd/camera.py | 6 +- python/mujoco/usd/demo.py | 11 +-- python/mujoco/usd/exporter.py | 133 ++++++++++++----------------- python/mujoco/usd/exporter_test.py | 1 + python/mujoco/usd/lights.py | 1 - python/mujoco/usd/objects.py | 17 ++-- python/mujoco/usd/shapes.py | 85 +++++++++--------- 7 files changed, 113 insertions(+), 141 deletions(-) diff --git a/python/mujoco/usd/camera.py b/python/mujoco/usd/camera.py index 01d2f521d0..e85361358d 100644 --- a/python/mujoco/usd/camera.py +++ b/python/mujoco/usd/camera.py @@ -16,10 +16,7 @@ from typing import List, Optional, Tuple -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 @@ -28,9 +25,10 @@ import mujoco.usd.utils as utils_module + class USDCamera: """Class that handles the cameras in the USD scene""" - + def __init__(self, stage: Usd.Stage, camera_name: str): self.stage = stage diff --git a/python/mujoco/usd/demo.py b/python/mujoco/usd/demo.py index 5bc4ec230c..5b9cf15232 100644 --- a/python/mujoco/usd/demo.py +++ b/python/mujoco/usd/demo.py @@ -37,16 +37,14 @@ def generate_usd_trajectory(local_args): ) cam = mujoco.MjvCamera() - + # step through the simulation for the given duration of time while d.time < local_args.duration: mujoco.mj_step(m, d) if exp.frame_count < d.time * local_args.framerate: exp.update_scene(data=d, camera=cam) - exp.add_light(pos=(0, 0, 0), - intensity=2000, - light_type='dome') + exp.add_light(pos=(0, 0, 0), intensity=2000, light_type='dome') exp.save_scene(filetype=local_args.export_extension) @@ -81,10 +79,7 @@ def generate_usd_trajectory(local_args): ) 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 405f4203e5..39149f3498 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -17,26 +17,27 @@ import os from typing import List, Optional, Tuple, 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 -import mujoco.usd.shapes as shapes_module -import mujoco.usd.utils as utils_module import numpy as np from PIL import Image as im from PIL import ImageOps - # TODO: b/288149332 - Remove once USD Python Binding works well with pytype. # pytype: disable=module-attr from pxr import Sdf from pxr import Usd from pxr import UsdGeom +import mujoco +from mujoco import _enums +from mujoco import _functions +from mujoco import _structs +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 + PRIMARY_CAMERA_NAME = "primary_camera" + class USDExporter: """MuJoCo to USD exporter for porting scenes to external renderers.""" @@ -139,48 +140,6 @@ def _initialize_output_directories(self): if self.verbose: print(f"Writing output frames and assets to {self.output_directory_path}") - def _update_scene( - self, - data: mujoco.MjData, - camera: Union[int, str] = -1, - scene_option: Optional[mujoco.MjvOption] = None, - ): - """Updates the scene.""" - camera_id = camera - if isinstance(camera_id, str): - camera_id = mujoco.mj_name2id( - self.model, mujoco.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 range [-1, {self.model.ncam})." - ) - - # Render camera. - camera = mujoco.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 = mujoco.mjtCamera.mjCAMERA_FREE - mujoco.mjv_defaultFreeCamera(self.model, camera) - else: - camera.type = mujoco.mjtCamera.mjCAMERA_FIXED - - scene_option = scene_option or self._scene_option - mujoco.mjv_updateScene( - self.model, - data, - scene_option, - None, - camera, - mujoco.mjtCatBit.mjCAT_ALL.value, - self._scene, - ) - def _update_scene( self, data: mujoco.MjData, @@ -189,7 +148,7 @@ def _update_scene( ) -> None: if not isinstance(camera, _structs.MjvCamera): if self.model.ncam == 0: - raise ValueError(f'No fixed cameras defined in mujoco model.') + raise ValueError(f"No fixed cameras defined in mujoco model.") camera_id = camera if isinstance(camera_id, str): camera_id = _functions.mj_name2id( @@ -198,10 +157,12 @@ def _update_scene( if camera_id == -1: raise ValueError(f'The camera "{camera}" does not exist.') if camera_id == -1: - raise ValueError('Free cameras are not supported during USD export.') + raise ValueError("Free cameras are not supported during USD export.") if camera_id < 0 or camera_id >= self.model.ncam: - raise ValueError(f'The camera id {camera_id} is out of' - f' range [-1, {self.model.ncam}).') + raise ValueError( + f"The camera id {camera_id} is out of" + f" range [-1, {self.model.ncam})." + ) assert camera_id != -1 @@ -216,7 +177,8 @@ def _update_scene( data, scene_option, None, - camera, _enums.mjtCatBit.mjCAT_ALL.value, + camera, + _enums.mjtCatBit.mjCAT_ALL.value, self._scene, ) @@ -274,9 +236,7 @@ def _load_textures(self): relative_path = os.path.relpath( self.assets_directory, self.frames_directory ) - img_path = os.path.join( - relative_path, texture_file_name - ) + img_path = os.path.join(relative_path, texture_file_name) self.texture_files.append(img_path) @@ -297,7 +257,9 @@ def _load_geom(self, geom: mujoco.MjvGeom): geom_textures = [] else: geom_textures = [ - (self.texture_files[i.item()], self.model.tex_type[i.item()]) if i.item() != -1 else None + (self.texture_files[i.item()], self.model.tex_type[i.item()]) + if i.item() != -1 + else None for i in self.model.mat_texid[geom.matid] ] @@ -319,7 +281,7 @@ def _load_geom(self, geom: mujoco.MjvGeom): name=geom_name, geom_type=geom.type, size=np.array([1.0, 1.0, 1.0]), - decouple=True + decouple=True, ) usd_geom = object_module.USDTendon( mesh_config=mesh_config, @@ -333,9 +295,7 @@ def _load_geom(self, geom: mujoco.MjvGeom): # handling primitives in our scene else: mesh_config = shapes_module.mesh_config_generator( - name=geom_name, - geom_type=geom.type, - size=geom.size + name=geom_name, geom_type=geom.type, size=geom.size ) usd_geom = object_module.USDPrimitiveMesh( mesh_config=mesh_config, @@ -408,15 +368,21 @@ def _update_lights(self) -> None: def _load_cameras(self) -> None: self.usd_cameras = {} - + # add a primary camera for which the scene will be rendered - self.usd_cameras[PRIMARY_CAMERA_NAME] = camera_module.USDCamera(stage=self.stage, camera_name=PRIMARY_CAMERA_NAME) + self.usd_cameras[PRIMARY_CAMERA_NAME] = camera_module.USDCamera( + stage=self.stage, camera_name=PRIMARY_CAMERA_NAME + ) if self.camera_names is not None: for camera_name in self.camera_names: - self.usd_cameras[camera_name] = camera_module.USDCamera(stage=self.stage, camera_name=camera_name) + self.usd_cameras[camera_name] = camera_module.USDCamera( + stage=self.stage, camera_name=camera_name + ) def _get_camera_orientation(self) -> Tuple[_structs.MjvGLCamera, np.ndarray]: - avg_camera = mujoco.mjv_averageCamera(self._scene.camera[0], self._scene.camera[1]) + avg_camera = mujoco.mjv_averageCamera( + self._scene.camera[0], self._scene.camera[1] + ) forward = avg_camera.forward up = avg_camera.up @@ -437,19 +403,23 @@ def _update_cameras( ) -> None: # first, update the primary camera given the new scene self._update_scene( - data, camera=camera, scene_option=scene_option, + data, + camera=camera, + scene_option=scene_option, ) avg_camera, R = self._get_camera_orientation() - self.usd_cameras[PRIMARY_CAMERA_NAME].update(cam_pos=avg_camera.pos, cam_mat=R, frame=self.updates) + self.usd_cameras[PRIMARY_CAMERA_NAME].update( + cam_pos=avg_camera.pos, cam_mat=R, frame=self.updates + ) # update the names of the fixed cameras in the scene if self.camera_names is not None: for camera_name in self.camera_names: - self._update_scene( - data, camera=camera_name, scene_option=scene_option - ) + self._update_scene(data, camera=camera_name, scene_option=scene_option) avg_camera, R = self._get_camera_orientation() - self.usd_cameras[camera_name].update(cam_pos=avg_camera.pos, cam_mat=R, frame=self.updates) + self.usd_cameras[camera_name].update( + cam_pos=avg_camera.pos, cam_mat=R, frame=self.updates + ) def add_light( self, @@ -479,7 +449,9 @@ def add_light( pos=np.array(pos), intensity=intensity, color=color, frame=0 ) elif light_type == "dome": - new_light = light_module.USDDomeLight(stage=self.stage, light_name=light_name) + new_light = light_module.USDDomeLight( + stage=self.stage, light_name=light_name + ) new_light.update(intensity=intensity, color=color) def add_camera( @@ -496,7 +468,8 @@ def add_camera( camera_name: name of the camera to be stored in USD """ new_camera = camera_module.USDCamera( - stage=self.stage, camera_name=camera_name) + stage=self.stage, camera_name=camera_name + ) rotation = np.zeros(9) quat = np.zeros(4) @@ -507,18 +480,18 @@ def add_camera( def save_scene(self, filetype: str = "usd") -> None: """Saves the current scene as a USD trajectory Args: - filetype: type of USD file to save the scene as (options include + filetype: type of USD file to save the scene as (options include "usd", "usda", and "usdc") """ self.stage.SetEndTimeCode(self.frame_count) # post-processing for visibility of geoms in scene for _, geom_ref in self.geom_refs.items(): - geom_ref.update_visibility(False, geom_ref.last_visible_frame+1) + geom_ref.update_visibility(False, geom_ref.last_visible_frame + 1) self.stage.Export( - f"{self.output_directory_root}/{self.output_directory}/" + - f"frames/frame_{self.frame_count}.{filetype}" + f"{self.output_directory_root}/{self.output_directory}/" + + f"frames/frame_{self.frame_count}.{filetype}" ) if self.verbose: print(f"Completed writing frame_{self.frame_count}.{filetype}") diff --git a/python/mujoco/usd/exporter_test.py b/python/mujoco/usd/exporter_test.py index 8e4c15ea7c..2bf815082e 100644 --- a/python/mujoco/usd/exporter_test.py +++ b/python/mujoco/usd/exporter_test.py @@ -19,6 +19,7 @@ from absl.testing import absltest from etils import epath + import mujoco try: diff --git a/python/mujoco/usd/lights.py b/python/mujoco/usd/lights.py index 4048010238..a1fa1515f2 100644 --- a/python/mujoco/usd/lights.py +++ b/python/mujoco/usd/lights.py @@ -17,7 +17,6 @@ from typing import Optional import numpy as np - # TODO: b/288149332 - Remove once USD Python Binding works well with pytype. # pytype: disable=module-attr from pxr import Gf diff --git a/python/mujoco/usd/objects.py b/python/mujoco/usd/objects.py index 208ec206d5..2f6fb276ca 100644 --- a/python/mujoco/usd/objects.py +++ b/python/mujoco/usd/objects.py @@ -18,12 +18,7 @@ import collections from typing import Any, Dict, Optional, Sequence, Tuple -import mujoco -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 @@ -33,6 +28,10 @@ from pxr import UsdShade from pxr import Vt +import mujoco +import mujoco.usd.shapes as shapes_module +import mujoco.usd.utils as utils_module + class USDObject(abc.ABC): """Abstract interface for all USD objects including meshes and primitives. @@ -55,7 +54,7 @@ def __init__( geom: mujoco.MjvGeom, obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), - geom_textures: Sequence[Optional[Tuple[str, mujoco.mjtTexture]]] = () + geom_textures: Sequence[Optional[Tuple[str, mujoco.mjtTexture]]] = (), ): self.stage = stage self.model = model @@ -225,7 +224,7 @@ def __init__( obj_name: str, dataid: int, rgba: np.ndarray = np.array([1, 1, 1, 1]), - geom_textures: Sequence[Optional[Tuple[str, mujoco.mjtTexture]]] = () + geom_textures: Sequence[Optional[Tuple[str, mujoco.mjtTexture]]] = (), ): super().__init__(stage, model, geom, obj_name, rgba, geom_textures) @@ -324,7 +323,7 @@ def __init__( geom: mujoco.MjvGeom, obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), - geom_textures: Sequence[Optional[Tuple[str, mujoco.mjtTexture]]] = () + geom_textures: Sequence[Optional[Tuple[str, mujoco.mjtTexture]]] = (), ): super().__init__(stage, model, geom, obj_name, rgba, geom_textures) @@ -424,7 +423,7 @@ def __init__( geom: mujoco.MjvGeom, obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), - geom_textures: Sequence[Optional[Tuple[str, mujoco.mjtTexture]]] = () + geom_textures: Sequence[Optional[Tuple[str, mujoco.mjtTexture]]] = (), ): super().__init__(stage, model, geom, obj_name, rgba, geom_textures) diff --git a/python/mujoco/usd/shapes.py b/python/mujoco/usd/shapes.py index 21f5f7571a..1d4e625b7e 100644 --- a/python/mujoco/usd/shapes.py +++ b/python/mujoco/usd/shapes.py @@ -16,14 +16,15 @@ from typing import Any, Dict, Optional, Tuple, Union -import mujoco import numpy as np +import mujoco + def get_triangle_uvs( vertices: np.ndarray, triangles: np.ndarray, - texture_type: Optional[mujoco.mjtTexture] + texture_type: Optional[mujoco.mjtTexture], ): """Returns UV coordinates for a given mesh.""" if texture_type is None: @@ -88,10 +89,12 @@ class TriangleMesh: triangle_uvs: A numpy array of UV coordinates. """ - def __init__(self, - vertices: np.ndarray, - triangles: np.ndarray, - triangle_uvs: np.ndarray): + def __init__( + self, + vertices: np.ndarray, + triangles: np.ndarray, + triangle_uvs: np.ndarray, + ): """Creates a TriangleMesh object. Args: @@ -109,30 +112,34 @@ def create_box( width: float, height: float, depth: float, - texture_type: Optional[mujoco.mjtTexture] + 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]]) + 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) @@ -143,13 +150,13 @@ def create_sphere( cls, radius: float, texture_type: Optional[mujoco.mjtTexture], - resolution: int + resolution: int, ) -> "TriangleMesh": """Creates a sphere.""" vertices = [] triangles = [] - for i in range(2*resolution + 1): - phi = np.pi * i / (2*resolution) + 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) @@ -157,7 +164,7 @@ def create_sphere( z = radius * np.cos(phi) vertices.append([x, y, z]) - for i in range(2*resolution): + for i in range(2 * resolution): for j in range(resolution): first = i * (resolution + 1) + j second = first + resolution + 1 @@ -217,7 +224,7 @@ def create_cylinder( radius: float, height: float, texture_type: Optional[mujoco.mjtTexture], - resolution: int + resolution: int, ) -> "TriangleMesh": """Creates a cylinder.""" vertices = [] @@ -330,8 +337,8 @@ def mesh_config_generator( }, "right_hemisphere": { "radius": size[0], - "transform": {"translate": (0, 0, 2*size[2])}, - } + "transform": {"translate": (0, 0, 2 * size[2])}, + }, } elif geom_type == mujoco.mjtGeom.mjGEOM_ELLIPSOID: sphere = mesh_config_generator( @@ -390,26 +397,26 @@ def mesh_factory( width=mesh_config[shape]["width"], height=mesh_config[shape]["height"], depth=mesh_config[shape]["depth"], - texture_type=texture_type + texture_type=texture_type, ) elif "hemisphere" in shape: prim_mesh = TriangleMesh.create_hemisphere( radius=mesh_config[shape]["radius"], texture_type=texture_type, - resolution=resolution + resolution=resolution, ) elif "sphere" in shape: prim_mesh = TriangleMesh.create_sphere( radius=mesh_config[shape]["radius"], texture_type=texture_type, - resolution=resolution + 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 + resolution=resolution, ) else: raise ValueError("Shape not supported")