implementtorch collision loss
This commit is contained in:
+31
-77
@@ -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))
|
||||
|
||||
Reference in New Issue
Block a user