@@ -1,3 +1,4 @@
|
|||||||
from .exporter import *
|
from .exporter import *
|
||||||
from .component import *
|
from .component import *
|
||||||
from .utils import *
|
from .utils import *
|
||||||
|
from .shapes import *
|
||||||
+42
-281
@@ -15,7 +15,10 @@
|
|||||||
from typing import List, Optional, Tuple
|
from typing import List, Optional, Tuple
|
||||||
|
|
||||||
import mujoco
|
import mujoco
|
||||||
|
|
||||||
import mujoco.usd.utils
|
import mujoco.usd.utils
|
||||||
|
import mujoco.usd.shapes as shapes_component
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
|
||||||
# TODO: b/288149332 - Remove once USD Python Binding works well with pytype.
|
# 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)
|
UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl)
|
||||||
|
|
||||||
def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int):
|
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
|
rotation_matrix=mat, translation_vector=pos
|
||||||
).T
|
).T
|
||||||
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
|
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
|
||||||
@@ -234,12 +237,14 @@ class USDPrimitiveMesh:
|
|||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
|
mesh_config: List[dict],
|
||||||
stage: Usd.Stage,
|
stage: Usd.Stage,
|
||||||
geom: mujoco.MjvGeom,
|
geom: mujoco.MjvGeom,
|
||||||
obj_name: str,
|
obj_name: str,
|
||||||
rgba: np.ndarray = np.array([1, 1, 1, 1]),
|
rgba: np.ndarray = np.array([1, 1, 1, 1]),
|
||||||
texture_file: Optional[str] = None,
|
texture_file: Optional[str] = None,
|
||||||
):
|
):
|
||||||
|
self.mesh_config = mesh_config
|
||||||
self.stage = stage
|
self.stage = stage
|
||||||
self.geom = geom
|
self.geom = geom
|
||||||
self.obj_name = obj_name
|
self.obj_name = obj_name
|
||||||
@@ -251,6 +256,40 @@ class USDPrimitiveMesh:
|
|||||||
self.prim_mesh = None
|
self.prim_mesh = None
|
||||||
self.transform_op = Gf.Matrix4d(1.)
|
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):
|
def _set_refinement_properties(self):
|
||||||
self.usd_prim.GetAttribute("subdivisionScheme").Set("none")
|
self.usd_prim.GetAttribute("subdivisionScheme").Set("none")
|
||||||
|
|
||||||
@@ -365,7 +404,7 @@ class USDPrimitiveMesh:
|
|||||||
UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl)
|
UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl)
|
||||||
|
|
||||||
def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int):
|
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
|
rotation_matrix=mat, translation_vector=pos
|
||||||
).T
|
).T
|
||||||
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
|
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
|
||||||
@@ -377,284 +416,6 @@ class USDPrimitiveMesh:
|
|||||||
else:
|
else:
|
||||||
self.usd_prim.GetAttribute("visibility").Set("invisible", frame)
|
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:
|
class USDSphereLight:
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
@@ -732,7 +493,7 @@ class USDCamera:
|
|||||||
|
|
||||||
def update(self, cam_pos: np.ndarray, cam_mat: np.ndarray, frame: int):
|
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
|
rotation_matrix=cam_mat, translation_vector=cam_pos
|
||||||
).T
|
).T
|
||||||
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
|
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
|
||||||
|
|||||||
@@ -1,19 +1,20 @@
|
|||||||
|
|
||||||
import mujoco
|
import mujoco
|
||||||
from mujoco.usd import exporter
|
from mujoco.usd import exporter
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|
||||||
# load a model to mujoco
|
# 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)
|
d = mujoco.MjData(m)
|
||||||
|
|
||||||
# create an instance of the USDExporter
|
# create an instance of the USDExporter
|
||||||
exp = exporter.USDExporter(model=m)
|
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="usd")
|
||||||
|
|
||||||
exp.save_scene(filetype="usda")
|
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -15,7 +15,10 @@
|
|||||||
import os
|
import os
|
||||||
|
|
||||||
import mujoco
|
import mujoco
|
||||||
|
|
||||||
|
import mujoco.usd.shapes as shapes_module
|
||||||
import mujoco.usd.component as component_module
|
import mujoco.usd.component as component_module
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import scipy
|
import scipy
|
||||||
import termcolor
|
import termcolor
|
||||||
@@ -230,11 +233,13 @@ class USDExporter:
|
|||||||
def _load_geom(self, geom: mujoco.MjvGeom):
|
def _load_geom(self, geom: mujoco.MjvGeom):
|
||||||
|
|
||||||
geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid)
|
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
|
assert geom_name not in self.geom_name2usd
|
||||||
|
|
||||||
texture_file = self.texture_files[geom.texid] if geom.texid != -1 else None
|
texture_file = self.texture_files[geom.texid] if geom.texid != -1 else None
|
||||||
|
|
||||||
# handles meshes in scene
|
|
||||||
if geom.type == mujoco.mjtGeom.mjGEOM_MESH:
|
if geom.type == mujoco.mjtGeom.mjGEOM_MESH:
|
||||||
usd_geom = component_module.USDMesh(
|
usd_geom = component_module.USDMesh(
|
||||||
stage=self.stage,
|
stage=self.stage,
|
||||||
@@ -245,56 +250,20 @@ class USDExporter:
|
|||||||
rgba=geom.rgba,
|
rgba=geom.rgba,
|
||||||
texture_file=texture_file,
|
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:
|
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
|
self.geom_name2usd[geom_name] = usd_geom
|
||||||
|
|
||||||
@@ -306,6 +275,8 @@ class USDExporter:
|
|||||||
for i in range(self.scene.ngeom):
|
for i in range(self.scene.ngeom):
|
||||||
geom = self.scene.geoms[i]
|
geom = self.scene.geoms[i]
|
||||||
geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid)
|
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:
|
if geom_name not in self.geom_name2usd:
|
||||||
self._load_geom(geom)
|
self._load_geom(geom)
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user