From f6ee435005df7fead03ecf474061ae123aa404ae Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Sat, 4 May 2024 13:47:02 -0500 Subject: [PATCH] updating geom visibility logic; fixing infinite plane texture bug; adding frame skip arg to demo --- python/mujoco/usd/component.py | 49 +++++++++++++++------- python/mujoco/usd/demo.py | 12 ++++-- python/mujoco/usd/exporter.py | 75 +++++++++++++++++++--------------- 3 files changed, 86 insertions(+), 50 deletions(-) diff --git a/python/mujoco/usd/component.py b/python/mujoco/usd/component.py index e91fc8e2..d7729e52 100644 --- a/python/mujoco/usd/component.py +++ b/python/mujoco/usd/component.py @@ -16,8 +16,10 @@ from typing import List, Optional, Tuple import mujoco -import mujoco.usd.utils as utils_component -import mujoco.usd.shapes as shapes_component +# import mujoco.usd.utils as utils_component +# import mujoco.usd.shapes as shapes_component +import utils as utils_component +import shapes as shapes_component import numpy as np @@ -80,6 +82,8 @@ class USDMesh: # defining ops required by update function self.transform_op = self.usd_xform.AddTransformOp() + self.last_visible_frame = -2 # not an arbitary value, forces difference greater than 1 for visibility on 0th frame + def get_facetexcoord_ranges(self, nmesh, arr): facetexcoords_ranges = [0] running_sum = 0 @@ -224,13 +228,18 @@ class USDMesh: rotation_matrix=mat, translation_vector=pos ).T self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame) - self.update_visibility(visible, frame) + + if visible and frame - self.last_visible_frame > 1: + # non consecutive visible frames + self.update_visibility(False, max(0, self.last_visible_frame)) + self.update_visibility(True, frame) + + if visible: + self.last_visible_frame = frame def update_visibility(self, visible: bool, frame: int): - if visible: - self.usd_prim.GetAttribute("visibility").Set("inherited", frame) - else: - self.usd_prim.GetAttribute("visibility").Set("invisible", frame) + visibility_setting = "inherited" if visible else "invisible" + self.usd_prim.GetAttribute("visibility").Set(visibility_setting, frame) class USDPrimitiveMesh: @@ -290,6 +299,8 @@ class USDPrimitiveMesh: # defining ops required by update function self.transform_op = self.usd_xform.AddTransformOp() + self.last_visible_frame = -2 + def _set_refinement_properties(self): self.usd_prim.GetAttribute("subdivisionScheme").Set("none") @@ -298,14 +309,17 @@ class USDPrimitiveMesh: assert self.prim_mesh x_scale, y_scale = self.geom.texrepeat - mesh_texcoord = np.array(self.prim_mesh.triangle_uvs) mesh_facetexcoord = np.asarray(self.prim_mesh.triangles) x_multiplier, y_multiplier = 1, 1 if self.geom.texuniform: - x_multiplier, y_multiplier = self.geom.size[:2] - + # x_multiplier, y_multiplier = self.geom.size[:2] + if "box" in self.mesh_config: + x_multiplier, y_multiplier = self.mesh_config["box"]["width"], self.mesh_config["box"]["height"] + elif "sphere" in self.mesh_config: + x_multiplier, y_multiplier = self.mesh_config["sphere"]["radius"], self.mesh_config["sphere"]["radius"] + mesh_texcoord[:, 0] *= x_scale * x_multiplier mesh_texcoord[:, 1] *= y_scale * y_multiplier @@ -408,13 +422,18 @@ class USDPrimitiveMesh: rotation_matrix=mat, translation_vector=pos ).T self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame) - self.update_visibility(visible, frame) + + if visible and frame - self.last_visible_frame > 1: + # non consecutive visible frames + self.update_visibility(False, max(0, self.last_visible_frame)) + self.update_visibility(True, frame) + + if visible: + self.last_visible_frame = frame def update_visibility(self, visible: bool, frame: int): - if visible: - self.usd_prim.GetAttribute("visibility").Set("inherited", frame) - else: - self.usd_prim.GetAttribute("visibility").Set("invisible", frame) + visibility_setting = "inherited" if visible else "invisible" + self.usd_prim.GetAttribute("visibility").Set(visibility_setting, frame) class USDSphereLight: diff --git a/python/mujoco/usd/demo.py b/python/mujoco/usd/demo.py index 13b6f21f..645fa2e4 100644 --- a/python/mujoco/usd/demo.py +++ b/python/mujoco/usd/demo.py @@ -3,7 +3,8 @@ from tqdm import tqdm from pathlib import Path import mujoco -from mujoco.usd import exporter +# from mujoco.usd import exporter +import exporter as exporter def generate_usd_trajectory(args): @@ -20,12 +21,12 @@ def generate_usd_trajectory(args): # step through the model for length steps for i in tqdm(range(args.length)): - mujoco.mj_step(m, d) + for i in range(args.steps_per_frame): + mujoco.mj_step(m, d) exp.update_scene(d) exp.save_scene(filetype=args.export_extension) - if __name__ == "__main__": parser = argparse.ArgumentParser() @@ -55,6 +56,11 @@ if __name__ == "__main__": default="usd", help='extension of exported file (can be usd, usda, or usdc)') + parser.add_argument('--steps_per_frame', + type=int, + default=1, + help='number of frames to skip for each rendering step') + args = parser.parse_args() generate_usd_trajectory(args) diff --git a/python/mujoco/usd/exporter.py b/python/mujoco/usd/exporter.py index d4b03fae..745f430a 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -16,8 +16,10 @@ import os import mujoco -import mujoco.usd.shapes as shapes_module -import mujoco.usd.component as component_module +# import mujoco.usd.shapes as shapes_module +# import mujoco.usd.component as component_module +import shapes as shapes_module +import component as component_module import numpy as np import scipy @@ -102,7 +104,8 @@ class USDExporter: self.frame_count = 0 # maintains how many times we have saved the scene self.updates = 0 - self.geom_name2usd = {} + self.geom_names = set() + self.geom_refs = {} # initializing rendering requirements self.renderer = mujoco.Renderer(model, height, width, max_geom) @@ -215,7 +218,7 @@ class USDExporter: ) img_path = os.path.join( relative_path, texture_file_name - ) # relative path, TODO: switch back to this + ) self.texture_files.append(img_path) @@ -232,9 +235,9 @@ class USDExporter: def _load_geom(self, geom: mujoco.MjvGeom): - geom_name = self._get_geom_name(geom.objtype, geom.objid) + geom_name = self._get_geom_name(geom) - assert geom_name not in self.geom_name2usd + assert geom_name not in self.geom_names texture_file = self.texture_files[geom.texid] if geom.texid != -1 else None @@ -263,42 +266,30 @@ class USDExporter: texture_file=texture_file, ) - self.geom_name2usd[geom_name] = usd_geom + self.geom_names.add(geom_name) + self.geom_refs[geom_name] = usd_geom def _update_geoms(self): - geom_names = set(self.geom_name2usd.keys()) - - for i in range(self.scene.ngeom): - geom = self.scene.geoms[i] - geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid) - if not geom_name: - geom_name = "None" - geom_name += f"_{geom.objid}" - # iterate through all geoms in the scene and makes update for i in range(self.scene.ngeom): geom = self.scene.geoms[i] - geom_name = self._get_geom_name(geom.objtype, geom.objid) + geom_name = self._get_geom_name(geom) - if geom_name not in self.geom_name2usd: - self._load_geom(geom) - if self.geom_name2usd[geom_name]: - self.geom_name2usd[geom_name].update_visibility(False, 0) + if "tendon" in geom_name: + # we handle tendon geoms separately to avoid naming conflicts + pass + else: + if geom_name not in self.geom_names: + # load a new object into USD + self._load_geom(geom) - if self.geom_name2usd[geom_name]: - self.geom_name2usd[geom_name].update( + self.geom_refs[geom_name].update( pos=geom.pos, mat=geom.mat, visible=geom.rgba[3] > 0, frame=self.updates, ) - if geom_name in geom_names: - geom_names.remove(geom_name) - - for geom_name in geom_names: - if self.geom_name2usd[geom_name]: - self.geom_name2usd[geom_name].update_visibility(False, self.updates) def _load_lights(self): # initializes an usd light object for every light in the scene @@ -399,15 +390,35 @@ class USDExporter: def save_scene(self, filetype: str = "usd"): assert filetype in ["usd", "usda", "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) + self.stage.Export( f"{self.output_directory_root}/{self.output_directory_name}/frames/frame_{self.frame_count}_.{filetype}" ) if self.verbose: print(termcolor.colored(f"Completed writing frame_{self.frame_count}.{filetype}", "green")) - def _get_geom_name(self, objtype, objid): - geom_name = mujoco.mj_id2name(self.model, objtype, objid) + def _get_geom_name(self, geom): + # adding id as part of name for USD file + geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid) if not geom_name: geom_name = "None" - geom_name += f"_{objid}" + geom_name += f"_id{geom.objid}" + + # adding additional naming information to differentiate between geoms and tendons + if geom.objtype == mujoco.mjtObj.mjOBJ_GEOM: + geom_name += "_geom" + elif geom.objtype == mujoco.mjtObj.mjOBJ_TENDON: + geom_name += "_tendon" + return geom_name + + # for debugging purposes, prints all geoms in scene including those part of tendons + def _print_scene_geom_info(self): + for i in range(self.scene.ngeom): + geom = self.scene.geoms[i] + geom_name = self._get_geom_name(geom) + print(i, geom_name) \ No newline at end of file