146 lines
4.7 KiB
Python
146 lines
4.7 KiB
Python
import os
|
|
import numpy as np
|
|
import trimesh
|
|
|
|
|
|
def get_mesh_bb(mesh):
|
|
"""
|
|
:param mesh - trimesh mesh object
|
|
returns bb_min (3), bb_max (3)
|
|
"""
|
|
bb_min = mesh.vertices.max(axis=0)
|
|
bb_max = mesh.vertices.min(axis=0)
|
|
return bb_min, bb_max
|
|
|
|
|
|
def get_scene_bb(meshes):
|
|
"""
|
|
:param mesh_seqs - (potentially nested) list of trimesh objects
|
|
returns bb_min (3), bb_max (3)
|
|
"""
|
|
if isinstance(meshes, trimesh.Trimesh):
|
|
return get_mesh_bb(meshes)
|
|
|
|
bb_mins, bb_maxs = zip(*[get_scene_bb(mesh) for mesh in meshes])
|
|
bb_mins = np.stack(bb_mins, axis=0)
|
|
bb_maxs = np.stack(bb_maxs, axis=0)
|
|
return bb_mins.min(axis=0), bb_maxs.max(axis=0)
|
|
|
|
|
|
def make_batch_mesh(verts, faces, colors):
|
|
"""
|
|
convenience function to make batch of meshes
|
|
meshs have same faces in batch, verts have same color in mesh
|
|
:param verts (B, V, 3)
|
|
:param faces (F, 3)
|
|
:param colors (B, 3)
|
|
"""
|
|
B, V, _ = verts.shape
|
|
return [make_mesh(verts[b], faces, colors[b, None].expand(V, -1)) for b in range(B)]
|
|
|
|
|
|
def make_mesh(verts, faces, colors=None, yup=True):
|
|
"""
|
|
create a trimesh object for the faces and vertices
|
|
:param verts (V, 3) tensor
|
|
:param faces (F, 3) tensor
|
|
:param colors (optional) (V, 3) tensor
|
|
:param yup (optional bool) whether or not to save with Y up
|
|
"""
|
|
verts = verts.detach().cpu().numpy()
|
|
faces = faces.detach().cpu().numpy()
|
|
if yup:
|
|
verts = np.array([1, -1, -1])[None, :] * verts
|
|
if colors is None:
|
|
colors = np.ones_like(verts) * 0.5
|
|
else:
|
|
colors = colors.detach().cpu().numpy()
|
|
return trimesh.Trimesh(
|
|
vertices=verts, faces=faces, vertex_colors=colors, process=False
|
|
)
|
|
|
|
|
|
def save_mesh_scenes(out_dir, scenes):
|
|
"""
|
|
:param scenes, list of scenes (list of meshes)
|
|
"""
|
|
assert isinstance(scenes, list)
|
|
assert isinstance(scenes[0], list)
|
|
B = len(scenes[0])
|
|
if B == 1:
|
|
save_meshes_to_obj(out_dir, [x[0] for x in scenes])
|
|
else:
|
|
save_scenes_to_glb(out_dir, scenes)
|
|
|
|
|
|
def save_scenes_to_glb(out_dir, scenes):
|
|
"""
|
|
Saves a list of scenes (list of meshes) each to glb files
|
|
"""
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
for t, meshes in enumerate(scenes):
|
|
save_meshes_to_glb(f"{out_dir}/scene_{t:03d}.glb", meshes)
|
|
|
|
|
|
def save_meshes_to_glb(path, meshes, names=None):
|
|
"""
|
|
put trimesh meshes in a scene and export to glb
|
|
"""
|
|
if names is not None:
|
|
assert len(meshes) == len(names)
|
|
|
|
scene = trimesh.Scene()
|
|
for i, mesh in enumerate(meshes):
|
|
name = f"mesh_{i:03d}" if names is None else names[i]
|
|
scene.add_geometry(mesh, node_name=name)
|
|
|
|
with open(path, "wb") as f:
|
|
f.write(trimesh.exchange.gltf.export_glb(scene, include_normals=True))
|
|
|
|
|
|
def save_meshes_to_obj(out_dir, meshes, names=None):
|
|
if names is not None:
|
|
assert len(meshes) == len(names)
|
|
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
for i, mesh in enumerate(meshes):
|
|
name = f"mesh_{i:03d}" if names is None else names[i]
|
|
path = os.path.join(out_dir, f"{name}.obj")
|
|
with open(path, "w") as f:
|
|
mesh.export(f, file_type="obj", include_color=False, include_normals=True)
|
|
|
|
def vertices_to_trimesh(vertices, faces, mesh_base_color=(1.0, 1.0, 0.9),
|
|
rot_axis=[1,0,0], rot_angle=0, is_right=1):
|
|
# material = pyrender.MetallicRoughnessMaterial(
|
|
# metallicFactor=0.0,
|
|
# alphaMode='OPAQUE',
|
|
# baseColorFactor=(*mesh_base_color, 1.0))
|
|
faces_new = np.array([[92, 38, 234],
|
|
[234, 38, 239],
|
|
[38, 122, 239],
|
|
[239, 122, 279],
|
|
[122, 118, 279],
|
|
[279, 118, 215],
|
|
[118, 117, 215],
|
|
[215, 117, 214],
|
|
[117, 119, 214],
|
|
[214, 119, 121],
|
|
[119, 120, 121],
|
|
[121, 120, 78],
|
|
[120, 108, 78],
|
|
[78, 108, 79]])
|
|
faces = np.concatenate([faces, faces_new], axis=0)
|
|
vertex_colors = np.array([(*mesh_base_color, 1.0)] * vertices.shape[0])
|
|
# print(vertices.shape, self.faces.shape)
|
|
mesh = trimesh.Trimesh(vertices.copy(), faces.copy(), vertex_colors=vertex_colors)
|
|
# mesh = trimesh.Trimesh(vertices.copy(), self.faces.copy())
|
|
|
|
rot = trimesh.transformations.rotation_matrix(
|
|
np.radians(rot_angle), rot_axis)
|
|
mesh.apply_transform(rot)
|
|
|
|
rot = trimesh.transformations.rotation_matrix(
|
|
np.radians(180), [1, 0, 0])
|
|
mesh.apply_transform(rot)
|
|
return mesh
|