tohid4n/PartCrafter

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