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..48b3e251 100644 --- a/python/mujoco/usd/component.py +++ b/python/mujoco/usd/component.py @@ -15,7 +15,10 @@ from typing import List, Optional, Tuple import mujoco + import mujoco.usd.utils +import mujoco.usd.shapes as shapes_component + import numpy as np # TODO: b/288149332 - Remove once USD Python Binding works well with pytype. @@ -217,7 +220,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 +237,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 +256,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 +404,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 +416,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 +493,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..e8e8deb1 100644 --- a/python/mujoco/usd/demo.py +++ b/python/mujoco/usd/demo.py @@ -1,19 +1,20 @@ + 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") + m = mujoco.MjModel.from_xml_path("/Users/abhishek/Documents/research/mujoco/test/engine/testdata/catenary.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..ff5b24f3 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -15,7 +15,10 @@ import os import mujoco + +import mujoco.usd.shapes as shapes_module import mujoco.usd.component as component_module + import numpy as np import scipy import termcolor @@ -230,11 +233,13 @@ class USDExporter: def _load_geom(self, geom: mujoco.MjvGeom): geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid) + if not geom_name: + geom_name = f"unnamed_geom{geom.objid}" + assert geom_name not in self.geom_name2usd 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 +250,20 @@ 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 @@ -306,6 +275,8 @@ class USDExporter: 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 = f"unnamed_geom{geom.objid}" if geom_name not in self.geom_name2usd: self._load_geom(geom) diff --git a/python/mujoco/usd/shapes.py b/python/mujoco/usd/shapes.py index e69de29b..0cd6842e 100644 --- a/python/mujoco/usd/shapes.py +++ b/python/mujoco/usd/shapes.py @@ -0,0 +1,173 @@ +import copy +import mujoco +import pprint +import numpy as np +import open3d as o3d + +def create_hemisphere( + radius: float, + resolution: int = 20, + theta_steps: int = 50, + phi_steps: int = 50 +): + + points = [] + for i in range(phi_steps + 1): + phi = np.pi / 2 * i / phi_steps + for j in range(theta_steps + 1): + theta = 2 * np.pi * j / theta_steps + x = radius * np.sin(phi) * np.cos(theta) + y = radius * np.sin(phi) * np.sin(theta) + z = radius * np.cos(phi) + points.append([x, y, z]) + + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(points) + + mesh = pcd.compute_convex_hull()[0] + + return mesh + +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) + return { + "name": name, + "cylinder": cylinder["cylinder"], + "left_hemisphere": { + "radius": size[0], + "transform": { + "translate": (0, 0, -size[2]), + "rotate": (np.pi, 0, 0) + } + }, + "right_hemisphere": { + "radius": size[0], + "transform": { + "translate": (0, 0, size[2]) + } + }, + } + 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 "hemisphere" in shape: + prim_mesh = create_hemisphere( + radius=mesh_config[shape]["radius"] + ) + 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: + + if "rotate" in config["transform"]: + R = mesh.get_rotation_matrix_from_xyz(config["transform"]["rotate"]) + prim_mesh.rotate(R, center=(0, 0, 0)) + if "scale" in config["transform"]: + prim_mesh.vertices = o3d.utility.Vector3dVector( + np.asarray(prim_mesh.vertices) * np.array(config["transform"]["scale"])) + if "translate" in config["transform"]: + prim_mesh.translate(config["transform"]["translate"]) + + # for transform, val in config["transform"].items(): + # if transform == "translate": + # prim_mesh.translate(val) + # elif transform == "scale": + # prim_mesh.vertices = o3d.utility.Vector3dVector( + # np.asarray(prim_mesh.vertices) * np.array(val)) + # elif transform == "rotate": + # R = mesh.get_rotation_matrix_from_xyz(val) + # prim_mesh.rotate(R, center=(0, 0, 0)) + + if not mesh: + mesh = prim_mesh + else: + mesh += prim_mesh + + return mesh_config["name"], mesh + + + + + + +