Merge pull request #4 from awesome-aj0123/shape-meshes

Shape meshes
This commit is contained in:
Abhishek Joshi
2024-04-10 14:22:08 -05:00
committed by GitHub
5 changed files with 244 additions and 337 deletions
+2 -1
View File
@@ -1,3 +1,4 @@
from .exporter import *
from .component import *
from .utils import *
from .utils import *
from .shapes import *
+42 -281
View File
@@ -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)
+6 -5
View File
@@ -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")
+21 -50
View File
@@ -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)
+173
View File
@@ -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