diff --git a/.gitignore b/.gitignore index c55f0c71..44654326 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,7 @@ *.egg-info/ build/ build_cmake/ +python/dist/ # Exclude macOS folder attributes .DS_Store @@ -28,4 +29,4 @@ Info.framework.plist # Python byte-compiled / optimized / DLL files __pycache__/ *.py[cod] -*$py.class +*$py.class \ No newline at end of file diff --git a/python/mujoco/usd/__init__.py b/python/mujoco/usd/__init__.py new file mode 100644 index 00000000..1bb0c952 --- /dev/null +++ b/python/mujoco/usd/__init__.py @@ -0,0 +1,3 @@ +from .exporter import * +from .component import * +from .utils import * \ No newline at end of file diff --git a/python/mujoco/usd/component.py b/python/mujoco/usd/component.py index 044689ab..fbd99a2f 100644 --- a/python/mujoco/usd/component.py +++ b/python/mujoco/usd/component.py @@ -378,123 +378,6 @@ class USDPrimitiveMesh: self.usd_prim.GetAttribute("visibility").Set("invisible", frame) -class USDPrimitive: - - def __init__( - self, - stage: Usd.Stage, - geom: mujoco.MjvGeom, - obj_name: str, - rgba: np.ndarray = np.array([1, 1, 1, 1]), - texture_file: Optional[str] = None, - ): - self.stage = stage - self.geom = geom - self.obj_name = obj_name - self.rgba = rgba - self.texture_file = texture_file - - self.usd_prim = Usd.Prim() - self.usd_primitive_shape = Usd.PrimitiveShape() - self.transform_op = Usd.TransformOp() - - def _set_refinement_properties(self): - self.usd_prim.GetAttribute("subdivisionScheme").Set("none") - - def _attach_material(self): - mtl_path = Sdf.Path(f"/World/_materials/Material_{self.obj_name}") - mtl = UsdShade.Material.Define(self.stage, mtl_path) - if self.texture_file: - bsdf_shader = UsdShade.Shader.Define( - self.stage, mtl_path.AppendPath("Principled_BSDF") - ) - image_shader = UsdShade.Shader.Define( - self.stage, mtl_path.AppendPath("Image_Texture") - ) - uvmap_shader = UsdShade.Shader.Define( - self.stage, mtl_path.AppendPath("uvmap") - ) - - # settings the bsdf shader attributes - bsdf_shader.CreateIdAttr("UsdPreviewSurface") - bsdf_shader.CreateInput( - "diffuseColor", Sdf.ValueTypeNames.Color3f - ).ConnectToSource(image_shader.ConnectableAPI(), "rgb") - bsdf_shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set( - float(self.rgba[-1]) - ) - bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set( - self.geom.shininess - ) - bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set( - 1.0 - self.geom.shininess - ) - - mtl.CreateSurfaceOutput().ConnectToSource( - bsdf_shader.ConnectableAPI(), "surface" - ) - - self.usd_primitive_shape.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI) - UsdShade.MaterialBindingAPI(self.usd_primitive_shape).Bind(mtl) - - # setting the image texture attributes - image_shader.CreateIdAttr("UsdUVTexture") - image_shader.CreateInput("file", Sdf.ValueTypeNames.Asset).Set( - self.texture_file - ) - image_shader.CreateInput( - "sourceColorSpace", Sdf.ValueTypeNames.Token - ).Set("sRGB") - image_shader.CreateInput("st", Sdf.ValueTypeNames.Float2).ConnectToSource( - uvmap_shader.ConnectableAPI(), "result" - ) - image_shader.CreateOutput("rgb", Sdf.ValueTypeNames.Float3) - - # setting uvmap shader attributes - uvmap_shader.CreateIdAttr("UsdPrimvarReader_float2") - uvmap_shader.CreateInput("varname", Sdf.ValueTypeNames.Token).Set("UVMap") - uvmap_shader.CreateOutput("results", Sdf.ValueTypeNames.Float2) - else: - bsdf_shader = UsdShade.Shader.Define( - self.stage, mtl_path.AppendPath("Principled_BSDF") - ) - - # settings the bsdf shader attributes - bsdf_shader.CreateIdAttr("UsdPreviewSurface") - bsdf_shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set( - tuple(self.rgba[:3]) - ) - bsdf_shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set( - float(self.rgba[-1]) - ) - bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set( - self.geom.shininess - ) - bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set( - 1.0 - self.geom.shininess - ) - - mtl.CreateSurfaceOutput().ConnectToSource( - bsdf_shader.ConnectableAPI(), "surface" - ) - - self.usd_primitive_shape.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI) - UsdShade.MaterialBindingAPI(self.usd_primitive_shape).Bind(mtl) - - def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int): - transformation_mat = mujoco.usd_util.create_transform_matrix( - rotation_matrix=mat, translation_vector=pos - ).T - self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame) - self.update_visibility(visible, 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) - - class USDCapsule(USDPrimitive): def __init__( @@ -522,7 +405,7 @@ class USDCapsule(USDPrimitive): self._set_size_attributes() self._attach_material() - self._set_refinement_properties() + # self._set_refinement_properties() def _set_size_attributes(self): self.usd_primitive_shape.GetRadiusAttr().Set(float(self.geom.size[0])) @@ -558,7 +441,7 @@ class USDEllipsoid(USDPrimitive): self._set_size_attributes() self._attach_material() - self._set_refinement_properties() + # self._set_refinement_properties() def _set_size_attributes(self): self.scale_op.Set(Gf.Vec3d(self.geom.size.tolist())) @@ -849,7 +732,7 @@ class USDCamera: def update(self, cam_pos: np.ndarray, cam_mat: np.ndarray, frame: int): - transformation_mat = mujoco.usd_util.create_transform_matrix( + transformation_mat = mujoco.usd.utils.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/demo.py b/python/mujoco/usd/demo.py new file mode 100644 index 00000000..0279a690 --- /dev/null +++ b/python/mujoco/usd/demo.py @@ -0,0 +1,19 @@ +import mujoco +from mujoco.usd import exporter + +if __name__ == "__main__": + + # load a model to mujoco + m = mujoco.MjModel.from_xml_path("/Users/abhishek/Documents/research/mujoco/model/humanoid/humanoid.xml") + d = mujoco.MjData(m) + + # create an instance of the USDExporter + exp = exporter.USDExporter(model=m) + + mujoco.mj_step(m, d) + + exp.update_scene(d) + + exp.save_scene(filetype="usda") + + diff --git a/python/mujoco/usd/exporter.py b/python/mujoco/usd/exporter.py index 33513ec2..cb069a05 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -395,7 +395,7 @@ class USDExporter: intensity: int, radius: Optional[float] = 1.0, color: Optional[np.ndarray] = np.array([0.3, 0.3, 0.3]), - objid: Optional[int] = 1, + obj_name: Optional[str] = "light_1", light_type: Optional[str] = "sphere", ): @@ -413,7 +413,7 @@ class USDExporter: self, pos: List[float], rotation_xyz: List[float], - objid: Optional[int] = 1, + obj_name: Optional[str] = "camera_1", ): new_camera = component_module.USDCamera( stage=self.stage, obj_name=str(objid)) diff --git a/python/mujoco/usd/shapes.py b/python/mujoco/usd/shapes.py new file mode 100644 index 00000000..e69de29b