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))