implementtorch collision loss

losses/losses.py

@@ -7,6 +7,7 @@ import scipy.io as sio
 
 from pytorch3d.structures import Meshes
 from scipy.spatial import cKDTree
+from torch_cluster import knn
 # from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals
 # from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
 
@@ -14,60 +15,45 @@ from scipy.spatial import cKDTree
 
 # from hbertiche/PBNS
 """ COLLISION FUNCTIONS START """
-F_coll = np.load("smpl/smpl_faces.npy")
+to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
-F_coll_skirt = np.load("smpl/smpl_faces_skirt.npy")
+F_coll = torch.tensor(np.load("smpl/smpl_faces.npy")[to_keep,:], device='cuda')
-F_coll_orig = np.load("smpl/smpl_faces_orig.npy")
+F_coll_skirt = torch.tensor(np.load("smpl/smpl_faces_skirt.npy")[to_keep,:], device='cuda')
-F_norm = np.tile(F_coll[np.newaxis,:,:], [4,1,1])
+F_coll_orig = torch.tensor(np.load("smpl/smpl_faces_orig.npy")[:6890,:], device='cuda')
-F_norm_skirt = np.tile(F_coll_skirt[np.newaxis,:,:], [4,1,1])
+F_norm = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
-F_norm_orig = np.tile(F_coll_orig[np.newaxis,:,:], [4,1,1])
+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')
+
 
 def _nearest_neighbour(V, B):
-	tree = cKDTree(B.numpy())
+    # 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
-	return tree.query(V.numpy(), workers=-1, distance_upper_bound=0.2)[1]
+    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 = (idx - batch_x * 14475).reshape(batch_size, 14475)
+    return idx
+
 
 # @tf.function
-# def _tf_nn_parallel(V, B):
+def collision_loss(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: float = .004):
-# 	return tf.map_fn(fn=lambda elem: tf.py_function(func=_nearest_neighbour, inp=elem, Tout=tf.int32), elems=(V,B), fn_output_signature=tf.int32, parallel_iterations=V.shape[0])
+    mesh = Meshes(B, F_norm)
-
-# @tf.function
-def collision_loss(V, B, mask, thr=.004):
-    mesh = Meshes(B, torch.tile(F_coll[None], [B.shape[0], 1, 1]))
     B_vn = torch.stack(mesh.verts_normals_list())
    
     # compute correspondences
-    idx = torch.stack([_nearest_neighbour(V[i], B[i]) for i in range(V.shape[0])])
+    idx = _nearest_neighbour(V, B)
-    N, M = V.shape[:2]
-    # _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M])
-    _idx = torch.ones((N, M), dtype=torch.int32, device=V.get_device()) * np.array(list(range(N)), np.int32)[:,None]
-    idx = torch.stack((_idx, idx), axis=-1)
     # compute loss
-    D = V - B[idx]
+    D = V - torch.stack([B[i,idx[i],:] for i in range(V.shape[0])])
-    b_vn = B_vn[idx]
+    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, 0)  * mask
+    thresholded = torch.minimum(dot - thr, torch.zeros_like(dot))  * mask
     loss = torch.sum(thresholded ** 2, axis=-1)
     ratio = torch.sum((thresholded < 0).to(torch.float32)) / torch.sum(mask)
 
     return loss, ratio
 
 
-# def collision_ratio(V, B, mask, thr=.004):
-    # B_vn = vertex_normals(B, tf.tile(F[None], [B.shape[0], 1, 1]))
-   
-    # # compute correspondences
-    # idx = _tf_nn_parallel(V, B)
-    # N, M = V.shape[:2]
-    # # _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M])
-    # _idx = tf.ones((N, M), dtype=tf.int32) * np.array(list(range(N)), np.int32)[:,None]
-    # idx = tf.stack((_idx, idx), axis=-1)
-    # # compute loss
-    # D = V - tf.gather_nd(B, idx)
-    # b_vn = tf.gather_nd(B_vn, idx)
-    # dot = tf.einsum('abc,abc->ab', D, b_vn)
-    # loss = tf.minimum(dot - thr, 0)  * mask # / tf.reduce_sum(mask, axis=-1)
-    # ratio = tf.reduce_sum(tf.cast(loss < 0, dtype=tf.float32)) / tf.reduce_sum(mask)
-
-    # return ratio
 
 """ COLLISION FUNCTIONS END """
 
@@ -116,41 +102,6 @@ def edge_loss_skirt(x, y, mask):
     return loss
 
 
-# class PairwiseDistance(tf.keras.layers.Layer):
-#     def __init__(self, **kwargs):
-#         super(PairwiseDistance, self).__init__(**kwargs)
-
-#     def call(self, A, B):
-#         rA = tf.reduce_sum(tf.square(A), axis=-1)
-#         rB = tf.reduce_sum(tf.square(B), axis=-1)
-#         transpose_axes = [0, 2, 1] 
-#         distances = - 2*tf.matmul(A, tf.transpose(B, transpose_axes)) + rA[:, :, tf.newaxis] + rB[:, tf.newaxis, :]
-#         return distances
-
-
-# class NearestNeighbour(tf.keras.layers.Layer):
-#     def __init__(self, **kwargs):
-#         super(NearestNeighbour, self).__init__(**kwargs)
-
-#     def call(self, A, B):
-#         distances = PairwiseDistance(dtype=self.dtype)(A, B)
-#         nearest_neighbour = tf.argmin(distances, axis=-1)
-#         return tf.cast(nearest_neighbour, dtype=tf.int32)
-
-#
-# @tf.function
-def collision_penalty(va, vb, mask, eps=1e-3):
-    nb = vertex_normals(vb, tf.tile(F[None], [vb.shape[0], 1, 1]))
-    batch_size = tf.cast(tf.shape(va)[0], va.dtype)
-
-    closest_vertices = NearestNeighbour(dtype=va.dtype)(va, vb)
-    vb = tf.gather(vb, closest_vertices, batch_dims=1)
-    nb = tf.gather(nb, closest_vertices, batch_dims=1)
-
-    distance = tf.reduce_sum(nb*(va - vb), axis=-1) 
-    interpenetration = tf.maximum(eps - distance * mask, 0)
-
-    return tf.reduce_sum(interpenetration**3) / batch_size
 
 whole_garment_uv_mask = np.load('custom_utils/mask_256x512.npy')
 
@@ -239,6 +190,9 @@ def normal_loss(pred, gt, mask, is_skirt=False, is_smpl_orig=False):
     else:
         raise NotImplementedError
 
-    pred_norm = vertex_normals(pred, f_norm)
+    mesh = Meshes(pred, f_norm)
-    gt_norm = vertex_normals(gt, f_norm)
+    pred_norm = torch.stack(mesh.verts_normals_list())
-    return tf.reduce_mean(tf.reduce_mean((tf.reduce_sum(tf.math.abs(pred_norm - gt_norm), axis=-1) * mask), axis=-1))
+    mesh = Meshes(gt, f_norm)
+    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))