Spaces:
Sleeping
Sleeping
File size: 4,347 Bytes
b5042f1 |
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 |
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
@torch.no_grad()
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
|