Spaces:
Sleeping
Sleeping
"""Mesh data utilities.""" | |
import random | |
import networkx as nx | |
import numpy as np | |
# import pyrr | |
from six.moves import range | |
import trimesh | |
from scipy.spatial.transform import Rotation | |
def to_mesh(vertices, faces, transpose=True, post_process=False): | |
if transpose: | |
vertices = vertices[:, [1, 2, 0]] | |
if faces.min() == 1: | |
faces = (np.array(faces) - 1).tolist() | |
mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False) | |
if post_process: | |
mesh.merge_vertices() | |
mesh.update_faces(mesh.unique_faces()) | |
mesh.fix_normals() | |
return mesh | |
def center_vertices(vertices): | |
"""Translate the vertices so that bounding box is centered at zero.""" | |
vert_min = vertices.min(axis=0) | |
vert_max = vertices.max(axis=0) | |
vert_center = 0.5 * (vert_min + vert_max) | |
# vert_center = np.mean(vertices, axis=0) | |
return vertices - vert_center | |
def face_to_cycles(face): | |
"""Find cycles in face.""" | |
g = nx.Graph() | |
for v in range(len(face) - 1): | |
g.add_edge(face[v], face[v + 1]) | |
g.add_edge(face[-1], face[0]) | |
return list(nx.cycle_basis(g)) | |
def block_index(vertex, block_size=32): | |
return (vertex[2] // block_size, vertex[1] // block_size, vertex[0] // block_size) | |
def block_id(block_index, num_blocks=4): | |
return block_index[0] * num_blocks**2 + block_index[1] * num_blocks + block_index[2] | |
def normalize_vertices_scale(vertices, scale=0.95): | |
"""Scale the vertices so that the long axis of the bounding box is one.""" | |
vert_min = vertices.min(axis=0) | |
vert_max = vertices.max(axis=0) | |
extents = (vert_max - vert_min).max() | |
return 2.0 * scale * vertices / (extents + 1e-6) | |
def quantize_process_mesh(vertices, faces, quantization_bits=8, block_first_order=True, block_size=32, num_blocks=4): | |
"""Quantize vertices, remove resulting duplicates and reindex faces.""" | |
vertices = discretize(vertices, num_discrete=2**quantization_bits) | |
vertices, inv = np.unique(vertices, axis=0, return_inverse=True) | |
if block_first_order: | |
block_indices = np.array([block_index(v, block_size) for v in vertices]) | |
block_ids = np.array([block_id(b, num_blocks) for b in block_indices]) | |
sort_inds = np.lexsort((vertices[:, 0], vertices[:, 1], vertices[:, 2], block_ids)) | |
else: | |
# Sort vertices by z then y then x. | |
sort_inds = np.lexsort(vertices.T) | |
vertices = vertices[sort_inds] | |
faces = [np.argsort(sort_inds)[inv[f]] for f in faces] | |
sub_faces = [] | |
for f in faces: | |
cliques = face_to_cycles(f) | |
for c in cliques: | |
c_length = len(c) | |
if c_length > 2: | |
d = np.argmin(f) | |
sub_faces.append([f[(d + i) % c_length] for i in range(c_length)]) | |
faces = sub_faces | |
# Sort faces by lowest vertex indices. If two faces have the same lowest | |
# index then sort by next lowest and so on. | |
faces.sort(key=lambda f: tuple(sorted(f))) | |
num_verts = vertices.shape[0] | |
vert_connected = np.equal( | |
np.arange(num_verts)[:, None], np.hstack(faces)[None] | |
).any(axis=-1) | |
vertices = vertices[vert_connected] | |
# Re-index faces to re-ordered vertices. | |
vert_indices = np.arange(num_verts) - np.cumsum(1 - vert_connected.astype("int")) | |
faces = [vert_indices[f].tolist() for f in faces] | |
return vertices, faces | |
def process_mesh(vertices, faces, quantization_bits=8, augment=True, augment_dict=None): | |
"""Process mesh vertices and faces.""" | |
# Transpose so that z-axis is vertical. | |
vertices = vertices[:, [2, 0, 1]] | |
# Translate the vertices so that bounding box is centered at zero. | |
vertices = center_vertices(vertices) | |
if augment: | |
vertices = augment_mesh(vertices, **augment_dict) | |
# Scale the vertices so that the long diagonal of the bounding box is equal | |
# to one. | |
vertices = normalize_vertices_scale(vertices) | |
# Quantize and sort vertices, remove resulting duplicates, sort and reindex | |
# faces. | |
vertices, faces = quantize_process_mesh( | |
vertices, faces, quantization_bits=quantization_bits | |
) | |
vertices = undiscretize(vertices, num_discrete=2**quantization_bits) | |
# Discard degenerate meshes without faces. | |
return { | |
"vertices": vertices, | |
"faces": faces, | |
} | |
def load_process_mesh(mesh_obj_path, quantization_bits=8, augment=False, augment_dict=None): | |
"""Load obj file and process.""" | |
# Load mesh | |
mesh = trimesh.load(mesh_obj_path, force='mesh', process=False) | |
return process_mesh(mesh.vertices, mesh.faces, quantization_bits, augment=augment, augment_dict=augment_dict) | |
def augment_mesh(vertices, scale_min=0.95, scale_max=1.05, rotation=0., jitter_strength=0.): | |
'''scale vertices by a factor in [0.75, 1.25]''' | |
# vertices [nv, 3] | |
for i in range(3): | |
# Generate a random scale factor | |
scale = random.uniform(scale_min, scale_max) | |
# independently applied scaling across each axis of vertices | |
vertices[:, i] *= scale | |
if rotation != 0.: | |
axis = [random.uniform(-1, 1), random.uniform(-1, 1), random.uniform(-1, 1)] | |
radian = np.pi / 180 * rotation | |
rotation = Rotation.from_rotvec(radian * np.array(axis)) | |
vertices =rotation.apply(vertices) | |
if jitter_strength != 0.: | |
jitter_amount = np.random.uniform(-jitter_strength, jitter_strength) | |
vertices += jitter_amount | |
return vertices | |
def discretize( | |
t, | |
continuous_range = (-1, 1), | |
num_discrete: int = 128 | |
): | |
lo, hi = continuous_range | |
assert hi > lo | |
t = (t - lo) / (hi - lo) | |
t *= num_discrete | |
t -= 0.5 | |
return t.round().astype(np.int32).clip(min = 0, max = num_discrete - 1) | |
def undiscretize( | |
t, | |
continuous_range = (-1, 1), | |
num_discrete: int = 128 | |
): | |
lo, hi = continuous_range | |
assert hi > lo | |
t = t.astype(np.float32) | |
t += 0.5 | |
t /= num_discrete | |
return t * (hi - lo) + lo | |