luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

static encoder runs, but doesn't work yet

Browse Source
This commit is contained in:
luni
2023-04-04 13:42:00 +00:00
parent 9394245a7e
commit 4d70ad5499
10 changed files with 1070 additions and 489 deletions
Download Patch File Download Diff File Expand all files Collapse all files
custom_utils/custom_utils.py
+80 -36
View File
@@ -1,12 +1,16 @@
import json
import numpy as np
import tensorflow as tf
# 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
# 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):
@@ -23,8 +27,8 @@ def m2uv_array(m2uv, is_smpl_orig=False):
return result
class UV2Mesh(tf.keras.Model):
def __init__(self, res=256, is_skirt=False, is_smpl_orig=False, is_adaptive=False):
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')
@@ -33,48 +37,40 @@ class UV2Mesh(tf.keras.Model):
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
self.is_adaptive = is_adaptive
if self.is_adaptive:
laplacian = sp.sparse.load_npz("custom_utils/laplacian/laplacian0" + ("_skirt" if is_skirt else "") + ".npz")
self.gc = GraphConvolution((3, 3), L=laplacian, batchSize=4)
@tf.function
def uv_to_mesh(self, uv):
if uv.shape[1] == 3:
uv = tf.transpose(uv, [0,2,3,1])
uv = torch.permute(uv, [0,2,3,1])
if uv.shape[1] != 512:
uv = tf.image.resize(uv, (512, 256), tf.image.ResizeMethod.LANCZOS5)
uv = Resize(uv[:,:,None], size=(512, 256), interpolation=InterpolationMode.BILINEAR)
batch_size = uv.shape[0]
# 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_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)
# 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)
# p1 = tf.stack((x1, y1, z1), axis=-1)
# p2 = tf.stack((x2, y2, z2), axis=-1)
mesh = (p1 + p2) / 2
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)
if self.is_adaptive:
x = self.gc(x)
return x
@@ -313,10 +309,16 @@ def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False
def get_summary_str(model):
lines = []
model.summary(print_fn=lines.append, line_length=150,expand_nested=True,show_trainable=True)
# 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
return lines
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)
@@ -341,7 +343,7 @@ def merge_verts(Vs, batchDim):
return V
class GraphConvolution(tf.keras.layers.Layer):
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
@@ -405,4 +407,46 @@ def group_inputs_for_levels(batch, config):
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
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)