280 lines
9.7 KiB
Python
280 lines
9.7 KiB
Python
import numpy as np
|
|
import tensorflow_graphics as tfg
|
|
import tensorflow as tf
|
|
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 custom_utils.custom_utils import maskFacesBatch, _maskFaces
|
|
from custom_utils.graph_utils import sparse_to_tensor
|
|
|
|
|
|
def d_logistic(real_images, generator, discriminator, z_dim, conditions):
|
|
batch_size = tf.shape(real_images)[0]
|
|
z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
|
|
|
|
# forward pass
|
|
fake_images = generator([z, conditions], training=True)
|
|
real_scores = discriminator(real_images, training=True)
|
|
fake_scores = discriminator(fake_images, training=True)
|
|
|
|
# gan loss
|
|
d_loss = tf.math.softplus(fake_scores)
|
|
d_loss += tf.math.softplus(-real_scores)
|
|
|
|
|
|
return d_loss
|
|
|
|
|
|
def d_logistic_r1_reg(real_images, generator, discriminator, z_dim, conditions):
|
|
batch_size = tf.shape(real_images)[0]
|
|
z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
|
|
|
|
# forward pass
|
|
fake_images = generator([z, conditions], training=True)
|
|
real_scores = discriminator(real_images, training=True)
|
|
fake_scores = discriminator(fake_images, training=True)
|
|
|
|
# gan loss
|
|
d_loss = tf.math.softplus(fake_scores)
|
|
d_loss = d_loss + tf.math.softplus(-real_scores)
|
|
|
|
# gradient penalty
|
|
with tf.GradientTape() as r1_tape:
|
|
r1_tape.watch([real_images, conditions])
|
|
real_scores = discriminator(real_images, training=True)
|
|
real_loss = tf.reduce_sum(real_scores)
|
|
|
|
real_grads = r1_tape.gradient(real_loss, real_images)
|
|
r1_penalty = tf.reduce_sum(tf.math.square(real_grads), axis=[1, 2, 3])
|
|
r1_penalty = tf.expand_dims(r1_penalty, axis=1)
|
|
return d_loss, r1_penalty
|
|
|
|
|
|
def g_logistic_non_saturating(real_images, generator, discriminator, z_dim, conditions):
|
|
batch_size = tf.shape(real_images)[0]
|
|
z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
|
|
|
|
# forward pass
|
|
fake_images = generator([z, conditions], training=True)
|
|
fake_scores = discriminator(fake_images, training=True)
|
|
|
|
# gan loss
|
|
g_loss = tf.math.softplus(-fake_scores)
|
|
return fake_images, g_loss
|
|
|
|
|
|
def g_logistic_ns_pathreg(real_images, generator, discriminator, z_dim,
|
|
pl_mean, pl_minibatch_shrink, pl_denorm, pl_decay,
|
|
conditions):
|
|
batch_size = tf.shape(real_images)[0]
|
|
z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
|
|
|
|
pl_minibatch = tf.maximum(1, tf.math.floordiv(batch_size, pl_minibatch_shrink))
|
|
pl_z = tf.random.normal(shape=[pl_minibatch, z_dim], dtype=tf.float32)
|
|
pl_conditions = { key: conditions[key][:pl_minibatch] for key in conditions.keys() }
|
|
# for key in conditions.keys():
|
|
# pl_conditions[key] = conditions[key][:pl_minibatch]
|
|
# pl_conditions = (conditions[0][:pl_minibatch], conditions[1][:pl_minibatch], conditions[2][:pl_minibatch], conditions[3][:pl_minibatch], conditions[4][:pl_minibatch])
|
|
|
|
# forward pass
|
|
fake_images = generator([z, conditions], ret_w_broadcasted=False, training=True)
|
|
fake_scores = discriminator(fake_images, training=True)
|
|
g_loss = tf.math.softplus(-fake_scores)
|
|
|
|
# Evaluate the regularization term using a smaller minibatch to conserve memory.
|
|
with tf.GradientTape() as pl_tape:
|
|
pl_tape.watch([pl_z, pl_conditions])
|
|
pl_fake_images, pl_w_broadcasted = generator([pl_z, pl_conditions], ret_w_broadcasted=True, training=True)
|
|
|
|
pl_noise = tf.random.normal(tf.shape(pl_fake_images)) * pl_denorm
|
|
pl_noise_applied = tf.reduce_sum(pl_fake_images * pl_noise)
|
|
|
|
pl_grads = pl_tape.gradient(pl_noise_applied, pl_w_broadcasted)
|
|
pl_lengths = tf.math.sqrt(tf.reduce_mean(tf.reduce_sum(tf.math.square(pl_grads), axis=2), axis=1))
|
|
|
|
# Track exponential moving average of |J*y|.
|
|
pl_mean_val = pl_mean + pl_decay * (tf.reduce_mean(pl_lengths) - pl_mean)
|
|
pl_mean.assign(pl_mean_val)
|
|
|
|
# Calculate (|J*y|-a)^2.
|
|
pl_penalty = tf.square(pl_lengths - pl_mean)
|
|
return fake_images, g_loss, pl_penalty
|
|
|
|
|
|
# from hbertiche/PBNS
|
|
""" COLLISION FUNCTIONS START """
|
|
F = np.load("custom_utils/smpl_faces.npy")
|
|
|
|
|
|
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[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.reduce_sum((tf.minimum(dot - thr, 0) * mask) ** 2, axis=-1) # / tf.reduce_sum(mask, axis=-1)
|
|
|
|
return loss
|
|
|
|
""" COLLISION FUNCTIONS END """
|
|
|
|
|
|
laplacian_loss = sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo()
|
|
edges = 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
|
|
|
|
|
|
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
|
|
|
|
|
|
if __name__ == '__main__':
|
|
V = tf.ones((4, 14475, 3))
|
|
B = tf.zeros((4, 14475, 3))
|
|
m = tf.ones((4, 14475))
|
|
loss_1 = collision_loss(V, B, m, thr=.004)
|
|
print()
|