save progress
This commit is contained in:
+270
@@ -0,0 +1,270 @@
|
||||
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()
|
||||
@@ -0,0 +1,137 @@
|
||||
import os
|
||||
|
||||
# os.environ["CUDA_VISIBLE_DEVICES"] = "3"
|
||||
import tensorflow as tf
|
||||
|
||||
import sys
|
||||
import shutil
|
||||
import tqdm
|
||||
import random
|
||||
import math
|
||||
import numpy as np
|
||||
import scipy as sp
|
||||
import scipy.io as sio
|
||||
import plotly.graph_objects as go
|
||||
import plotly.express as ex
|
||||
|
||||
from custom_utils.custom_utils import _maskFaces
|
||||
|
||||
from scipy.spatial import cKDTree
|
||||
from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals
|
||||
|
||||
|
||||
# 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])
|
||||
|
||||
|
||||
def fix_collision(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
|
||||
B_offset_closest = tf.gather_nd(B, idx) + (b_vn / 2e2)
|
||||
output = tf.where(loss[:,:,tf.newaxis] < 0, B_offset_closest, V)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
""" COLLISION FUNCTIONS END """
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
F = np.load("custom_utils/smpl_faces.npy")
|
||||
posed = np.load('/home/workspace/local-data/train/mesh_posed/00016_10.npz')['mesh_posed'][np.newaxis,:,:]
|
||||
body = np.load('/home/workspace/local-data/train/mesh_body_posed/00016_10.npy')[np.newaxis,:,:]
|
||||
mask = np.load('/home/workspace/local-data/train/mesh_mask/00016.npy')
|
||||
posed_2 = np.load('/home/workspace/local-data/train/mesh_posed/00016_15.npz')['mesh_posed'][np.newaxis,:,:]
|
||||
|
||||
fixed = fix_collision(posed_2, body, mask).numpy()
|
||||
|
||||
x = posed_2[0, mask,:]
|
||||
y = body[0]
|
||||
_, faces = _maskFaces(np.zeros((14475,3), np.float32), F, mask)
|
||||
fig = go.Figure(data=[
|
||||
go.Mesh3d(
|
||||
x=x[:,0],
|
||||
y=x[:,1],
|
||||
z=x[:,2],
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2],
|
||||
),
|
||||
go.Mesh3d(
|
||||
x=y[:,0],
|
||||
y=y[:,1],
|
||||
z=y[:,2],
|
||||
i = F[:,0],
|
||||
j = F[:,1],
|
||||
k = F[:,2],
|
||||
),
|
||||
])
|
||||
|
||||
fig.update_layout(scene_aspectmode='manual',
|
||||
scene_aspectratio=dict(x=1, y=1, z=1))
|
||||
fig.update_layout(
|
||||
scene = dict(
|
||||
xaxis = dict(nticks=4, range=[-1,1],),
|
||||
yaxis = dict(nticks=4, range=[-1,1],),
|
||||
zaxis = dict(nticks=4, range=[-1,1],),)
|
||||
)
|
||||
|
||||
fig.show()
|
||||
|
||||
x = fixed[0, mask,:]
|
||||
y = body[0]
|
||||
_, faces = _maskFaces(np.zeros((14475,3), np.float32), F, mask)
|
||||
fig = go.Figure(data=[
|
||||
go.Mesh3d(
|
||||
x=x[:,0],
|
||||
y=x[:,1],
|
||||
z=x[:,2],
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2],
|
||||
),
|
||||
go.Mesh3d(
|
||||
x=y[:,0],
|
||||
y=y[:,1],
|
||||
z=y[:,2],
|
||||
i = F[:,0],
|
||||
j = F[:,1],
|
||||
k = F[:,2],
|
||||
),
|
||||
])
|
||||
|
||||
fig.update_layout(scene_aspectmode='manual',
|
||||
scene_aspectratio=dict(x=1, y=1, z=1))
|
||||
fig.update_layout(
|
||||
scene = dict(
|
||||
xaxis = dict(nticks=4, range=[-1,1],),
|
||||
yaxis = dict(nticks=4, range=[-1,1],),
|
||||
zaxis = dict(nticks=4, range=[-1,1],),)
|
||||
)
|
||||
|
||||
fig.show()
|
||||
|
||||
print()
|
||||
@@ -0,0 +1,27 @@
|
||||
import os
|
||||
import sys
|
||||
import numpy as np
|
||||
|
||||
|
||||
|
||||
class Accumulator():
|
||||
def __init__(self):
|
||||
self.dict = dict()
|
||||
|
||||
def update(self, dict):
|
||||
for key in dict.keys():
|
||||
if key in self.dict.keys():
|
||||
self.dict[key].append(dict[key])
|
||||
else:
|
||||
self.dict[key] = [dict[key]]
|
||||
|
||||
def get(self):
|
||||
result = dict()
|
||||
for key in self.dict.keys():
|
||||
result[key] = np.mean(self.dict[key])
|
||||
return result
|
||||
|
||||
def reset(self):
|
||||
self.dict = dict()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user