223 lines
8.7 KiB
Python
223 lines
8.7 KiB
Python
import numpy as np
|
|
import tensorflow as tf
|
|
import plotly.graph_objects as go
|
|
from plotly.subplots import make_subplots
|
|
|
|
|
|
def m2uv_array(m2uv):
|
|
result = np.zeros((14475, 2), dtype=np.int32)
|
|
for k, v in m2uv.items():
|
|
if len(v) == 2:
|
|
result[k] = list(v)
|
|
elif len(v) == 1:
|
|
x = v.pop()
|
|
result[k] = [x, x]
|
|
else:
|
|
assert False
|
|
|
|
return result
|
|
|
|
|
|
class UV2Mesh():
|
|
def __init__(self, res=256, is_skirt=False):
|
|
super(UV2Mesh, self).__init__()
|
|
self.res = res
|
|
uv_map = np.load('custom_utils/stretched_uv_map_' + ('skirt' if is_skirt else '3') + '.npy')
|
|
self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32)
|
|
m2uv = np.load('custom_utils/m2uv' + ('_skirt' if is_skirt else '') + '.npy', allow_pickle=True).item()
|
|
m2uv = m2uv_array(m2uv)
|
|
self.uv_pixels = self.uv_map_discreet[m2uv]
|
|
# self.batch_size = batch_size
|
|
|
|
# self.x1 = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.zeros((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
|
|
# self.y1 = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
|
|
# self.z1 = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32) * 2), axis=-1)
|
|
# self.x2 = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.zeros((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
|
|
# self.y2 = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
|
|
# self.z2 = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32) * 2), axis=-1)
|
|
|
|
@tf.function
|
|
def uv_to_mesh(self, uv):
|
|
if uv.shape[1] == 3:
|
|
uv = tf.transpose(uv, [0,2,3,1])
|
|
if uv.shape[1] != 512:
|
|
uv = tf.image.resize(uv, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
|
|
|
batch_size = uv.shape[0]
|
|
|
|
x1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
|
|
y1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
|
|
z1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, 14475, 1), dtype=tf.int32) * 2), axis=-1)
|
|
x2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
|
|
y2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
|
|
z2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, 14475, 1), dtype=tf.int32) * 2), axis=-1)
|
|
|
|
|
|
x1 = tf.gather_nd(uv, x1_ind, batch_dims=1)
|
|
y1 = tf.gather_nd(uv, y1_ind, batch_dims=1)
|
|
z1 = tf.gather_nd(uv, z1_ind, batch_dims=1)
|
|
x2 = tf.gather_nd(uv, x2_ind, batch_dims=1)
|
|
y2 = tf.gather_nd(uv, y2_ind, batch_dims=1)
|
|
z2 = tf.gather_nd(uv, z2_ind, batch_dims=1)
|
|
|
|
p1 = tf.stack((x1, y1, z1), axis=-1)
|
|
p2 = tf.stack((x2, y2, z2), axis=-1)
|
|
|
|
mesh = (p1 + p2) / 2
|
|
return mesh
|
|
|
|
def __call__(self, uv):
|
|
return self.uv_to_mesh(uv)
|
|
|
|
|
|
class UV2MeshLabel():
|
|
def __init__(self, res=256):
|
|
super(UV2MeshLabel, self).__init__()
|
|
self.res = res
|
|
uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
|
|
self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32)
|
|
m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
|
|
m2uv = m2uv_array(m2uv)
|
|
self.uv_pixels = self.uv_map_discreet[m2uv]
|
|
|
|
# input = channels_first
|
|
# @tf.function
|
|
def uv_to_mesh(self, uv):
|
|
batch_size = uv.shape[0]
|
|
|
|
x1_ind = tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1])
|
|
x2_ind = tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1])
|
|
|
|
x1 = tf.gather_nd(uv, x1_ind, batch_dims=1)
|
|
x2 = tf.gather_nd(uv, x2_ind, batch_dims=1)
|
|
|
|
mesh = (x1 + x2) / 2
|
|
return mesh
|
|
|
|
def __call__(self, uv):
|
|
return self.uv_to_mesh(uv)
|
|
|
|
|
|
def _maskFaces(V, F, mask):
|
|
V_mask = mask
|
|
F_mask = np.take(V_mask, F.reshape(-1), axis=0)
|
|
F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
|
|
|
|
V_ind = np.zeros(V_mask.shape)
|
|
V_ind[V_mask] = np.arange(V_mask.sum())
|
|
F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
|
|
V_new = V[V_mask, :]
|
|
|
|
return V_new, F_new
|
|
|
|
|
|
# def _maskFaces_tf(V, F, mask):
|
|
# V_mask = mask
|
|
# F_mask = tf.gather(V_mask, F.reshape(-1), axis=0)
|
|
# F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
|
|
# F_mask = tf.reduce_sum(tf.reshape(F_mask, (-1, 3)), axis=1) == 3
|
|
|
|
# V_ind = tf.zeros(V_mask.shape)
|
|
# V_ind[V_mask] = tf.range(V_mask.sum())
|
|
# F_new = tf.gather(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
|
|
# V_new = V[V_mask, :]
|
|
|
|
# return V_new, F_new
|
|
|
|
|
|
def maskFacesBatch(V, F, mask):
|
|
Vs = []
|
|
Fs = []
|
|
for i in range(mask.shape[0]):
|
|
v, f = _maskFaces(V[i], F, mask[i])
|
|
Vs.append(v)
|
|
Fs.append(f)
|
|
return Vs, Fs
|
|
|
|
|
|
def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
|
|
fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ])
|
|
|
|
fig.add_trace(go.Mesh3d(
|
|
# 8 vertices of a cube
|
|
x=mesh_gt[:,0],
|
|
y=mesh_gt[:,2],
|
|
z=mesh_gt[:,1],
|
|
# i, j and k give the vertices of triangles
|
|
i = faces[:,0],
|
|
j = faces[:,1],
|
|
k = faces[:,2]
|
|
),
|
|
row=1, col=1)
|
|
|
|
fig.add_trace(go.Mesh3d(
|
|
# 8 vertices of a cube
|
|
x=mesh_pred_unposed[:,0],
|
|
y=mesh_pred_unposed[:,2],
|
|
z=mesh_pred_unposed[:,1],
|
|
# i, j and k give the vertices of triangles
|
|
i = faces[:,0],
|
|
j = faces[:,1],
|
|
k = faces[:,2]
|
|
),
|
|
row=1, col=2)
|
|
|
|
fig.add_trace(go.Mesh3d(
|
|
# 8 vertices of a cube
|
|
x=collision[:,0],
|
|
y=collision[:,2],
|
|
z=collision[:,1],
|
|
# i, j and k give the vertices of triangles
|
|
i = collision_faces[:,0],
|
|
j = collision_faces[:,1],
|
|
k = collision_faces[:,2],
|
|
vertexcolor=collision_color
|
|
),
|
|
row=1, col=3)
|
|
|
|
|
|
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
|
|
scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
|
|
scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
|
|
)
|
|
helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
|
|
yaxis = dict(nticks=4, range=[-1,1],),
|
|
zaxis = dict(nticks=4, range=[-1,1],),)
|
|
fig.update_layout(scene = helper,
|
|
scene2 = helper,
|
|
scene3 = helper,
|
|
)
|
|
|
|
return fig
|
|
|
|
|
|
def batch_to_dict(batch, signature):
|
|
result = dict()
|
|
for i, key in enumerate(signature.keys()):
|
|
result[key] = batch[i]
|
|
return result
|
|
|
|
conditions = ['pose_shape', 'uv_static', 'uv_mask', 'uv_shaped_body']
|
|
def nest_conditions(batch):
|
|
batch['conditions'] = {d: batch.pop(d) for d in conditions}
|
|
return batch
|
|
|
|
|
|
def sparse_to_tensor(L):
|
|
idx = np.mat([L.row, L.col]).transpose().astype(np.float32)
|
|
data = L.data.astype(dtype=np.float32)
|
|
shape = L.shape
|
|
return tf.SparseTensor(idx, data, shape)
|
|
|
|
|
|
def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False):
|
|
ckpt = tf.train.Checkpoint(model=model)
|
|
manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
|
|
if not disable_assert:
|
|
assert manager.latest_checkpoint is not None, checkpoint
|
|
assert manager.latest_checkpoint.split('/')[-1] != 'ckpt-0', checkpoint
|
|
if assert_consumed_on:
|
|
ckpt.restore(manager.latest_checkpoint).assert_consumed()
|
|
else:
|
|
ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed()
|