save all

losses/losses.py

@@ -32,7 +32,7 @@ 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_x = torch.arange(batch_size, device=V.device).repeat_interleave(V.shape[1])
     V = V.reshape(-1, V.shape[-1])
     B = B.reshape(-1, B.shape[-1])
     idx = knn(B.to(V.dtype), V, k=1, batch_x=batch_x, batch_y=batch_x)[1]
@@ -58,7 +58,9 @@ def collision_loss(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: fl
         F_norm_local = F_norm
     elif B.shape[1] == 3638:
         F_norm_local = F_norm_orig
-    mesh = Meshes(B, F_norm_local)
+    if V.device != F_norm_local.device:
+        F_norm_local = F_norm_local.to(V.device)
+    mesh = Meshes(B, F_norm_local[:B.shape[0],:,:])
     B_vn = torch.stack(mesh.verts_normals_list())
    
     # compute correspondences
@@ -74,6 +76,54 @@ def collision_loss(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: fl
     return loss, ratio
 
 
+def _nearest_neighbour_error(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.device).repeat_interleave(B.shape[1])
+    batch_y = torch.arange(batch_size, device=V.device).repeat_interleave(V.shape[1])
+    V = V.reshape(-1, V.shape[-1])
+    B = B.reshape(-1, B.shape[-1])
+    idx = knn(B.to(V.dtype), V, k=1, batch_x=batch_x, batch_y=batch_y)[1]
+    idx = (idx - batch_y * vertices_count).reshape(batch_size, vertices_count)
+    return idx
+
+# def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor):
+#     # create an empty tensor to store the indices of the nearest neighbors
+#     batch_size = V.shape[0]
+#     vertices_count = V.shape[1]
+#     idx = torch.zeros((batch_size, vertices_count), dtype=torch.long, device=V.device)
+
+#     # loop over the batch dimension and find the nearest neighbor for each point in B
+#     for i in range(batch_size):
+#         tree = KDTree(B[i].cpu().numpy())
+#         _, nearest = tree.query(V[i].cpu().numpy(), k=1)
+#         idx[i] = torch.from_numpy(nearest)
+
+#     return idx
+
+def collision_error(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: float = 0.0005):
+    if B.shape[1] == 14475:
+        F_norm_local = F_norm
+    elif B.shape[1] == 3638:
+        F_norm_local = F_norm_orig
+    if V.device != F_norm_local.device:
+        F_norm_local = F_norm_local.to(V.device)
+    mesh = Meshes(B, F_norm_local[:B.shape[0],:,:])
+    B_vn = torch.stack(mesh.verts_normals_list())
+   
+    # compute correspondences
+    idx = _nearest_neighbour_error(V, B)
+    # compute loss
+    D = V - torch.stack([B[i,idx[i],:] for i in range(V.shape[0])])
+    b_vn = torch.stack([B_vn[i,idx[i],:] for i in range(V.shape[0])])
+    dot = torch.einsum('abc,abc->ab', D, b_vn)
+    thresholded = torch.minimum(dot - thr, torch.zeros_like(dot))
+    loss = torch.sum(thresholded ** 2, axis=-1)
+    ratio = torch.sum((thresholded < 0).to(torch.float32)) / V.shape[1]
+
+    return loss, ratio
+
 
 """ COLLISION FUNCTIONS END """
 
@@ -218,9 +268,9 @@ def normal_loss(pred, gt, mask, is_skirt=False, is_smpl_orig=False):
     else:
         raise NotImplementedError
 
-    mesh = Meshes(pred, f_norm)
+    mesh = Meshes(pred, f_norm[:pred.shape[0], :, :])
     pred_norm = torch.stack(mesh.verts_normals_list())
-    mesh = Meshes(gt, f_norm)
+    mesh = Meshes(gt, f_norm[:pred.shape[0], :, :])
     gt_norm = torch.stack(mesh.verts_normals_list())
 
     return torch.mean(torch.mean((torch.sum(torch.abs(pred_norm - gt_norm), axis=-1) * mask), axis=-1))