diff --git a/python/mujoco/usd/__init__.py b/python/mujoco/usd/__init__.py index 1bb0c952..b3c238f5 100644 --- a/python/mujoco/usd/__init__.py +++ b/python/mujoco/usd/__init__.py @@ -1,3 +1,4 @@ from .exporter import * from .component import * -from .utils import * \ No newline at end of file +from .utils import * +from .shapes import * \ No newline at end of file diff --git a/python/mujoco/usd/component.py b/python/mujoco/usd/component.py index fbd99a2f..b883e3b1 100644 --- a/python/mujoco/usd/component.py +++ b/python/mujoco/usd/component.py @@ -15,7 +15,12 @@ from typing import List, Optional, Tuple import mujoco -import mujoco.usd.utils + +# import mujoco.usd.utils +# import mujoco.usd.shapes as shapes_component +import utils as utils_component +import shapes as shapes_component + import numpy as np # TODO: b/288149332 - Remove once USD Python Binding works well with pytype. @@ -217,7 +222,7 @@ class USDMesh: UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl) def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int): - transformation_mat = mujoco.usd.utils.create_transform_matrix( + transformation_mat = utils_component.create_transform_matrix( rotation_matrix=mat, translation_vector=pos ).T self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame) @@ -234,12 +239,14 @@ class USDPrimitiveMesh: def __init__( self, + mesh_config: List[dict], stage: Usd.Stage, geom: mujoco.MjvGeom, obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): + self.mesh_config = mesh_config self.stage = stage self.geom = geom self.obj_name = obj_name @@ -251,6 +258,40 @@ class USDPrimitiveMesh: self.prim_mesh = None self.transform_op = Gf.Matrix4d(1.) + _, self.prim_mesh = shapes_component.mesh_generator(mesh_config) + + xform_path = f"/World/{self.obj_name}_Xform" + mesh_path = f"{xform_path}/{obj_name}" + self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) + self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path) + self.usd_prim = stage.GetPrimAtPath(mesh_path) + + self.prim_mesh.translate(-self.prim_mesh.get_center()) + + mesh_vert, mesh_face, mesh_facenum = self._get_mesh_geometry() + self.usd_mesh.GetPointsAttr().Set(mesh_vert) + self.usd_mesh.GetFaceVertexCountsAttr().Set( + [3 for _ in range(mesh_facenum)] + ) + self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face) + + # setting mesh uv properties + mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry() + self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar( + "UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying + ) + + self.texcoords.Set(mesh_texcoord) + self.texcoords.SetIndices(Vt.IntArray([i for i in range(mesh_facenum * 3)])) + + self._set_refinement_properties() + + # setting attributes for the shape + self._attach_material() + + # defining ops required by update function + self.transform_op = self.usd_xform.AddTransformOp() + def _set_refinement_properties(self): self.usd_prim.GetAttribute("subdivisionScheme").Set("none") @@ -365,7 +406,7 @@ class USDPrimitiveMesh: UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl) def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int): - transformation_mat = mujoco.usd.utils.create_transform_matrix( + transformation_mat = utils_component.create_transform_matrix( rotation_matrix=mat, translation_vector=pos ).T self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame) @@ -377,284 +418,6 @@ class USDPrimitiveMesh: else: self.usd_prim.GetAttribute("visibility").Set("invisible", frame) - -class USDCapsule(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, - ): - - super().__init__(stage, geom, obj_name, rgba, texture_file) - - xform_path = f"/World/Capsule_Xform_{obj_name}" - capsule_path = f"{xform_path}/Capsule_{obj_name}" - self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) - self.usd_primitive_shape = UsdGeom.Capsule.Define(stage, capsule_path) - self.usd_prim = stage.GetPrimAtPath(capsule_path) - - # defining ops required by update function - self.transform_op = self.usd_xform.AddTransformOp() - self.scale_op = self.usd_xform.AddScaleOp() - - # setting attributes for the shape - self._set_size_attributes() - self._attach_material() - - # self._set_refinement_properties() - - def _set_size_attributes(self): - self.usd_primitive_shape.GetRadiusAttr().Set(float(self.geom.size[0])) - self.usd_primitive_shape.GetHeightAttr().Set( - float(self.geom.size[2] * 2) - ) # mujoco gives the half length - - -class USDEllipsoid(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, - ): - - super().__init__(stage, geom, obj_name, rgba, texture_file) - - xform_path = f"/World/Ellipsoid_Xform_{obj_name}" - ellipsoid_path = f"{xform_path}/Ellipsoid_{obj_name}" - self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) - self.usd_primitive_shape = UsdGeom.Sphere.Define(stage, ellipsoid_path) - self.usd_prim = stage.GetPrimAtPath(ellipsoid_path) - - # defining ops required by update function - self.transform_op = self.usd_xform.AddTransformOp() - self.scale_op = self.usd_xform.AddScaleOp() - - # setting attributes for the shape - self._set_size_attributes() - self._attach_material() - - # self._set_refinement_properties() - - def _set_size_attributes(self): - self.scale_op.Set(Gf.Vec3d(self.geom.size.tolist())) - - -class USDCubeMesh(USDPrimitiveMesh): - - 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, - ): - - super().__init__(stage, geom, obj_name, rgba, texture_file) - - xform_path = f"/World/CubeMesh_Xform_{obj_name}" - mesh_path = f"{xform_path}/CubeMesh_{obj_name}" - self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) - self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path) - self.usd_prim = stage.GetPrimAtPath(mesh_path) - - self.prim_mesh = o3d.geometry.TriangleMesh.create_box( - width=self.geom.size[0] * 2, - height=self.geom.size[1] * 2, - depth=self.geom.size[2] * 2, - ) - - self.prim_mesh.translate(-self.prim_mesh.get_center()) - - mesh_vert, mesh_face, mesh_facenum = self._get_mesh_geometry() - self.usd_mesh.GetPointsAttr().Set(mesh_vert) - self.usd_mesh.GetFaceVertexCountsAttr().Set( - [3 for _ in range(mesh_facenum)] - ) - self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face) - - # setting mesh uv properties - mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry() - self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar( - "UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying - ) - - self.texcoords.Set(mesh_texcoord) - self.texcoords.SetIndices(Vt.IntArray([i for i in range(mesh_facenum * 3)])) - - self._set_refinement_properties() - - # setting attributes for the shape - self._attach_material() - - # defining ops required by update function - self.transform_op = self.usd_xform.AddTransformOp() - - -class USDSphereMesh(USDPrimitiveMesh): - - 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, - ): - - super().__init__(stage, geom, obj_name, rgba, texture_file) - - xform_path = f"/World/SphereMesh_Xform_{obj_name}" - mesh_path = f"{xform_path}/SphereMesh_{obj_name}" - self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) - self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path) - self.usd_prim = stage.GetPrimAtPath(mesh_path) - - self.prim_mesh = o3d.geometry.TriangleMesh.create_sphere( - radius=float(self.geom.size[0]), create_uv_map=True - ) - - self.prim_mesh.translate(-self.prim_mesh.get_center()) - - mesh_vert, mesh_face, mesh_facenum = self._get_mesh_geometry() - self.usd_mesh.GetPointsAttr().Set(mesh_vert) - self.usd_mesh.GetFaceVertexCountsAttr().Set( - [3 for _ in range(mesh_facenum)] - ) - self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face) - - # setting mesh uv properties - mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry() - self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar( - "UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying - ) - - self.texcoords.Set(mesh_texcoord) - self.texcoords.SetIndices(Vt.IntArray([i for i in range(mesh_facenum * 3)])) - - self._set_refinement_properties() - - # setting attributes for the shape - self._attach_material() - - # defining ops required by update function - self.transform_op = self.usd_xform.AddTransformOp() - - -class USDCylinderMesh(USDPrimitiveMesh): - - 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, - ): - - super().__init__(stage, geom, obj_name, rgba, texture_file) - - xform_path = f"/World/CylinderMesh_Xform_{obj_name}" - mesh_path = f"{xform_path}/CylinderMesh_{obj_name}" - self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) - self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path) - self.usd_prim = stage.GetPrimAtPath(mesh_path) - - self.prim_mesh = o3d.geometry.TriangleMesh.create_cylinder( - radius=self.geom.size[0], - height=self.geom.size[2] * 2, - create_uv_map=True, - ) - - self.prim_mesh.translate(-self.prim_mesh.get_center()) - - mesh_vert, mesh_face, mesh_facenum = self._get_mesh_geometry() - self.usd_mesh.GetPointsAttr().Set(mesh_vert) - self.usd_mesh.GetFaceVertexCountsAttr().Set( - [3 for _ in range(mesh_facenum)] - ) - self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face) - - # setting mesh uv properties - mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry() - self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar( - "UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying - ) - - self.texcoords.Set(mesh_texcoord) - self.texcoords.SetIndices(Vt.IntArray([i for i in range(mesh_facenum * 3)])) - - self._set_refinement_properties() - - # setting attributes for the shape - self._attach_material() - - # defining ops required by update function - self.transform_op = self.usd_xform.AddTransformOp() - - -class USDPlaneMesh(USDPrimitiveMesh): - - 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, - ): - - super().__init__(stage, geom, obj_name, rgba, texture_file) - - xform_path = f"/World/Plane_Xform_{obj_name}" - plane_path = f"{xform_path}/PlaneMesh_{obj_name}" - self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) - self.usd_mesh = UsdGeom.Mesh.Define(stage, plane_path) - self.usd_prim = stage.GetPrimAtPath(plane_path) - - self.prim_mesh = o3d.geometry.TriangleMesh.create_box( - width=self.geom.size[0] * 2 if self.geom.size[0] > 0 else 100, - height=self.geom.size[1] * 2 if self.geom.size[1] > 0 else 100, - depth=0.001, - create_uv_map=True, - map_texture_to_each_face=True, - ) - - self.prim_mesh.translate(-self.prim_mesh.get_center()) - - mesh_vert, mesh_face, mesh_facenum = self._get_mesh_geometry() - self.usd_mesh.GetPointsAttr().Set(mesh_vert) - self.usd_mesh.GetFaceVertexCountsAttr().Set( - [3 for _ in range(mesh_facenum)] - ) - self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face) - - # setting mesh uv properties - mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry() - self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar( - "UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying - ) - - self.texcoords.Set(mesh_texcoord) - self.texcoords.SetIndices(Vt.IntArray([i for i in range(mesh_facenum * 3)])) - - self._set_refinement_properties() - - # setting attributes for the shape - self._attach_material() - - # defining ops required by update function - self.transform_op = self.usd_xform.AddTransformOp() - - class USDSphereLight: def __init__( @@ -732,7 +495,7 @@ class USDCamera: def update(self, cam_pos: np.ndarray, cam_mat: np.ndarray, frame: int): - transformation_mat = mujoco.usd.utils.create_transform_matrix( + transformation_mat = utils_component.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 index 0279a690..19b4c358 100644 --- a/python/mujoco/usd/demo.py +++ b/python/mujoco/usd/demo.py @@ -1,19 +1,21 @@ + import mujoco -from mujoco.usd import exporter +# from mujoco.usd import exporter +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") + m = mujoco.MjModel.from_xml_path("/Users/abhishek/Documents/research/mujoco/model/car/car.xml") d = mujoco.MjData(m) # create an instance of the USDExporter exp = exporter.USDExporter(model=m) - mujoco.mj_step(m, d) + for i in range(100): + mujoco.mj_step(m, d) + exp.update_scene(d) - exp.update_scene(d) - - exp.save_scene(filetype="usda") + exp.save_scene(filetype="usd") diff --git a/python/mujoco/usd/exporter.py b/python/mujoco/usd/exporter.py index cb069a05..92ae7e02 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -15,7 +15,12 @@ import os import mujoco -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 import termcolor @@ -234,7 +239,6 @@ class USDExporter: texture_file = self.texture_files[geom.texid] if geom.texid != -1 else None - # handles meshes in scene if geom.type == mujoco.mjtGeom.mjGEOM_MESH: usd_geom = component_module.USDMesh( stage=self.stage, @@ -245,56 +249,22 @@ class USDExporter: rgba=geom.rgba, texture_file=texture_file, ) - elif geom.type == mujoco.mjtGeom.mjGEOM_PLANE: - usd_geom = component_module.USDPlaneMesh( - stage=self.stage, - geom=geom, - obj_name=geom_name, - rgba=geom.rgba, - texture_file=texture_file, - ) - elif geom.type == mujoco.mjtGeom.mjGEOM_SPHERE: - usd_geom = component_module.USDSphereMesh( - stage=self.stage, - geom=geom, - obj_name=geom_name, - rgba=geom.rgba, - texture_file=texture_file, - ) - elif geom.type == mujoco.mjtGeom.mjGEOM_CAPSULE: - usd_geom = component_module.USDCapsule( - stage=self.stage, - geom=geom, - obj_name=geom_name, - rgba=geom.rgba, - texture_file=texture_file, - ) - elif geom.type == mujoco.mjtGeom.mjGEOM_ELLIPSOID: - usd_geom = component_module.USDEllipsoid( - stage=self.stage, - geom=geom, - obj_name=geom_name, - rgba=geom.rgba, - texture_file=texture_file, - ) - elif geom.type == mujoco.mjtGeom.mjGEOM_CYLINDER: - usd_geom = component_module.USDCylinderMesh( - stage=self.stage, - geom=geom, - obj_name=geom_name, - rgba=geom.rgba, - texture_file=texture_file, - ) - elif geom.type == mujoco.mjtGeom.mjGEOM_BOX: - usd_geom = component_module.USDCubeMesh( - stage=self.stage, - geom=geom, - obj_name=geom_name, - rgba=geom.rgba, - texture_file=texture_file, - ) else: - usd_geom = None + + mesh_config = shapes_module.mesh_config_generator( + name=geom_name, + geom_type=geom.type, + size=geom.size + ) + + usd_geom = component_module.USDPrimitiveMesh( + mesh_config=mesh_config, + stage=self.stage, + geom=geom, + obj_name=geom_name, + rgba=geom.rgba, + texture_file=texture_file, + ) self.geom_name2usd[geom_name] = usd_geom diff --git a/python/mujoco/usd/shapes.py b/python/mujoco/usd/shapes.py index e69de29b..2959e7da 100644 --- a/python/mujoco/usd/shapes.py +++ b/python/mujoco/usd/shapes.py @@ -0,0 +1,129 @@ +import copy +import mujoco +import pprint +import numpy as np +import open3d as o3d + +def mesh_config_generator( + name: str, + geom_type: mujoco.mjtGeom, + size: np.ndarray +): + + if geom_type == mujoco.mjtGeom.mjGEOM_PLANE: + return { + "name": name, + "box": { + "width": size[0] * 2 if size[0] > 0 else 100, + "height": size[1] * 2 if size[1] > 0 else 100, + "depth": 0.001, + "map_texture_to_each_face": True, + } + } + elif geom_type == mujoco.mjtGeom.mjGEOM_SPHERE: + return { + "name": name, + "sphere": { + "radius": float(size[0]) + } + } + elif geom_type == mujoco.mjtGeom.mjGEOM_CAPSULE: + cylinder = mesh_config_generator(name, mujoco.mjtGeom.mjGEOM_CYLINDER, size) + left_sphere = mesh_config_generator(name, mujoco.mjtGeom.mjGEOM_SPHERE, size) + right_sphere = copy.deepcopy(left_sphere) + left_sphere["sphere"]["transform"] = { + "translate": (0, 0, -size[2]) + } + right_sphere["sphere"]["transform"] = { + "translate": (0, 0, size[2]) + } + return { + "name": name, + "cylinder": cylinder["cylinder"], + "left_sphere": left_sphere["sphere"], + "right_sphere": right_sphere["sphere"], + } + elif geom_type == mujoco.mjtGeom.mjGEOM_ELLIPSOID: + sphere = mesh_config_generator(name, mujoco.mjtGeom.mjGEOM_SPHERE, [1.0]) + sphere["sphere"]["transform"] = { + "scale": tuple(size) + } + return { + "name": name, + "sphere": sphere["sphere"], + } + elif geom_type == mujoco.mjtGeom.mjGEOM_CYLINDER: + return { + "name": name, + "cylinder": { + "radius": size[0], + "height": size[2] * 2, + } + } + elif geom_type == mujoco.mjtGeom.mjGEOM_BOX: + return { + "name": name, + "box": { + "width": size[0] * 2, + "height": size[1] * 2, + "depth": size[2] * 2, + } + } + else: + raise NotImplemented(f"{geom_type} primitive geom type not implemented with USD integration") + + +def mesh_generator( + mesh_config: dict +): + + assert "name" in mesh_config + + mesh = None + + for shape, config in mesh_config.items(): + + if shape == "name": + continue + + if "box" in shape: + prim_mesh = o3d.geometry.TriangleMesh.create_box( + width=mesh_config[shape]["width"], + height=mesh_config[shape]["height"], + depth=mesh_config[shape]["depth"], + create_uv_map=True, + map_texture_to_each_face=True, + ) + elif "sphere" in shape: + prim_mesh = o3d.geometry.TriangleMesh.create_sphere( + radius=mesh_config[shape]["radius"], + create_uv_map=True + ) + elif "cylinder" in shape: + prim_mesh = o3d.geometry.TriangleMesh.create_cylinder( + radius=mesh_config[shape]["radius"], + height=mesh_config[shape]["height"], + create_uv_map=True, + ) + + if "transform" in config: + for transform, val in config["transform"].items(): + if transform == "translate": + prim_mesh.translate(val) + if transform == "scale": + prim_mesh.vertices = o3d.utility.Vector3dVector( + np.asarray(prim_mesh.vertices) * np.array(val)) + + if not mesh: + mesh = prim_mesh + else: + mesh += prim_mesh + + return mesh_config["name"], mesh + + + + + + +