chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Merge pull request #1827 from abhihjoshi:open3d-removal

Browse Source 
PiperOrigin-RevId: 662510404
Change-Id: I79b356eecddb3402c444cb16a305516bd724aef6
This commit is contained in:
Copybara-Service
2024-08-13 07:33:23 -07:00
parent 39896f8064 e0483c2f86
commit 98cd49d529
5 changed files with 401 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/usd/camera.py
+2 -2
View File
@@ -14,7 +14,7 @@
# ==============================================================================
"""Camera handling for USD exporter."""
import mujoco.usd.utils as utils_component
import mujoco.usd.utils as utils_module
import numpy as np
@@ -49,7 +49,7 @@ class USDCamera:
def update(self, cam_pos: np.ndarray, cam_mat: np.ndarray, frame: int):
"""Updates the position and orientation of the camera in the scene."""
transformation_mat = utils_component.create_transform_matrix(
transformation_mat = utils_module.create_transform_matrix(
rotation_matrix=cam_mat, translation_vector=cam_pos
).T
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
python/mujoco/usd/demo.py
+4
View File
@@ -42,6 +42,10 @@ def generate_usd_trajectory(local_args):
if exp.frame_count < d.time * local_args.framerate:
exp.update_scene(data=d)
exp.add_light(pos=(0, 0, 0),
intensity=2000,
light_type='dome')
exp.save_scene(filetype=local_args.export_extension)
python/mujoco/usd/exporter.py
+13 -12
View File
@@ -27,7 +27,6 @@ from PIL import Image as im
from PIL import ImageOps
import scipy
import termcolor
import tqdm
# TODO: b/288149332 - Remove once USD Python Binding works well with pytype.
# pytype: disable=module-attr
@@ -209,7 +208,7 @@ class USDExporter:
"""Load textures."""
data_adr = 0
self.texture_files = []
for texture_id in tqdm.tqdm(range(self.model.ntex)):
for texture_id in range(self.model.ntex):
texture_height = self.model.tex_height[texture_id]
texture_width = self.model.tex_width[texture_id]
texture_nchannel = self.model.tex_nchannel[texture_id]
@@ -251,13 +250,13 @@ class USDExporter:
assert geom_name not in self.geom_names
texture_file = (
self.texture_files[
self.model.mat_texid[geom.matid][mujoco.mjTEXROLE_RGB]
]
if geom.matid != -1
else None
)
if geom.matid == -1:
geom_textures = []
else:
geom_textures = [
(self.texture_files[i], self.model.tex_type[i]) if i != -1 else None
for i in self.model.mat_texid[geom.matid]
]
# handling meshes in our scene
if geom.type == mujoco.mjtGeom.mjGEOM_MESH:
@@ -268,7 +267,7 @@ class USDExporter:
obj_name=geom_name,
dataid=self.model.geom_dataid[geom.objid],
rgba=geom.rgba,
texture_file=texture_file,
geom_textures=geom_textures,
)
else:
# handling tendons in our scene
@@ -282,10 +281,11 @@ class USDExporter:
usd_geom = object_module.USDTendon(
mesh_config=mesh_config,
stage=self.stage,
model=self.model,
geom=geom,
obj_name=geom_name,
rgba=geom.rgba,
texture_file=texture_file,
geom_textures=geom_textures,
)
# handling primitives in our scene
else:
@@ -297,10 +297,11 @@ class USDExporter:
usd_geom = object_module.USDPrimitiveMesh(
mesh_config=mesh_config,
stage=self.stage,
model=self.model,
geom=geom,
obj_name=geom_name,
rgba=geom.rgba,
texture_file=texture_file,
geom_textures=geom_textures,
)
self.geom_names.add(geom_name)
python/mujoco/usd/objects.py
+104 -60
View File
@@ -16,13 +16,14 @@
import abc
import collections
from typing import Optional, Dict, Any
from typing import Any, Dict, List, Optional, Tuple
import mujoco
import mujoco.usd.shapes as shapes_component
import mujoco.usd.utils as utils_component
import mujoco.usd.shapes as shapes_module
import mujoco.usd.utils as utils_module
import numpy as np
# TODO: b/288149332 - Remove once USD Python Binding works well with pytype.
# pytype: disable=module-attr
from pxr import Gf
@@ -50,16 +51,18 @@ class USDObject(abc.ABC):
def __init__(
self,
stage: Usd.Stage,
model: mujoco.MjModel,
geom: mujoco.MjvGeom,
obj_name: str,
rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None,
geom_textures: List[Optional[Tuple[str, mujoco.mjtTexture]]] = None
):
self.stage = stage
self.model = model
self.geom = geom
self.obj_name = obj_name
self.rgba = rgba
self.texture_file = texture_file
self.geom_textures = geom_textures
self.xform_path = f"/World/Mesh_Xform_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, self.xform_path)
@@ -117,7 +120,7 @@ class USDObject(abc.ABC):
# setting the image texture attributes
image_shader.CreateIdAttr("UsdUVTexture")
image_shader.CreateInput("file", Sdf.ValueTypeNames.Asset).Set(
self.texture_file
self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB][0]
)
image_shader.CreateInput("sourceColorSpace", Sdf.ValueTypeNames.Token).Set(
"sRGB"
@@ -185,7 +188,7 @@ class USDObject(abc.ABC):
scale: Optional[np.ndarray] = None,
):
"""Updates the position and orientation of an object."""
transformation_mat = utils_component.create_transform_matrix(
transformation_mat = utils_module.create_transform_matrix(
rotation_matrix=mat, translation_vector=pos
).T
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
@@ -222,11 +225,10 @@ class USDMesh(USDObject):
obj_name: str,
dataid: int,
rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None,
geom_textures: List[Optional[Tuple[str, mujoco.mjtTexture]]] = None
):
super().__init__(stage, geom, obj_name, rgba, texture_file)
super().__init__(stage, model, geom, obj_name, rgba, geom_textures)
self.model = model
self.dataid = dataid
mesh_path = f"{self.xform_path}/Mesh_{obj_name}"
@@ -241,15 +243,19 @@ class USDMesh(USDObject):
)
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(mesh_facetexcoord.tolist()))
if self.texture_file:
if (
geom.matid != -1
and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]
):
# 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(mesh_facetexcoord.tolist()))
self.attach_image_material(self.usd_mesh)
else:
self.attach_solid_material(self.usd_mesh)
@@ -314,12 +320,13 @@ class USDPrimitiveMesh(USDObject):
self,
mesh_config: Dict[Any, Any],
stage: Usd.Stage,
model: mujoco.MjModel,
geom: mujoco.MjvGeom,
obj_name: str,
rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None,
geom_textures: List[Optional[Tuple[str, mujoco.mjtTexture]]] = None
):
super().__init__(stage, geom, obj_name, rgba, texture_file)
super().__init__(stage, model, geom, obj_name, rgba, geom_textures)
self.mesh_config = mesh_config
self.prim_mesh = self.generate_primitive_mesh()
@@ -334,42 +341,66 @@ class USDPrimitiveMesh(USDObject):
[3 for _ in range(mesh_facenum)]
)
self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
# setting mesh uv properties
mesh_texcoord, _ = 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(list(range(mesh_facenum * 3))))
self._set_refinement_properties(self.usd_prim)
if self.texture_file:
if (
geom.matid != -1
and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]
):
# setting mesh uv properties
mesh_texcoord, _ = 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(list(range(mesh_facenum * 3))))
self.attach_image_material(self.usd_mesh)
else:
self.attach_solid_material(self.usd_mesh)
def generate_primitive_mesh(self):
"""Generates the mesh for the primitive USD object."""
_, prim_mesh = shapes_component.mesh_generator(self.mesh_config)
tex_role = mujoco.mjtTextureRole
geom_rgb_texture = (
self.geom_textures[tex_role.mjTEXROLE_RGB]
if self.geom_textures
else None
)
texture_type = geom_rgb_texture[1] if geom_rgb_texture else None
_, prim_mesh = shapes_module.mesh_factory(self.mesh_config, texture_type)
prim_mesh.translate(-prim_mesh.get_center())
return prim_mesh
def _get_uv_geometry(self):
assert self.prim_mesh
assert self.prim_mesh and self.prim_mesh.triangle_uvs is not None
mesh_texcoord = np.array(self.prim_mesh.triangle_uvs)
mesh_facetexcoord = np.asarray(self.prim_mesh.triangles)
# TODO(etom): bring back support for rescaling the texture coordinates.
tex_role = mujoco.mjtTextureRole
geom_rgb_texture = self.geom_textures[tex_role.mjTEXROLE_RGB][1]
if geom_rgb_texture == mujoco.mjtTexture.mjTEXTURE_2D:
s_scale, t_scale = self.model.mat_texrepeat[self.geom.matid]
if self.model.mat_texuniform[self.geom.matid]:
if self.geom.size[0] > 0:
s_scale *= self.geom.size[0]
if self.geom.size[1] > 0:
t_scale *= self.geom.size[1]
mesh_texcoord[:, 0] *= s_scale / (self.geom.size[0] * 2)
mesh_texcoord[:, 1] *= t_scale / (self.geom.size[1] * 2)
return mesh_texcoord, mesh_facetexcoord.flatten()
def _get_mesh_geometry(self):
assert self.prim_mesh
# get mesh geometry from the open3d mesh model
# get mesh geometry
mesh_vert = np.asarray(self.prim_mesh.vertices)
mesh_face = np.asarray(self.prim_mesh.triangles)
@@ -383,12 +414,13 @@ class USDTendon(USDObject):
self,
mesh_config: Dict[Any, Any],
stage: Usd.Stage,
model: mujoco.MjModel,
geom: mujoco.MjvGeom,
obj_name: str,
rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None,
geom_textures: List[Optional[Tuple[str, mujoco.mjtTexture]]] = None
):
super().__init__(stage, geom, obj_name, rgba, texture_file)
super().__init__(stage, model, geom, obj_name, rgba, geom_textures)
self.mesh_config = mesh_config
self.tendon_parts = self.generate_primitive_mesh()
@@ -407,6 +439,7 @@ class USDTendon(USDObject):
# setting mesh geometry properties for each of the parts in the tendon
part_geometries = self._get_mesh_geometry()
part_geometry = None
for name, part_geometry in part_geometries.items():
self.usd_refs[name]["usd_mesh"].GetPointsAttr().Set(
part_geometry["mesh_vert"]
@@ -418,33 +451,43 @@ class USDTendon(USDObject):
part_geometry["mesh_face"]
)
# setting uv properties for each of the parts in the tendon
part_uv_geometries = self._get_uv_geometry()
for name, part_uv_geometry in part_uv_geometries.items():
self.texcoords = UsdGeom.PrimvarsAPI(
self.usd_refs[name]["usd_mesh"]
).CreatePrimvar(
"UVMap",
Sdf.ValueTypeNames.TexCoord2fArray,
UsdGeom.Tokens.faceVarying,
)
self.texcoords.Set(part_uv_geometry["mesh_texcoord"])
self.texcoords.SetIndices(
Vt.IntArray(list(range(part_geometry["mesh_facenum"] * 3)))
)
for _, ref in self.usd_refs.items():
self._set_refinement_properties(ref["usd_prim"])
if self.texture_file:
self.attach_image_material(ref["usd_mesh"])
else:
tex_role = mujoco.mjtTextureRole
if geom.matid != -1 and self.geom_textures[tex_role.mjTEXROLE_RGB]:
# setting uv properties for each of the parts in the tendon
part_uv_geometries = self._get_uv_geometry()
for name, part_uv_geometry in part_uv_geometries.items():
self.texcoords = UsdGeom.PrimvarsAPI(
self.usd_refs[name]["usd_mesh"]
).CreatePrimvar(
"UVMap",
Sdf.ValueTypeNames.TexCoord2fArray,
UsdGeom.Tokens.faceVarying,
)
self.texcoords.Set(part_uv_geometry["mesh_texcoord"])
self.texcoords.SetIndices(
Vt.IntArray(list(range(part_geometry["mesh_facenum"] * 3)))
)
for _, ref in self.usd_refs.items():
self._set_refinement_properties(ref["usd_prim"])
self.attach_image_material(ref["usd_mesh"])
else:
for _, ref in self.usd_refs.items():
self._set_refinement_properties(ref["usd_prim"])
self.attach_solid_material(ref["usd_mesh"])
def generate_primitive_mesh(self):
"""Generates the tendon mesh using primitives."""
mesh_parts = {}
geom_rgb_texture = (
self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]
if self.geom_textures
else None
)
texture_type = geom_rgb_texture[1] if geom_rgb_texture else None
for part_config in self.mesh_config:
mesh_name, prim_mesh = shapes_component.mesh_generator(part_config)
mesh_name, prim_mesh = shapes_module.mesh_factory(
part_config, texture_type
)
prim_mesh.translate(-prim_mesh.get_center())
mesh_parts[mesh_name] = prim_mesh
return mesh_parts
@@ -452,6 +495,7 @@ class USDTendon(USDObject):
def _get_uv_geometry(self):
part_uv_geometries = collections.defaultdict(dict)
for name, mesh in self.tendon_parts.items():
assert mesh.triangle_uvs is not None
mesh_texcoord = np.array(mesh.triangle_uvs)
mesh_facetexcoord = np.asarray(mesh.triangles)
part_uv_geometries[name] = {
@@ -463,7 +507,7 @@ class USDTendon(USDObject):
def _get_mesh_geometry(self):
part_geometries = collections.defaultdict(dict)
for name, mesh in self.tendon_parts.items():
# get mesh geometry from the open3d mesh model
# get mesh geometry
mesh_vert = np.asarray(mesh.vertices)
mesh_face = np.asarray(mesh.triangles)
part_geometries[name] = {
python/mujoco/usd/shapes.py
+278 -38
View File
@@ -14,33 +14,270 @@
# ==============================================================================
"""Built-in shapes for USD exporter."""
from typing import Dict, Any
from typing import Any, Dict, Optional, Tuple, Union
import mujoco
import numpy as np
from open3d import open3d as o3d
def create_hemisphere(
radius: float, theta_steps: int = 50, phi_steps: int = 50
def get_triangle_uvs(
vertices: np.ndarray,
triangles: np.ndarray,
texture_type: Optional[mujoco.mjtTexture]
):
"""Creates a hemisphere mesh from a point cloud."""
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])
"""Returns UV coordinates for a given mesh."""
if not texture_type:
return None
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(points)
triangle_uvs = []
if texture_type == mujoco.mjtTexture.mjTEXTURE_2D:
triangle_uvs = [
[vertices[i][0], vertices[i][1]] for i in np.nditer(triangles)
]
mesh = pcd.compute_convex_hull()[0]
elif texture_type == mujoco.mjtTexture.mjTEXTURE_CUBE:
center = np.mean(vertices, axis=0)
for vertex_id in np.nditer(triangles):
x, y, z = vertices[vertex_id] - center
return mesh
abs_x, abs_y, abs_z = abs(x), abs(y), abs(z)
u = 0
v = 0
if x > 0 and abs_x >= abs_y and abs_x >= abs_z:
u = -z / abs_x
v = y / abs_x
elif x <= 0 and abs_x >= abs_y and abs_x >= abs_z:
u = z / abs_x
v = y / abs_x
elif y > 0 and abs_y >= abs_x and abs_y >= abs_z:
u = x / abs_y
v = -z / abs_y
elif y <= 0 and abs_y >= abs_x and abs_y >= abs_z:
u = x / abs_y
v = z / abs_y
elif z > 0 and abs_z >= abs_x and abs_z >= abs_y:
u = x / abs_z
v = y / abs_z
elif z <= 0 and abs_z >= abs_x and abs_z >= abs_y:
u = -x / abs_z
v = y / abs_z
u = (u + 1.0) / 2.0
v = (v + 1.0) / 2.0
v /= 6
assert 0 <= u and u <= 1 and 0 <= v and v <= 1
triangle_uvs.append([u, v])
elif texture_type == mujoco.mjtTexture.mjTEXTURE_SKYBOX:
# defaults to 2D mapping temporarily
triangle_uvs = [
[vertices[i][0], vertices[i][1]] for i in np.nditer(triangles)
]
return np.array(triangle_uvs)
class TriangleMesh:
"""Store UV and geometry information for a primitive mesh.
Attributes:
vertices: A numpy array of vertices.
triangles: A numpy array of triangles.
triangle_uvs: A numpy array of UV coordinates.
"""
def __init__(self,
vertices: np.ndarray,
triangles: np.ndarray,
triangle_uvs: np.ndarray):
"""Creates a TriangleMesh object.
Args:
vertices: A numpy array of vertices.
triangles: A numpy array of triangles.
triangle_uvs: A numpy array of UV coordinates.
"""
self.vertices = vertices
self.triangles = triangles
self.triangle_uvs = triangle_uvs
@classmethod
def create_box(
cls,
width: float,
height: float,
depth: float,
texture_type: Optional[mujoco.mjtTexture]
) -> TriangleMesh:
"""Creates a box."""
vertices = np.array([[0.0, 0.0, 0.0],
[width, 0.0, 0.0],
[0.0, 0.0, depth],
[width, 0.0, depth],
[0.0, height, 0.0],
[width, height, 0.0],
[0.0, height, depth],
[width, height, depth]])
triangles = np.array([[4, 7, 5],
[4, 6, 7],
[0, 2, 4],
[2, 6, 4],
[0, 1, 2],
[1, 3, 2],
[1, 5, 7],
[1, 7, 3],
[2, 3, 7],
[2, 7, 6],
[0, 4, 1],
[1, 4, 5]])
triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type)
return TriangleMesh(vertices, triangles, triangle_uvs)
@classmethod
def create_sphere(
cls,
radius: float,
texture_type: Optional[mujoco.mjtTexture],
resolution: int
) -> TriangleMesh:
"""Creates a sphere."""
vertices = []
triangles = []
for i in range(2*resolution + 1):
phi = np.pi * i / (2*resolution)
for j in range(resolution + 1):
theta = 2 * np.pi * j / resolution
x = radius * np.sin(phi) * np.cos(theta)
y = radius * np.sin(phi) * np.sin(theta)
z = radius * np.cos(phi)
vertices.append([x, y, z])
for i in range(2*resolution):
for j in range(resolution):
first = i * (resolution + 1) + j
second = first + resolution + 1
triangles.append([first, second, first + 1])
triangles.append([second, second + 1, first + 1])
vertices = np.array(vertices)
triangles = np.array(triangles)
triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type)
return TriangleMesh(vertices, triangles, triangle_uvs)
@classmethod
def create_hemisphere(
cls,
radius: float,
texture_type: Optional[mujoco.mjtTexture],
resolution: int,
) -> TriangleMesh:
"""Creates a hemisphere."""
vertices = []
triangles = []
for i in range(resolution + 1):
phi = np.pi / 2 * i / (resolution)
for j in range(resolution + 1):
theta = 2 * np.pi * j / resolution
x = radius * np.sin(phi) * np.cos(theta)
y = radius * np.sin(phi) * np.sin(theta)
z = radius * np.cos(phi)
vertices.append([x, y, z])
vertices.append([0, 0, 0])
for i in range(resolution):
for j in range(resolution):
first = i * (resolution + 1) + j
second = first + resolution + 1
triangles.append([first, second, first + 1])
triangles.append([second, second + 1, first + 1])
for i in range(resolution):
first = resolution * (resolution + 1) + i
triangles.append([first, first + 1, len(vertices) - 1])
vertices = np.array(vertices)
triangles = np.array(triangles)
triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type)
return TriangleMesh(vertices, triangles, triangle_uvs)
@classmethod
def create_cylinder(
cls,
radius: float,
height: float,
texture_type: Optional[mujoco.mjtTexture],
resolution: int
) -> TriangleMesh:
"""Creates a cylinder."""
vertices = []
triangles = []
# adding all the vertices for the cylinder including
# two center vertices at ends
for i in range(2):
z = 0 if i == 0 else height
for j in range(resolution + 1):
theta = 2 * np.pi * j / resolution
x = radius * np.cos(theta)
y = radius * np.sin(theta)
vertices.append([x, y, z])
vertices.append([0, 0, 0])
vertices.append([0, 0, height])
# constructing the end faces for the cylinder
for i in range(2):
for j in range(resolution):
first = (resolution + 1) * i + j
triangles.append([first, first + 1, len(vertices) - (2 - i)])
# constructing side of cylinder
for i in range(resolution):
second = resolution + 1 + i
triangles.append([i, second, second + 1])
triangles.append([i, i + 1, second + 1])
vertices = np.array(vertices)
triangles = np.array(triangles)
triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type)
return TriangleMesh(vertices, triangles, triangle_uvs)
def translate(self, translation: np.array):
self.vertices = self.vertices + translation
def rotate(self, rotation: np.array, center: Tuple[float, ...]):
translated_point = self.vertices - center
self.vertices = np.dot(translated_point, rotation) + center
def scale(self, scale: np.array):
self.vertices = self.vertices * scale
def get_center(self):
center = np.mean(self.vertices, axis=0)
return center
def __add__(self, other):
if isinstance(other, TriangleMesh):
new_vertices = np.vstack((self.vertices, other.vertices))
other_triangles = other.triangles + len(self.vertices)
new_triangles = np.vstack((self.triangles, other_triangles))
new_triangle_uvs = None
if self.triangle_uvs is not None:
new_triangle_uvs = np.vstack((self.triangle_uvs, other.triangle_uvs))
return TriangleMesh(new_vertices, new_triangles, new_triangle_uvs)
raise TypeError(f"Cannot add TriangleMesh with {type(other)}")
def decouple_config(config: Dict[str, Any]):
@@ -60,7 +297,7 @@ def decouple_config(config: Dict[str, Any]):
def mesh_config_generator(
name: str,
geom_type: int | mujoco.mjtGeom,
geom_type: Union[int, mujoco.mjtGeom],
size: np.ndarray,
decouple: bool = False,
):
@@ -79,20 +316,22 @@ def mesh_config_generator(
config = {"name": name, "sphere": {"radius": float(size[0])}}
elif geom_type == mujoco.mjtGeom.mjGEOM_CAPSULE:
cylinder = mesh_config_generator(name, mujoco.mjtGeom.mjGEOM_CYLINDER, size)
cylinder["cylinder"]["transform"] = {
"transform": {"translate": (0, 0, size[2])}
}
config = {
"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])},
},
"transform": {"translate": (0, 0, 2*size[2])},
}
}
elif geom_type == mujoco.mjtGeom.mjGEOM_ELLIPSOID:
sphere = mesh_config_generator(
@@ -131,8 +370,9 @@ def mesh_config_generator(
return config
def mesh_generator(
def mesh_factory(
mesh_config: Dict[str, Any],
texture_type: Optional[mujoco.mjtTexture],
resolution: int = 100,
):
"""Generates a mesh given a config consisting of shapes."""
@@ -146,27 +386,30 @@ def mesh_generator(
continue
if "box" in shape:
prim_mesh = o3d.geometry.TriangleMesh.create_box(
prim_mesh = 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,
texture_type=texture_type
)
elif "hemisphere" in shape:
prim_mesh = create_hemisphere(radius=mesh_config[shape]["radius"])
elif "sphere" in shape:
prim_mesh = o3d.geometry.TriangleMesh.create_sphere(
prim_mesh = TriangleMesh.create_hemisphere(
radius=mesh_config[shape]["radius"],
resolution=resolution,
create_uv_map=True,
texture_type=texture_type,
resolution=resolution
)
elif "sphere" in shape:
prim_mesh = TriangleMesh.create_sphere(
radius=mesh_config[shape]["radius"],
texture_type=texture_type,
resolution=resolution
)
elif "cylinder" in shape:
prim_mesh = o3d.geometry.TriangleMesh.create_cylinder(
prim_mesh = TriangleMesh.create_cylinder(
radius=mesh_config[shape]["radius"],
height=mesh_config[shape]["height"],
resolution=resolution,
create_uv_map=True,
texture_type=texture_type,
resolution=resolution
)
else:
raise ValueError("Shape not supported")
@@ -182,10 +425,7 @@ def mesh_generator(
rotation = rotation.reshape((3, 3))
prim_mesh.rotate(rotation, 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"])
)
prim_mesh.scale(config["transform"]["scale"])
if "translate" in config["transform"]:
prim_mesh.translate(config["transform"]["translate"])