luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

pyramid ready for cape

Browse Source
This commit is contained in:
luni
2023-05-22 12:44:43 +02:00
parent 6b81c0d1a6
commit 60844f8011
5 changed files with 321 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
losses/losses.py
+12 -5
View File
@@ -15,6 +15,9 @@ from torch_cluster import knn
# from hbertiche/PBNS
""" COLLISION FUNCTIONS START """
smpl_to_keep_orig = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
smpl_to_keep_orig = smpl_to_keep_orig[smpl_to_keep_orig < 6890]
to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
F_coll = torch.tensor(np.load("smpl/smpl_faces.npy")[to_keep,:], device='cuda')
F_coll_skirt = torch.tensor(np.load("smpl/smpl_faces_skirt.npy")[to_keep,:], device='cuda')
@@ -22,23 +25,29 @@ F_coll_orig = torch.tensor(np.load("smpl/smpl_faces_orig.npy")[:6890,:], device=
F_norm = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
F_norm_orig = torch.tensor(np.tile(np.load("smpl/smpl_faces_orig.npy")[:6890,:][None,:,:], [4,1,1]), device='cuda')
F_norm_orig = F_norm_orig[:,smpl_to_keep_orig,:]
def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor):
# flatten V and B along the batch dimension, create a batch_x and batch_y variable to restore the batch dimension after the nearest neighbour search using pytorch_geomtry knn function
batch_size = V.shape[0]
vertices_count = V.shape[1]
batch_x = torch.arange(batch_size, device=V.get_device()).repeat_interleave(V.shape[1])
batch_y = torch.arange(batch_size, device=B.get_device()).repeat_interleave(B.shape[1])
V = V.reshape(-1, V.shape[-1])
B = B.reshape(-1, B.shape[-1])
idx = knn(V, B.to(V.dtype), k=1, batch_x=batch_x, batch_y=batch_y)[1]
idx = (idx - batch_x * 14475).reshape(batch_size, 14475)
idx = (idx - batch_x * vertices_count).reshape(batch_size, vertices_count)
return idx
# @tf.function
def collision_loss(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: float = .004):
mesh = Meshes(B, F_norm)
if B.shape[1] == 14475:
F_norm_local = F_norm
elif B.shape[1] == 3638:
F_norm_local = F_norm_orig
mesh = Meshes(B, F_norm_local)
B_vn = torch.stack(mesh.verts_normals_list())
# compute correspondences
@@ -64,9 +73,7 @@ edges: torch.Tensor = torch.tensor(np.stack(np.nonzero(laplacian_loss.todense())
laplacian_loss_skirt = sio.loadmat("custom_utils/ssmpl_skirt0.mat")['L'].tocoo()
edges_skirt: torch.Tensor = torch.tensor(np.stack(np.nonzero(laplacian_loss.todense())).transpose(), device='cuda')
smpl_to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
smpl_to_keep = smpl_to_keep[smpl_to_keep < 6890]
laplacian_loss_orig = laplacian_loss.todense()[:6890,:6890][smpl_to_keep,:][:,smpl_to_keep]
laplacian_loss_orig = laplacian_loss.todense()[:6890,:6890][smpl_to_keep_orig,:][:,smpl_to_keep_orig]
edges_cape = torch.tensor(np.stack(np.nonzero(laplacian_loss_orig)).transpose(), device='cuda')
def get_edge_mask(edges: torch.Tensor, masks: torch.Tensor):