import tensorflow as tf import numpy as np import tensorflow_graphics as tfg import scipy as sp import scipy.io as sio from scipy.spatial import cKDTree 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 """ F_coll = np.load("custom_utils/smpl_faces.npy") F_coll_skirt = np.load("custom_utils/smpl_faces_skirt.npy") F_coll_orig = np.load("custom_utils/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]) def _nearest_neighbour(V, B): tree = cKDTree(B.numpy()) return tree.query(V.numpy(), workers=-1)[1] @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): B_vn = vertex_normals(B, tf.tile(F_coll[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) thresholded = tf.minimum(dot - thr, 0) * mask loss = tf.reduce_sum(thresholded ** 2, axis=-1) ratio = tf.reduce_sum(tf.cast(thresholded < 0, dtype=tf.float32)) / tf.reduce_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 """ laplacian_loss = sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo() edges = np.stack(np.nonzero(laplacian_loss.todense())).transpose() laplacian_loss_skirt = sio.loadmat("custom_utils/ssmpl_skirt0.mat")['L'].tocoo() edges_skirt = np.stack(np.nonzero(laplacian_loss.todense())).transpose() def get_edge_mask(masks): return tf.reduce_min(tf.gather(masks, edges, axis=1), axis=-1) def edge_loss(x, y, mask): e_mask = get_edge_mask(mask) x_e = tf.gather(x, edges[:,0], axis=1) - tf.gather(x, edges[:,1], axis=1) x_e = tf.reduce_sum(x_e ** 2, -1) x_e = tf.sqrt(x_e + 1e-9) y_e = tf.gather(y, edges[:,0], axis=1) - tf.gather(y, edges[:,1], axis=1) y_e = tf.reduce_sum(y_e ** 2, -1) y_e = tf.sqrt(y_e + 1e-9) d_e = y_e - x_e d_e = d_e * e_mask loss = tf.reduce_sum(d_e ** 2, axis=-1)# / tf.reduce_sum(e_mask, axis=-1) return loss def get_edge_mask_skirt(masks): return tf.reduce_min(tf.gather(masks, edges_skirt, axis=1), axis=-1) def edge_loss_skirt(x, y, mask): e_mask = get_edge_mask_skirt(mask) x_e = tf.gather(x, edges_skirt[:,0], axis=1) - tf.gather(x, edges_skirt[:,1], axis=1) x_e = tf.reduce_sum(x_e ** 2, -1) x_e = tf.sqrt(x_e + 1e-9) y_e = tf.gather(y, edges_skirt[:,0], axis=1) - tf.gather(y, edges_skirt[:,1], axis=1) y_e = tf.reduce_sum(y_e ** 2, -1) y_e = tf.sqrt(y_e + 1e-9) d_e = y_e - x_e d_e = d_e * e_mask loss = tf.reduce_sum(d_e ** 2, axis=-1)# / tf.reduce_sum(e_mask, axis=-1) 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') 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 - tf.square(mean) - tf.exp(logvar) kl_loss = tf.reduce_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 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)) # vgg = tf.keras.applications.vgg16.VGG16(include_top=False, input_shape=(512,256,3)) # vgg_block_1 = tf.keras.Sequential(vgg.layers[:4]) # vgg_block_2 = tf.keras.Sequential(vgg.layers[4:7]) # vgg_block_3 = tf.keras.Sequential(vgg.layers[7:11]) # def concept_loss(pred, gt): # pred = tf.keras.applications.vgg16.preprocess_input((pred + 1) * 128) # gt = tf.keras.applications.vgg16.preprocess_input((gt + 1) * 128) # pred_features_1 = vgg_block_1(pred, training=False) # pred_features_2 = vgg_block_2(pred_features_1, training=False) # pred_features_3 = vgg_block_3(pred_features_2, training=False) # gt_features_1 = vgg_block_1(gt, training=False) # gt_features_2 = vgg_block_2(gt_features_1, training=False) # gt_features_3 = vgg_block_3(gt_features_2, training=False) # l1_1 = tf.reduce_mean(tf.math.abs(pred_features_1 - gt_features_1)) # l1_2 = tf.reduce_mean(tf.math.abs(pred_features_2 - gt_features_2)) # l1_3 = tf.reduce_mean(tf.math.abs(pred_features_3 - gt_features_3)) # return l1_1 + l1_2 + l1_3 if __name__ == '__main__': V = tf.ones((4, 512, 256, 3)) B = tf.ones((4, 512, 256, 3)) m = tf.ones((4, 14475)) loss_1 = concept_loss(V, B) print()