227 lines
8.1 KiB
Python
227 lines
8.1 KiB
Python
import torch
|
|
|
|
import numpy as np
|
|
# import tensorflow_graphics as tfg
|
|
import scipy as sp
|
|
import scipy.io as sio
|
|
from scipy.spatial import KDTree
|
|
|
|
from pytorch3d.structures import Meshes
|
|
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
|
|
|
|
|
|
|
|
# from hbertiche/PBNS
|
|
""" COLLISION FUNCTIONS START """
|
|
smpl_to_keep_orig = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
|
|
smpl_to_keep_orig = smpl_to_keep_orig[smpl_to_keep_orig < 6890]
|
|
|
|
# to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
|
|
F_coll = torch.tensor(np.load("smpl/smpl_faces.npy"), device='cuda')
|
|
F_coll_skirt = torch.tensor(np.load("smpl/smpl_faces_skirt.npy"), 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")[None,:,:], [4,1,1]), device='cuda')
|
|
F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[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')
|
|
F_norm_orig = F_norm_orig[:,smpl_to_keep_orig,:]
|
|
|
|
|
|
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])
|
|
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]
|
|
idx = (idx - batch_x * 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_loss(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: float = .004):
|
|
if B.shape[1] == 14475:
|
|
F_norm_local = F_norm
|
|
elif B.shape[1] == 3638:
|
|
F_norm_local = F_norm_orig
|
|
mesh = Meshes(B, F_norm_local)
|
|
B_vn = torch.stack(mesh.verts_normals_list())
|
|
|
|
# compute correspondences
|
|
idx = _nearest_neighbour(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)) * mask
|
|
loss = torch.sum(thresholded ** 2, axis=-1)
|
|
ratio = torch.sum((thresholded < 0).to(torch.float32)) / torch.sum(mask)
|
|
|
|
return loss, ratio
|
|
|
|
|
|
|
|
""" COLLISION FUNCTIONS END """
|
|
|
|
|
|
laplacian_loss = sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo()
|
|
edges: torch.Tensor = torch.tensor(np.stack(np.nonzero(laplacian_loss.todense())).transpose(), device='cuda')
|
|
|
|
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')
|
|
|
|
laplacian_loss_orig = laplacian_loss.todense()[:6890,:6890][smpl_to_keep_orig,:][:,smpl_to_keep_orig]
|
|
edges_cape = torch.tensor(np.stack(np.nonzero(laplacian_loss_orig)).transpose(), device='cuda')
|
|
|
|
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):
|
|
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_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
|
|
d_e = d_e * e_mask
|
|
loss = torch.sum(d_e ** 2, dim=-1)# / tf.reduce_sum(e_mask, axis=-1)
|
|
return loss
|
|
|
|
def get_edge_mask_skirt(masks: torch.Tensor):
|
|
if edges.get_device() != masks.get_device():
|
|
edges.to(masks.get_device())
|
|
return torch.min(masks[:,edges_skirt], dim=-1)[0]
|
|
|
|
def edge_loss_skirt(x, y, mask):
|
|
e_mask = get_edge_mask_skirt(mask)
|
|
x_e = x[:,edges_skirt[:,0],:] - x[:,edges_skirt[:,1],:]
|
|
x_e = torch.sum(x_e ** 2, -1)
|
|
x_e = torch.sqrt(x_e + 1e-9)
|
|
y_e = y[:,edges_skirt[:,0],:] - y[:,edges_skirt[:,1],:]
|
|
y_e = torch.sum(y_e ** 2, -1)
|
|
y_e = torch.sqrt(y_e + 1e-9)
|
|
d_e = y_e - x_e
|
|
d_e = d_e * e_mask
|
|
loss = torch.sum(d_e ** 2, dim=-1)# / tf.reduce_sum(e_mask, axis=-1)
|
|
return loss
|
|
|
|
|
|
|
|
whole_garment_uv_mask = np.load('custom_utils/mask_256x512.npy')
|
|
|
|
def uv_mask_loss(pred, uv_mask):
|
|
inv_mask = (1 - uv_mask) * whole_garment_uv_mask
|
|
pred = tf.reduce_sum(pred, axis=1)
|
|
invalid = pred * inv_mask
|
|
invalid = invalid > 0.01
|
|
invalid_count = tf.math.count_nonzero(invalid, dtype=tf.float32)
|
|
total_count = tf.math.count_nonzero(inv_mask, dtype=tf.float32)
|
|
return invalid_count / total_count
|
|
|
|
|
|
|
|
|
|
def merge_verts(Vs, batchDim):
|
|
V = tf.concat(tf.unstack(Vs, num=batchDim), axis=0)
|
|
return V
|
|
|
|
|
|
def _bend_loss_single(x, F, L, mask):
|
|
verts = tf.boolean_mask(x, mask, axis=0)
|
|
VN = vertex_normals(verts, F)
|
|
bend = edge_convolution_template(
|
|
VN, L, sizes=None,
|
|
edge_function=lambda x,y: (x - y)**2,
|
|
reduction='weighted',
|
|
edge_function_kwargs={}
|
|
)
|
|
loss = tf.reduce_sum(bend)
|
|
return loss
|
|
|
|
def bend_loss(x, F, L, mask):
|
|
# x: predicted outfits
|
|
# F: faces
|
|
# L: laplacian
|
|
loss = tf.stack([_bend_loss_single(x[i], F[i], L[i], mask[i]) for i in range(x.shape[0])])
|
|
return loss
|
|
|
|
|
|
# def bend_loss(x, F, L, mask):
|
|
# # x: predicted outfits
|
|
# # F: faces
|
|
# # L: laplacian
|
|
# # x = x * mask[:,:,tf.newaxis]
|
|
# batch_size = x.shape[0]
|
|
# verts, faces = maskFacesBatch(x, F, mask)
|
|
# nr_valid_verts = [tf.reduce_sum(mask[i]) for i in range(batch_size)]
|
|
# x = tf.concat(verts, axis=0)
|
|
# F = tf.concat([faces[i] + tf.reduce_sum(nr_valid_verts[:i]) for i in range(1, batch_size)])
|
|
# F = tf.concat((faces[0], F), axis=0)
|
|
# VN = vertex_normals(x, F)
|
|
# bend = edge_convolution_template(
|
|
# VN, L, sizes=None,
|
|
# edge_function=lambda x,y: (x - y)**2,
|
|
# reduction='weighted',
|
|
# edge_function_kwargs={}
|
|
# )
|
|
# loss = tf.reduce_sum(bend, axis=[1])
|
|
# loss = tf.reshape(loss, (batch_size, 14475))
|
|
# loss = tf.reduce_mean(loss, axis=1)
|
|
# return loss
|
|
|
|
|
|
def log_normal_pdf(sample, mean, logvar, raxis=1):
|
|
log2pi = tf.math.log(2. * np.pi)
|
|
res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis)
|
|
return res
|
|
|
|
|
|
def kl_loss(mean, logvar):
|
|
kl_loss = 1 + logvar - torch.square(mean) - torch.exp(logvar)
|
|
kl_loss = torch.mean(kl_loss)
|
|
kl_loss *= -0.5
|
|
return kl_loss
|
|
|
|
|
|
def normal_loss(pred, gt, mask, is_skirt=False, is_smpl_orig=False):
|
|
f_norm = None
|
|
if not is_skirt and not is_smpl_orig:
|
|
f_norm = F_norm
|
|
elif is_skirt:
|
|
f_norm = F_norm_skirt
|
|
elif is_smpl_orig:
|
|
f_norm = F_norm_orig
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
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))
|