Spaces:
Sleeping
Sleeping
| import os | |
| import torch | |
| import trimesh | |
| import xatlas | |
| import numpy as np | |
| from sklearn.decomposition import PCA | |
| from torchvision import transforms | |
| from tqdm import tqdm | |
| from pytorch3d.io import ( | |
| load_obj, | |
| load_objs_as_meshes | |
| ) | |
| def compute_principle_directions(model_path, num_points=20000): | |
| mesh = trimesh.load_mesh(model_path, force="mesh") | |
| pc, _ = trimesh.sample.sample_surface_even(mesh, num_points) | |
| pc -= np.mean(pc, axis=0, keepdims=True) | |
| principle_directions = PCA(n_components=3).fit(pc).components_ | |
| return principle_directions | |
| def init_mesh(input_path, cache_path, device): | |
| print("=> parameterizing target mesh...") | |
| mesh = trimesh.load_mesh(input_path, force='mesh') | |
| try: | |
| vertices, faces = mesh.vertices, mesh.faces | |
| except AttributeError: | |
| print("multiple materials in {} are not supported".format(input_path)) | |
| exit() | |
| vmapping, indices, uvs = xatlas.parametrize(vertices, faces) | |
| xatlas.export(str(cache_path), vertices[vmapping], indices, uvs) | |
| print("=> loading target mesh...") | |
| # principle_directions = compute_principle_directions(cache_path) | |
| principle_directions = None | |
| _, faces, aux = load_obj(cache_path, device=device) | |
| mesh = load_objs_as_meshes([cache_path], device=device) | |
| num_verts = mesh.verts_packed().shape[0] | |
| # make sure mesh center is at origin | |
| bbox = mesh.get_bounding_boxes() | |
| mesh_center = bbox.mean(dim=2).repeat(num_verts, 1) | |
| mesh = apply_offsets_to_mesh(mesh, -mesh_center) | |
| # make sure mesh size is normalized | |
| box_size = bbox[..., 1] - bbox[..., 0] | |
| box_max = box_size.max(dim=1, keepdim=True)[0].repeat(num_verts, 3) | |
| mesh = apply_scale_to_mesh(mesh, 1 / box_max) | |
| return mesh, mesh.verts_packed(), faces, aux, principle_directions, mesh_center, box_max | |
| def apply_offsets_to_mesh(mesh, offsets): | |
| new_mesh = mesh.offset_verts(offsets) | |
| return new_mesh | |
| def apply_scale_to_mesh(mesh, scale): | |
| new_mesh = mesh.scale_verts(scale) | |
| return new_mesh | |
| def adjust_uv_map(faces, aux, init_texture, uv_size): | |
| """ | |
| adjust UV map to be compatiable with multiple textures. | |
| UVs for different materials will be decomposed and placed horizontally | |
| +-----+-----+-----+-- | |
| | 1 | 2 | 3 | | |
| +-----+-----+-----+-- | |
| """ | |
| textures_ids = faces.textures_idx | |
| materials_idx = faces.materials_idx | |
| verts_uvs = aux.verts_uvs | |
| num_materials = torch.unique(materials_idx).shape[0] | |
| new_verts_uvs = verts_uvs.clone() | |
| for material_id in range(num_materials): | |
| # apply offsets to horizontal axis | |
| faces_ids = textures_ids[materials_idx == material_id].unique() | |
| new_verts_uvs[faces_ids, 0] += material_id | |
| new_verts_uvs[:, 0] /= num_materials | |
| init_texture_tensor = transforms.ToTensor()(init_texture) | |
| init_texture_tensor = torch.cat([init_texture_tensor for _ in range(num_materials)], dim=-1) | |
| init_texture = transforms.ToPILImage()(init_texture_tensor).resize((uv_size, uv_size)) | |
| return new_verts_uvs, init_texture | |
| def update_face_angles(mesh, cameras, fragments): | |
| def get_angle(x, y): | |
| x = torch.nn.functional.normalize(x) | |
| y = torch.nn.functional.normalize(y) | |
| inner_product = (x * y).sum(dim=1) | |
| x_norm = x.pow(2).sum(dim=1).pow(0.5) | |
| y_norm = y.pow(2).sum(dim=1).pow(0.5) | |
| cos = inner_product / (x_norm * y_norm) | |
| angle = torch.acos(cos) | |
| angle = angle * 180 / 3.14159 | |
| return angle | |
| # face normals | |
| face_normals = mesh.faces_normals_padded()[0] | |
| # view vector (object center -> camera center) | |
| camera_center = cameras.get_camera_center() | |
| face_angles = get_angle( | |
| face_normals, | |
| camera_center.repeat(face_normals.shape[0], 1) | |
| ) # (F) | |
| face_angles_rev = get_angle( | |
| face_normals, | |
| -camera_center.repeat(face_normals.shape[0], 1) | |
| ) # (F) | |
| face_angles = torch.minimum(face_angles, face_angles_rev) | |
| # Indices of unique visible faces | |
| visible_map = fragments.pix_to_face.unique() # (num_visible_faces) | |
| invisible_mask = torch.ones_like(face_angles) | |
| invisible_mask[visible_map] = 0 | |
| face_angles[invisible_mask == 1] = 10000. # angles of invisible faces are ignored | |
| return face_angles | |