Removing Open3D dependency
This commit is contained in:
Abhishek Joshi
+188
-41
@@ -14,34 +14,184 @@
|
||||
# ==============================================================================
|
||||
"""Built-in shapes for USD exporter."""
|
||||
|
||||
from typing import Dict, Any
|
||||
from typing import Dict, Any, Tuple
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
from open3d import open3d as o3d
|
||||
|
||||
def get_triangle_uvs(vertices: np.array, triangles: np.array):
|
||||
# (jabhi) asusming all mappings are 2d mapping temporarily
|
||||
triangle_uvs = np.array([
|
||||
[vertices[i][0], vertices[i][1]] for i in np.nditer(triangles)]
|
||||
)
|
||||
return triangle_uvs
|
||||
|
||||
def create_hemisphere(
|
||||
radius: float, theta_steps: int = 50, phi_steps: int = 50
|
||||
):
|
||||
"""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])
|
||||
class TriangleMesh():
|
||||
""" Store UV and geometry information for a primitve mesh
|
||||
"""
|
||||
def __init__(self,
|
||||
vertices: np.array,
|
||||
triangles: np.array,
|
||||
triangle_uvs: np.array):
|
||||
self.vertices = vertices
|
||||
self.triangles = triangles
|
||||
self.triangle_uvs = triangle_uvs
|
||||
|
||||
pcd = o3d.geometry.PointCloud()
|
||||
pcd.points = o3d.utility.Vector3dVector(points)
|
||||
@classmethod
|
||||
def create_box(
|
||||
cls, width: float, height: float, depth: float
|
||||
):
|
||||
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)
|
||||
|
||||
return TriangleMesh(vertices, triangles, triangle_uvs)
|
||||
|
||||
mesh = pcd.compute_convex_hull()[0]
|
||||
@classmethod
|
||||
def create_sphere(
|
||||
cls, radius: float, resolution: int
|
||||
):
|
||||
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])
|
||||
|
||||
return mesh
|
||||
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)
|
||||
|
||||
return TriangleMesh(vertices, triangles, triangle_uvs)
|
||||
|
||||
@classmethod
|
||||
def create_hemisphere(
|
||||
cls, radius: float, resolution: int
|
||||
):
|
||||
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)
|
||||
|
||||
return TriangleMesh(vertices, triangles, triangle_uvs)
|
||||
|
||||
@classmethod
|
||||
def create_cylinder(
|
||||
cls, radius: float, height: float, resolution: int
|
||||
):
|
||||
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)
|
||||
|
||||
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 = get_triangle_uvs(new_vertices, new_triangles)
|
||||
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]):
|
||||
"""Breaks a shape config into is subcomponent shapes."""
|
||||
@@ -79,20 +229,20 @@ 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(
|
||||
@@ -130,10 +280,9 @@ def mesh_config_generator(
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def mesh_generator(
|
||||
def mesh_factory(
|
||||
mesh_config: Dict[str, Any],
|
||||
resolution: int = 100,
|
||||
resolution: int = 20,
|
||||
):
|
||||
"""Generates a mesh given a config consisting of shapes."""
|
||||
assert "name" in mesh_config
|
||||
@@ -145,28 +294,29 @@ def mesh_generator(
|
||||
if "name" in shape:
|
||||
continue
|
||||
|
||||
prim_mesh = None
|
||||
|
||||
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,
|
||||
depth=mesh_config[shape]["depth"]
|
||||
)
|
||||
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,
|
||||
resolution=resolution
|
||||
)
|
||||
elif "sphere" in shape:
|
||||
prim_mesh = TriangleMesh.create_sphere(
|
||||
radius=mesh_config[shape]["radius"],
|
||||
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,
|
||||
resolution=resolution
|
||||
)
|
||||
else:
|
||||
raise ValueError("Shape not supported")
|
||||
@@ -182,10 +332,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