Merge pull request #1827 from abhihjoshi:open3d-removal
PiperOrigin-RevId: 662510404 Change-Id: I79b356eecddb3402c444cb16a305516bd724aef6
This commit is contained in:
Copybara-Service
+278
-38
@@ -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"])
|
||||
|
||||
|
||||
Reference in New Issue
Block a user