Abhishek Joshi
130 lines
3.4 KiB
Python
130 lines
3.4 KiB
Python
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
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|