152 lines
4.6 KiB
Python
152 lines
4.6 KiB
Python
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() |