Spaces:
Running
on
Zero
Running
on
Zero
File size: 6,283 Bytes
bef5729 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 |
from src.utils.typing_utils import *
import os
import numpy as np
import trimesh
import torch
def normalize_mesh(
mesh: Union[trimesh.Trimesh, trimesh.Scene],
scale: float = 2.0,
):
# if not isinstance(mesh, trimesh.Trimesh) and not isinstance(mesh, trimesh.Scene):
# raise ValueError("Input mesh is not a trimesh.Trimesh or trimesh.Scene object.")
bbox = mesh.bounding_box
translation = -bbox.centroid
scale = scale / bbox.primitive.extents.max()
mesh.apply_translation(translation)
mesh.apply_scale(scale)
return mesh
def remove_overlapping_vertices(mesh: trimesh.Trimesh, reserve_material: bool = False):
if not isinstance(mesh, trimesh.Trimesh):
raise ValueError("Input mesh is not a trimesh.Trimesh object.")
vertices = mesh.vertices
faces = mesh.faces
unique_vertices, index_map, inverse_map = np.unique(
vertices, axis=0, return_index=True, return_inverse=True
)
clean_faces = inverse_map[faces]
clean_mesh = trimesh.Trimesh(vertices=unique_vertices, faces=clean_faces, process=True)
if reserve_material:
uv = mesh.visual.uv
material = mesh.visual.material
clean_uv = uv[index_map]
clean_visual = trimesh.visual.TextureVisuals(uv=clean_uv, material=material)
clean_mesh.visual = clean_visual
return clean_mesh
RGB = [
(82, 170, 220),
(215, 91, 78),
(45, 136, 117),
(247, 172, 83),
(124, 121, 121),
(127, 171, 209),
(243, 152, 101),
(145, 204, 192),
(150, 59, 121),
(181, 206, 78),
(189, 119, 149),
(199, 193, 222),
(200, 151, 54),
(236, 110, 102),
(238, 182, 212),
]
def get_colored_mesh_composition(
meshes: Union[List[trimesh.Trimesh], trimesh.Scene],
is_random: bool = True,
is_sorted: bool = False,
RGB: List[Tuple] = RGB
):
if isinstance(meshes, trimesh.Scene):
meshes = meshes.dump()
if is_sorted:
volumes = []
for mesh in meshes:
try:
volume = mesh.volume
except:
volume = 0.0
volumes.append(volume)
# sort by volume from large to small
meshes = [x for _, x in sorted(zip(volumes, meshes), key=lambda pair: pair[0], reverse=True)]
colored_scene = trimesh.Scene()
for idx, mesh in enumerate(meshes):
if is_random:
color = (np.random.rand(3) * 256).astype(int)
else:
color = np.array(RGB[idx % len(RGB)])
mesh.visual = trimesh.visual.ColorVisuals(
mesh=mesh,
vertex_colors=color,
)
colored_scene.add_geometry(mesh)
return colored_scene
def mesh_to_surface(
mesh: trimesh.Trimesh,
num_pc: int = 204800,
clip_to_num_vertices: bool = False,
return_dict: bool = False,
):
# if not isinstance(mesh, trimesh.Trimesh):
# raise ValueError("mesh must be a trimesh.Trimesh object")
if clip_to_num_vertices:
num_pc = min(num_pc, mesh.vertices.shape[0])
points, face_indices = mesh.sample(num_pc, return_index=True)
normals = mesh.face_normals[face_indices]
if return_dict:
return {
"surface_points": points,
"surface_normals": normals,
}
return points, normals
def scene_to_parts(
mesh: trimesh.Scene,
normalize: bool = True,
scale: float = 2.0,
num_part_pc: int = 204800,
clip_to_num_part_vertices: bool = False,
return_type: Literal["mesh", "point"] = "mesh",
) -> Union[List[trimesh.Geometry], List[Dict[str, np.ndarray]]]:
if not isinstance(mesh, trimesh.Scene):
raise ValueError("mesh must be a trimesh.Scene object")
if normalize:
mesh = normalize_mesh(mesh, scale=scale)
parts: List[trimesh.Geometry] = mesh.dump()
if return_type == "point":
datas: List[Dict[str, np.ndarray]] = []
for geom in parts:
data = mesh_to_surface(
geom,
num_pc=num_part_pc,
clip_to_num_vertices=clip_to_num_part_vertices,
return_dict=True,
)
datas.append(data)
return datas
elif return_type == "mesh":
return parts
else:
raise ValueError("return_type must be 'mesh' or 'point'")
def get_center(mesh: trimesh.Trimesh, method: Literal['mass', 'bbox']):
if method == 'mass':
return mesh.center_mass
elif method =='bbox':
return mesh.bounding_box.centroid
else:
raise ValueError('type must be mass or bbox')
def get_direction(vector: np.ndarray):
return vector / np.linalg.norm(vector)
def move_mesh_by_center(mesh: trimesh.Trimesh, scale: float, method: Literal['mass', 'bbox'] = 'mass'):
offset = scale - 1
center = get_center(mesh, method)
direction = get_direction(center)
translation = direction * offset
mesh = mesh.copy()
mesh.apply_translation(translation)
return mesh
def move_meshes_by_center(meshes: Union[List[trimesh.Trimesh], trimesh.Scene], scale: float):
if isinstance(meshes, trimesh.Scene):
meshes = meshes.dump()
moved_meshes = []
for mesh in meshes:
moved_mesh = move_mesh_by_center(mesh, scale)
moved_meshes.append(moved_mesh)
moved_meshes = trimesh.Scene(moved_meshes)
return moved_meshes
def get_series_splited_meshes(meshes: List[trimesh.Trimesh], scale: float, num_steps: int) -> List[trimesh.Scene]:
series_meshes = []
for i in range(num_steps):
temp_scale = 1 + (scale - 1) * i / (num_steps - 1)
temp_meshes = move_meshes_by_center(meshes, temp_scale)
series_meshes.append(temp_meshes)
return series_meshes
def load_surface(data, num_pc=204800):
surface = data["surface_points"] # Nx3
normal = data["surface_normals"] # Nx3
rng = np.random.default_rng()
ind = rng.choice(surface.shape[0], num_pc, replace=False)
surface = torch.FloatTensor(surface[ind])
normal = torch.FloatTensor(normal[ind])
surface = torch.cat([surface, normal], dim=-1)
return surface
def load_surfaces(surfaces, num_pc=204800):
surfaces = [load_surface(surface, num_pc) for surface in surfaces]
surfaces = torch.stack(surfaces, dim=0)
return surfaces |