From aab1953e00c69cbac3cb2311361f71b0bad1a05c Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sat, 29 Oct 2022 20:11:07 +0000 Subject: [PATCH] add fix_collision --- postprocessing.py | 152 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 152 insertions(+) create mode 100644 postprocessing.py diff --git a/postprocessing.py b/postprocessing.py new file mode 100644 index 0000000..50f4093 --- /dev/null +++ b/postprocessing.py @@ -0,0 +1,152 @@ +import os + +# os.environ["CUDA_VISIBLE_DEVICES"] = "3" +import tensorflow as tf + +# gpus = tf.config.experimental.list_physical_devices('GPU') +# if gpus: +# try: +# # Currently, memory growth needs to be the same across GPUs +# for gpu in gpus: +# tf.config.experimental.set_memory_growth(gpu, True) +# logical_gpus = tf.config.experimental.list_logical_devices('GPU') +# print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs") +# except RuntimeError as e: +# # Memory growth must be set before GPUs have been initialized +# print(e) + +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.supersmpl_io import readPC2 +from custom_utils.smpl import SMPL +from custom_utils.custom_utils import _maskFaces + +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 = 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() \ No newline at end of file