adapt ae_body to cape

losses/losses.py

@@ -24,14 +24,14 @@ F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[to_keep
 F_norm_orig = torch.tensor(np.tile(np.load("smpl/smpl_faces_orig.npy")[:6890,:][None,:,:], [4,1,1]), device='cuda')
 
 
-def _nearest_neighbour(V, B):
+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]
     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, k=1, batch_x=batch_x, batch_y=batch_y)[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)
     return idx
 
@@ -64,18 +64,28 @@ 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]
+edges_cape = torch.tensor(np.stack(np.nonzero(laplacian_loss_orig)).transpose(), device='cuda')
 
-def get_edge_mask(masks: torch.Tensor):
+def get_edge_mask(edges: torch.Tensor, masks: torch.Tensor):
     if edges.get_device() != masks.get_device():
         edges.to(masks.get_device())
     return torch.min(masks[:,edges], dim=-1)[0]
 
 def edge_loss(x, y, mask):
-    e_mask = get_edge_mask(mask)
-    x_e = x[:,edges[:,0],:] - x[:,edges[:,1],:]
+    if x.shape[1] == 14475:
+        edges_local = edges
+    elif x.shape[1] == 3638:
+        edges_local = edges_cape
+    else:
+        raise ValueError("Invalid number of vertices")
+    e_mask = get_edge_mask(edges_local, mask)
+    x_e = x[:,edges_local[:,0],:] - x[:,edges_local[:,1],:]
     x_e = torch.sum(x_e ** 2, -1)
     x_e = torch.sqrt(x_e + 1e-9)
-    y_e = y[:,edges[:,0],:] - y[:,edges[:,1],:]
+    y_e = y[:,edges_local[:,0],:] - y[:,edges_local[:,1],:]
     y_e = torch.sum(y_e ** 2, -1)
     y_e = torch.sqrt(y_e + 1e-9)
     d_e = y_e - x_e
@@ -173,8 +183,8 @@ def log_normal_pdf(sample, mean, logvar, raxis=1):
 
 
 def kl_loss(mean, logvar):
-    kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar)
-    kl_loss = tf.reduce_mean(kl_loss)
+    kl_loss = 1 + logvar - torch.square(mean) - torch.exp(logvar)
+    kl_loss = torch.mean(kl_loss)
     kl_loss *= -0.5
     return kl_loss