diff --git a/losses/losses.py b/losses/losses.py index 18c948e..2cc7bcb 100644 --- a/losses/losses.py +++ b/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") -F_coll_skirt = np.load("smpl/smpl_faces_skirt.npy") -F_coll_orig = np.load("smpl/smpl_faces_orig.npy") -F_norm = np.tile(F_coll[np.newaxis,:,:], [4,1,1]) -F_norm_skirt = np.tile(F_coll_skirt[np.newaxis,:,:], [4,1,1]) -F_norm_orig = np.tile(F_coll_orig[np.newaxis,:,:], [4,1,1]) +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') +F_coll_orig = torch.tensor(np.load("smpl/smpl_faces_orig.npy")[:6890,:], device='cuda') +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') + def _nearest_neighbour(V, B): - tree = cKDTree(B.numpy()) - return tree.query(V.numpy(), workers=-1, distance_upper_bound=0.2)[1] + # 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 = (idx - batch_x * 14475).reshape(batch_size, 14475) + return idx + # @tf.function -# def _tf_nn_parallel(V, B): -# 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]) - -# @tf.function -def collision_loss(V, B, mask, thr=.004): - mesh = Meshes(B, torch.tile(F_coll[None], [B.shape[0], 1, 1])) +def collision_loss(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: float = .004): + mesh = Meshes(B, F_norm) 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])]) - 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) + idx = _nearest_neighbour(V, B) # compute loss - D = V - B[idx] - b_vn = B_vn[idx] + 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, 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) - gt_norm = vertex_normals(gt, f_norm) - return tf.reduce_mean(tf.reduce_mean((tf.reduce_sum(tf.math.abs(pred_norm - gt_norm), axis=-1) * mask), axis=-1)) + mesh = Meshes(pred, f_norm) + pred_norm = torch.stack(mesh.verts_normals_list()) + 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))