436 lines
15 KiB
Python
436 lines
15 KiB
Python
import json
|
|
import numpy as np
|
|
# import tensorflow as tf
|
|
import torch
|
|
import torch.nn as nn
|
|
import scipy as sp
|
|
import plotly.graph_objects as go
|
|
from plotly.subplots import make_subplots
|
|
from torchvision.transforms import Resize
|
|
from torchvision.transforms.functional import InterpolationMode
|
|
|
|
# from tensorflow_graphics.geometry.convolution.graph_pooling import pool, unpool
|
|
# from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
|
|
|
|
|
|
def m2uv_array(m2uv, is_smpl_orig=False):
|
|
result = np.zeros((14475 if not is_smpl_orig else 3638, 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(nn.Module):
|
|
def __init__(self, res=256, is_skirt=False, is_smpl_orig=False):
|
|
super(UV2Mesh, self).__init__()
|
|
self.res = res
|
|
uv_map = np.load('custom_utils/stretched_uv_map_' + ('skirt' if is_skirt else '3') + ('_orig' if is_smpl_orig else '') + '.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 '') + ('_orig' if is_smpl_orig else '') + '.npy', allow_pickle=True).item()
|
|
m2uv = m2uv_array(m2uv, is_smpl_orig=is_smpl_orig)
|
|
self.uv_pixels = self.uv_map_discreet[m2uv]
|
|
self.smpl_res = 14475 if not is_smpl_orig else 3638
|
|
|
|
|
|
def uv_to_mesh(self, uv):
|
|
if uv.shape[1] == 3:
|
|
uv = torch.permute(uv, [0,2,3,1])
|
|
if uv.shape[1] != 512:
|
|
uv = Resize(uv[:,:,None], size=(512, 256), interpolation=InterpolationMode.BILINEAR)
|
|
|
|
# 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, self.smpl_res, 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, self.smpl_res, 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, self.smpl_res, 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, self.smpl_res, 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, self.smpl_res, 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, self.smpl_res, 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 = torch.mean(uv[:, self.uv_pixels[:,:,0], self.uv_pixels[:,:,1]], dim=2)
|
|
return mesh
|
|
|
|
|
|
def __call__(self, uv):
|
|
x = self.uv_to_mesh(uv)
|
|
return x
|
|
|
|
|
|
|
|
class UV2MeshLabel():
|
|
def __init__(self, res=256, is_skirt=False, is_smpl_orig=False):
|
|
super(UV2MeshLabel, self).__init__()
|
|
self.res = res
|
|
uv_map = np.load('custom_utils/stretched_uv_map_' + ('skirt' if is_skirt else '3') + ('_orig' if is_smpl_orig else '') + '.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 '') + ('_orig' if is_smpl_orig else '') + '.npy', allow_pickle=True).item()
|
|
m2uv = m2uv_array(m2uv, is_smpl_orig=is_smpl_orig)
|
|
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, return_unfixed=False):
|
|
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, :]
|
|
|
|
if return_unfixed:
|
|
return V_new, F_new, F[F_mask, :]
|
|
else:
|
|
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 flip_axis(x, ind=None):
|
|
y = x.copy()
|
|
if ind is not None:
|
|
y[ind:,1] = x[ind:,2]
|
|
y[ind:,2] = x[ind:,1]
|
|
else:
|
|
y[:,1] = x[:,2]
|
|
y[:,2] = x[:,1]
|
|
return y
|
|
|
|
|
|
def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color, flip_xy=False):
|
|
fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ])
|
|
|
|
if flip_xy:
|
|
mesh_gt = flip_axis(mesh_gt)
|
|
mesh_pred_unposed = flip_axis(mesh_pred_unposed)
|
|
collision = flip_axis(collision, ind=14475)
|
|
|
|
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 plot_3d_collision(gt, collision_1, collision_2, collision_faces, collision_color, flip_xy=False):
|
|
fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ])
|
|
|
|
if flip_xy:
|
|
collision_1 = flip_axis(collision_1, ind=14475)
|
|
collision_2 = flip_axis(collision_2, ind=14475)
|
|
gt = flip_axis(gt, ind=14475)
|
|
|
|
fig.add_trace(go.Mesh3d(
|
|
x=gt[:,0],
|
|
y=gt[:,2],
|
|
z=gt[:,1],
|
|
i = collision_faces[:,0],
|
|
j = collision_faces[:,1],
|
|
k = collision_faces[:,2],
|
|
vertexcolor=collision_color
|
|
),
|
|
row=1, col=1)
|
|
|
|
fig.add_trace(go.Mesh3d(
|
|
x=collision_1[:,0],
|
|
y=collision_1[:,2],
|
|
z=collision_1[:,1],
|
|
i = collision_faces[:,0],
|
|
j = collision_faces[:,1],
|
|
k = collision_faces[:,2],
|
|
vertexcolor=collision_color
|
|
),
|
|
row=1, col=2)
|
|
|
|
fig.add_trace(go.Mesh3d(
|
|
x=collision_2[:,0],
|
|
y=collision_2[:,2],
|
|
z=collision_2[:,1],
|
|
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 get_summary_str(model):
|
|
# from torchview import draw_graph
|
|
# draw_graph(model, input_size=((4,3,512,256), (4,512,256)), expand_nested=True)
|
|
|
|
# lines = []
|
|
# model.summary(print_fn=lines.append, line_length=150,expand_nested=True,show_trainable=True)
|
|
# Add initial spaces to avoid markdown formatting in TensorBoard
|
|
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
|
total_params = sum(param.numel() for param in model.parameters())
|
|
|
|
return [f'all_params: {total_params}, trainable: {trainable_params}']
|
|
|
|
def get_summary_html(model):
|
|
lines = get_summary_str(model)
|
|
html = ''
|
|
for line in lines:
|
|
html += "<pre>" + line + "</pre>"
|
|
return html
|
|
|
|
|
|
def repeat_laplacian(L, times):
|
|
idx = np.repeat(np.mat([L.row, L.col]), repeats=times, axis=1)
|
|
idx_offset = np.linspace(0, L.shape[0] * (times-1), times).astype(np.int32)
|
|
idx_offset = np.repeat(idx_offset, L.row.shape[0], axis=0)
|
|
idx += idx_offset[np.newaxis,:]
|
|
data = np.repeat(L.data, repeats=times, axis=0)
|
|
shape = (L.shape[0] * times, L.shape[1] * times)
|
|
return sp.sparse.coo_matrix((data, (np.squeeze(np.asarray(idx[0,:])), np.squeeze(np.asarray(idx[1,:])))), shape)
|
|
|
|
|
|
def merge_verts(Vs, batchDim):
|
|
V = tf.concat(tf.unstack(Vs, num=batchDim), axis=0)
|
|
return V
|
|
|
|
|
|
class GraphConvolution(nn.Module):
|
|
def __init__(self, shape, activation=None, name='gcn', L=None, batchSize=None, normalization=None):
|
|
super(GraphConvolution, self).__init__()
|
|
self.shape = shape
|
|
self.w = self.add_weight(name=name + '/weight', shape=(2, *shape), initializer=tf.initializers.glorot_normal, trainable=True)
|
|
self.b = self.add_weight(name=name + '/bias', shape=shape[-1], initializer=tf.initializers.zeros, trainable=True)
|
|
self.activation = activation or (lambda x: x)
|
|
self.normalization = normalization
|
|
self.batchSize = batchSize
|
|
self.indices = tf.cast(tf.linspace(0, self.batchSize * L.shape[1], self.batchSize+1), dtype=tf.int32)
|
|
self.L = repeat_laplacian(L, self.batchSize)
|
|
|
|
def gather(self):
|
|
return [self.w, self.b]
|
|
|
|
def _single(self, X):
|
|
X0 = tf.einsum('ab,bc->ac', X, self.w[0]) # Node
|
|
X1 = tf.einsum('ab,bc->ac', X, self.w[1]) # Neigh
|
|
# Neigh conv
|
|
X1 = edge_convolution_template(
|
|
X1, sparse_to_tensor(self.L), sizes=None,
|
|
edge_function=lambda x, y: y,
|
|
reduction='weighted',
|
|
edge_function_kwargs={}
|
|
)
|
|
X = X0 + X1 + self.b
|
|
return X
|
|
|
|
def call(self, X, training):
|
|
X = merge_verts(X, self.batchSize)
|
|
X = self._single(X)
|
|
X = tf.stack([X[self.indices[i-1]:self.indices[i]] for i in range(1, len(self.indices))])
|
|
X = self.activation(X)
|
|
if self.normalization:
|
|
X = self.normalization(X, training=training)
|
|
return X
|
|
|
|
|
|
def replace_arrays(string):
|
|
out = string[:1]
|
|
i = 0
|
|
while i < len(string) and string[i+1:].find('[') != -1:
|
|
start = string[i+1:].find('[')
|
|
end = string[i+1+start:].find(']')
|
|
out += string[i+1:i+1+start]
|
|
array = string[i+1+start:i+1+start+end+1].replace('\n', '')
|
|
out += array
|
|
i = i + 1 + start + end
|
|
|
|
out += string[i+1:]
|
|
|
|
return out
|
|
|
|
|
|
def group_inputs_for_levels(batch, config):
|
|
levels_input = []
|
|
for res in config['pyramid']['resolutions']:
|
|
res_modifier = '_' + str(res) if res != 512 else ''
|
|
if config['use_skinning']:
|
|
uv_orig = batch['uv_unposed' + res_modifier]
|
|
else:
|
|
uv_orig = batch['uv_posed' + res_modifier]
|
|
levels_input.append([uv_orig, batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier], batch['uv_mask' + res_modifier], batch['uv_mask']])
|
|
conds = [batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature']]
|
|
return levels_input, conds
|
|
|
|
|
|
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()
|
|
|
|
|
|
from stl import mesh as mesh_stl
|
|
|
|
def save_stl(file, mesh, faces):
|
|
obj = mesh_stl.Mesh(np.zeros(faces.shape[0], dtype=mesh_stl.Mesh.dtype))
|
|
for i, f in enumerate(faces):
|
|
for j in range(3):
|
|
obj.vectors[i][j] = mesh[f[j],:]
|
|
|
|
obj.save(file)
|
|
|
|
|
|
def write_obj(file, vertices, faces, colors):
|
|
with open(file, 'w') as file:
|
|
for v, c in zip(vertices, colors):
|
|
line = 'v ' + ' '.join([str(_) for _ in v]) + ' ' + ' '.join([str(_) for _ in c]) + '\n'
|
|
file.write(line)
|
|
for f in faces:
|
|
line = 'f ' + ' '.join([str(_) for _ in f]) + '\n'
|
|
file.write(line) |