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

backup

Browse Source
This commit is contained in:
luni
2022-11-04 16:17:58 +00:00
parent ebe8522905
commit 1c6b805519
18 changed files with 1413 additions and 486 deletions
Download Patch File Download Diff File Expand all files Collapse all files
calculate_error.py
+151
View File
@@ -0,0 +1,151 @@
import os
import sys
from configs.config_abl_abs_skirt import config
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
# 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 csv
import glob
import time
import argparse
import wandb
import tqdm
import random
import numpy as np
import plotly.graph_objects as go
import plotly.express as ex
# if config['is_skirt']:
# from losses import edge_loss_skirt as edge_loss
# else:
# from losses import edge_loss
# from postprocessing import fix_collision
# from losses import collision_ratio, kl_loss #, concept_loss
# # from losses import collision_penalty as collision_penalty
# from losses import collision_loss as collision_penalty
# from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
# from custom_utils.smpl import SMPL
# from custom_utils.supersmpl_io import writeOBJ
# uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
# smpl = SMPL(use_one_hot=config['use_one_hot'])
skirt_modifier = '_skirt' if config['is_skirt'] else ''
folder = config['test_dir']
checkpoint_name = config['name']
iteration = '70000'
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
garment_type_table = np.load('custom_data/garment_type_table_test.npy', allow_pickle=True).item()
output_folder = '/home/workspace/data/results/' + checkpoint_name + '-' + iteration
faces_folder = '/home/workspace/data/results/support/faces' + ('_skirt' if config['is_skirt'] else '')
ids = os.listdir(faces_folder)
ids = [id.split('.')[0] for id in ids]
faces = {id : np.load(faces_folder + '/' + id + '.npy') for id in ids }
masks_folder = '/home/workspace/data/dataset/test' + ('_skirt' if config['is_skirt'] else '') + '_ready/mesh_mask'
ids = os.listdir(masks_folder)
ids = [id.split('.')[0] for id in ids]
masks = {id : np.load(masks_folder + '/' + id + '.npy') for id in ids }
frames = os.listdir(output_folder)
frames = [frame for frame in frames if not frame.endswith('_fixed_coll.npy')]
if config['remove_global_orientation']:
gt_name = 'mesh_posed_no_orient'
else:
gt_name = 'mesh_posed'
orig_folder = '/home/workspace/data/dataset/test' + ('_skirt' if config['is_skirt'] else '') + '_ready' + ('_2' if config['is_skirt'] else '') + '/' + gt_name
errors = dict()
if config['is_skirt']:
per_garment_errors = { 'Tshirt': [], 'Dress': [], 'Skirt': [], 'Top': [] }
else:
per_garment_errors = { 'Tshirt': [], 'Jumpsuit': [], 'Trousers': [], 'Top': [] }
for frame in tqdm.tqdm(frames):
id = frame.split('_')[0]
frame_nr = frame.split('.')[0].split('_')[1]
gt = np.load(os.path.join(orig_folder, id + '_' + frame_nr + '.npz'))[gt_name][masks[id], :].astype(np.float32)
pred = np.load(os.path.join(output_folder, id + '_' + frame_nr + '.npy'))
metric = np.mean(np.sqrt(np.sum((gt - pred) ** 2, axis=-1)))
errors[id + '_' + frame_nr] = metric
per_garment_errors[garment_type_table[id]].append(metric)
for key in per_garment_errors.keys():
print(key + ': ' + str(np.mean(per_garment_errors[key])))
with open('/home/workspace/data/results/errors/' + checkpoint_name + '-' + iteration + '.csv', mode='w') as f:
file_writer = csv.writer(f, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
for key in sorted(errors.keys()):
file_writer.writerow([key, errors[key]])
# x = gt
# y = pred
# fig = go.Figure(data=[
# go.Mesh3d(
# x=x[:,0],
# y=x[:,1],
# z=x[:,2],
# i = faces[id][:,0],
# j = faces[id][:,1],
# k = faces[id][:,2],
# ),
# go.Mesh3d(
# x=y[:,0],
# y=y[:,1],
# z=y[:,2],
# i = faces[id][:,0],
# j = faces[id][:,1],
# k = faces[id][:,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()
# print()
convert_to_uv_local_single.py
+18 -10
View File
@@ -195,7 +195,8 @@ def convert_to_uv_batch(res, m2uv, uv_map, verts, neighbours=15, mask=None, meth
if garment_uv_mask is not None:
result *= garment_uv_mask[:,:,np.newaxis]
np.savez_compressed(os.path.join(local_folder, 'uv_unposed_no_shape_one_hot_2', id + '_' + str(i) + '.npz'), uv_unposed_no_shape=result)
np.savez_compressed(os.path.join(local_folder, 'uv_posed_no_orient', id + '_' + str(i) + '.npz'), uv_posed_no_orient=result)
# np.savez_compressed(os.path.join(local_folder, 'uv_unposed_no_shape_one_hot_2', id + '_' + str(i) + '.npz'), uv_unposed_no_shape=result)
@@ -383,17 +384,18 @@ def test(verts, faces, values):
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/local-data/train')
parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/local-data/valid')
parser.add_argument("-s", "--start", help = "start of range", default=0)
parser.add_argument("-e", "--end", help = "end of range", default=10)
parser.add_argument("-e", "--end", help = "end of range", default=1500)
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
parser.add_argument("-sk", "--skirt", help = "end of range", default=0)
parser.add_argument("-sk", "--skirt", help = "end of range", default=1)
args = parser.parse_args()
# smpl = SMPL('FeatureLines/superSMPL')
local_folder = args.folder
os.makedirs(os.path.join(local_folder, 'uv_unposed_no_shape_one_hot_2'), exist_ok=True)
os.makedirs(os.path.join(local_folder, 'uv_posed_no_orient'), exist_ok=True)
# os.makedirs(os.path.join(local_folder, 'uv_unposed_no_shape_one_hot_2'), exist_ok=True)
print("skirt" if args.skirt == 1 else "not skirt")
@@ -432,8 +434,11 @@ if __name__ == '__main__':
# rbf_dest[:, [0,1]] = rbf_dest[:, [1, 0]]
# rbf_dest = rbf_dest[unique_pairs(rbf_dest), :] # not good, unique_pairs doesn't consider coord order
ids = os.listdir(os.path.join(local_folder, 'mesh_static'))[int(args.start):int(args.end)]
ids = [id.split('.')[0] for id in ids]
# ids = os.listdir(os.path.join(local_folder, 'mesh_static'))[int(args.start):int(args.end)]
# ids = [id.split('.')[0] for id in ids]
ids = os.listdir(os.path.join(local_folder, 'pose_shape'))
ids = [id.split('_')[0] for id in ids]
ids = sorted(list(set(ids)))[int(args.start):int(args.end)]
for id in tqdm.tqdm(ids):
@@ -463,8 +468,10 @@ if __name__ == '__main__':
# mesh_unposed_no_shape = np.stack([mesh_unposed_no_shape[i]['mesh_unposed_no_shape'] for i in range(gt.shape[0])])
# body_no_orient = [np.load(os.path.join(folder, id, 'mesh_posed_body_no_orient', id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])]
# body_no_orient = np.stack([body_no_orient[i]['mesh_posed_body_no_orient'] for i in range(gt.shape[0])])
mesh_unposed_no_shape_one_hot = [np.load(os.path.join(local_folder, 'mesh_unposed_no_shape_one_hot', id + '_' + str(i) + '.npz')) for i in range(frames)]
mesh_unposed_no_shape_one_hot = [mesh_unposed_no_shape_one_hot[i]['mesh_unposed_no_shape'] for i in range(frames)]
# mesh_unposed_no_shape_one_hot = [np.load(os.path.join(local_folder, 'mesh_unposed_no_shape_one_hot', id + '_' + str(i) + '.npz')) for i in range(frames)]
# mesh_unposed_no_shape_one_hot = [mesh_unposed_no_shape_one_hot[i]['mesh_unposed_no_shape'] for i in range(frames)]
mesh_posed_no_orient = [np.load(os.path.join(local_folder, 'mesh_posed_no_orient', id + '_' + str(i) + '.npz')) for i in range(frames)]
mesh_posed_no_orient = [mesh_posed_no_orient[i]['mesh_posed_no_orient'] for i in range(frames)]
# verts_dynamic = gt
@@ -487,7 +494,8 @@ if __name__ == '__main__':
# uv_curvature = convert_to_uv_batch_label(uv_res, m2uv, uv_map, curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
# uv_unposed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
# uv_unposed_no_shape = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed_no_shape[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
uv_unposed_no_shape_one_hot = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed_no_shape_one_hot, neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border, id=id)
# uv_unposed_no_shape_one_hot = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed_no_shape_one_hot, neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border, id=id)
uv_posed_no_orient = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_posed_no_orient, neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border, id=id)
# uv_shaped_body = convert_to_uv(uv_res, m2uv, uv_map, shaped_body, neighbours=3, mask=None, method='rbf', garment_uv_mask=None).astype(np.float32)
print(time.perf_counter() - start)
custom_utils/custom_utils.py
+59 -7
View File
@@ -144,9 +144,25 @@ def maskFacesBatch(V, F, mask):
return Vs, Fs
def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
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],
@@ -200,8 +216,24 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli
return fig
def plot_3d_collision(collision_1, collision_2, collision_faces, collision_color):
fig = make_subplots(rows=1, cols=2, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(2)] ])
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],
@@ -212,7 +244,7 @@ def plot_3d_collision(collision_1, collision_2, collision_faces, collision_color
k = collision_faces[:,2],
vertexcolor=collision_color
),
row=1, col=1)
row=1, col=2)
fig.add_trace(go.Mesh3d(
x=collision_2[:,0],
@@ -223,17 +255,19 @@ def plot_3d_collision(collision_1, collision_2, collision_faces, collision_color
k = collision_faces[:,2],
vertexcolor=collision_color
),
row=1, col=2)
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
@@ -258,7 +292,7 @@ def sparse_to_tensor(L):
return tf.SparseTensor(idx, data, shape)
def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False):
def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False, iteration=''):
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
if not disable_assert:
@@ -267,7 +301,12 @@ def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False
if assert_consumed_on:
ckpt.restore(manager.latest_checkpoint).assert_consumed()
else:
ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed()
if iteration != '':
name_base = manager.latest_checkpoint.split('-')[0]
ckpt_ind = [i for i in range(len(manager.checkpoints)) if manager.checkpoints[i].endswith(name_base + '-' + iteration)][0]
ckpt.restore(manager.checkpoints[ckpt_ind]).assert_existing_objects_matched()#.assert_consumed()
else:
ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed()
def get_summary_str(model):
@@ -351,3 +390,16 @@ def replace_arrays(string):
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
custom_utils/smpl.py
+19 -2
View File
@@ -34,12 +34,14 @@ eval_joints = {
class SMPL:
def __init__(self, orig_smpl=False, use_one_hot=False):
def __init__(self, orig_smpl=False, use_one_hot=False, is_skirt=False):
self.J_regressor = []
self.v_template = []
self.shapedirs = []
self.posedirs = []
self.weights = []
self.use_one_hot = use_one_hot
self.is_skirt = is_skirt
# Load SMPL models
# path = os.path.abspath(os.path.dirname(__file__))
with open(os.path.join('custom_utils/super_model_f.pkl'), 'rb') as f:
@@ -299,7 +301,22 @@ class SMPL:
def run_garment(self, pose, betas, v_cloth, v_shaped=None):
J = self.get_joints_template(betas, v_shaped)
R = self.run_rodrigues(pose)
if self.use_one_hot and self.is_skirt:
# p_avg = (pose[:, 3:6] + pose[:, 6:9]) / 2
# pose_dress = tf.concat([pose[:, :3], p_avg, p_avg, pose[:, 9:]], axis=-1)
# R = self.run_rodrigues(pose_dress).numpy()
# for b in range(betas.shape[0]):
# R[b][[4,5,7,8,10,11],:,:] = np.eye(3)
pose_dress = pose
# p_avg = (pose_dress[:, 3:6] + pose_dress[:, 6:9]) / 2
# pose_dress[:,3:9] = np.concatenate([p_avg, p_avg], axis=-1)
R = self.run_rodrigues(pose_dress).numpy()
for b in range(betas.shape[0]):
R[b][[4,5,7,8,10,11],:,:] = np.eye(3)
else:
R = self.run_rodrigues(pose)
# if filter is not None:
# weights = tf.boolean_mask(self.weights, filter)
cvae.py
+122 -68
View File
@@ -1,6 +1,6 @@
import os
import sys
from configs.config_cvae_test import config
from configs.config_abl_norm import config
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
import tensorflow as tf
@@ -26,7 +26,6 @@ import numpy as np
import plotly.graph_objects as go
import plotly.express as ex
from PIL import Image
from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
from postprocessing import fix_collision
if config['is_skirt']:
@@ -36,30 +35,30 @@ else:
from losses import collision_ratio, kl_loss #, concept_loss
from losses import collision_ratio, kl_loss, normal_loss #, concept_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
from dataset import DatasetFactory, StaticDataset
from custom_utils.smpl import SMPL
from models.cvae_model import CVAE
def init_wandb():
run = wandb.init(project='cvae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp')
run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('cvae.py')
artifact.add_file('models/cvae_model.py')
artifact.add_file('losses.py')
artifact.add_file('configs/config_base.py')
artifact.add_file('configs/config_cvae.py')
# artifact.add_file('configs/config_base.py')
# artifact.add_file('configs/config_cvae.py')
artifact.add_file('models/base.py')
artifact.add_file('models/cond_encoder.py')
artifact.add_file('dataset.py')
artifact.add_file('convnext/models/convnext_tf.py')
config_file_name = [key for key in sys.modules.keys() if key.startswith('configs.')]
assert len(config_file_name) == 1
artifact.add_file('/'.join(config_file_name[0].split('.')) + '.py')
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')]
for config_file_name in config_file_names:
artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False)
@@ -76,28 +75,42 @@ def kl_loss(mean, logvar):
def train_step(batch):
with tf.GradientTape() as tape:
pred, mean, logvar = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
if config['use_skinning']:
mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
else:
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed']
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
if config['use_skinning']:
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
else:
mesh_pred_posed = mesh_pred
# is the output already masked?
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
n_loss = normal_loss(mesh_pred, mesh_orig, batch['mesh_mask'], is_skirt=config['is_skirt'])
# conc_loss = concept_loss(pred, batch['uv_unposed'])
kl = kl_loss(mean, logvar)
reconstr_loss = (l2_loss + l2_mesh_loss) * config['l2_coeff'] + (e_loss + c_loss) * config['reg_loss_coeff']
reconstr_loss = l2_loss
if config['use_extra_losses']:
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
if config['use_norm_loss'] and step > 2000:
reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
loss = reconstr_loss + g_balancing_coeff * kl
mean = tf.reduce_mean(mean)
@@ -105,41 +118,54 @@ def train_step(batch):
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
# c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/euclid_error': metric,
'train/reconstr_loss': reconstr_loss,
'train/reconstr_loss': reconstr_loss, 'train/collision_ratio': c_ratio,
'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
'train/l2_loss': l2_loss, 'train/l2_mesh_loss': l2_mesh_loss
# 'train/ms_ssim': ms_ssim, 'train/collision_ratio': c_ratio,
'train/l2_loss': l2_loss, 'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss,
# 'train/ms_ssim': ms_ssim,
}
@tf.function
def valid_step(batch):
pred, mean, logvar = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
if config['use_skinning']:
mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
else:
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed']
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
if config['use_skinning']:
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
else:
mesh_pred_posed = mesh_pred
# is the output already masked?
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
n_loss = normal_loss(mesh_pred, mesh_orig, batch['mesh_mask'], is_skirt=config['is_skirt'])
# conc_loss = concept_loss(pred, batch['uv_unposed'])
kl = kl_loss(mean, logvar)
reconstr_loss = (l2_loss + l2_mesh_loss) * config['l2_coeff'] + (e_loss + c_loss) * config['reg_loss_coeff']
reconstr_loss = l2_loss
if config['use_extra_losses']:
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
if config['use_norm_loss']:
reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
loss = reconstr_loss + g_balancing_coeff * kl
mean = tf.reduce_mean(mean)
@@ -150,49 +176,62 @@ def valid_step(batch):
# c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/euclid_error': metric,
'valid/reconstr_loss': reconstr_loss, 'valid/l2_mesh_loss': l2_mesh_loss,
'valid/reconstr_loss': reconstr_loss, 'valid/l2_mesh_loss': l2_mesh_loss, 'valid/collision_ratio': c_ratio,
'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar,
# 'valid/ms_ssim': ms_ssim, 'valid/collision_ratio': c_ratio,
'valid/normal_loss': n_loss,# 'valid/ms_ssim': ms_ssim,
}
@tf.function
def test_step(batch):
pred, mean, logvar = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
if config['use_skinning']:
mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
else:
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed']
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
if config['use_skinning']:
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
mesh_orig_shaped = mesh_orig + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_orig_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_orig_shaped)
else:
mesh_pred_posed = mesh_pred
mesh_orig_posed = mesh_orig
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
mesh_pred_posed_fixed = fix_collision(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
mesh_orig_shaped = mesh_orig + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_orig_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_orig_shaped)
# is the output already masked?
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
# conc_loss = concept_loss(pred, batch['uv_unposed'])
kl = kl_loss(mean, logvar)
reconstr_loss = (l2_loss + l2_mesh_loss) * config['l2_coeff'] + (e_loss + c_loss) * config['reg_loss_coeff']
reconstr_loss = l2_loss
if config['use_extra_losses']:
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
loss = reconstr_loss + g_balancing_coeff * kl
mean = tf.reduce_mean(mean)
logvar = tf.reduce_mean(logvar)
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, uv_orig + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
euclid_posed = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_posed - mesh_pred_posed) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
euclid_posed_fixed = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_posed - mesh_pred_posed_fixed) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
return { 'test/loss': loss, 'test/edge_loss': e_loss, 'test/collision_loss': c_loss, 'test/euclid_error': metric, 'test/ms_ssim': ms_ssim,
'test/reconstr_loss': reconstr_loss, 'test/collision_ratio': c_ratio, 'test/l2_mesh_loss': l2_mesh_loss,
@@ -201,32 +240,46 @@ def test_step(batch):
def plot_everything(batch: dict, type: str):
pred, mean, logvar = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
if config['use_skinning']:
mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
else:
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
mesh_pred = uv2mesh(pred).numpy()
mesh_orig = batch['mesh_unposed']
orig_img = batch['uv_unposed']
orig_img = uv_orig
uv_static_img = batch['uv_static']
uv_body_posed_img = batch['uv_body_posed']
pred_img = pred
img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy")
mask = batch['mesh_mask'].numpy().astype(bool)
garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
gt = [mesh_orig[i].numpy()[mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(pred.shape[0])]
if config['use_skinning']:
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
gt = [mesh_orig[i].numpy()[mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(pred.shape[0])]
for i in range(mesh_orig.shape[0]):
mesh_pred[i, mask[i], :] += shape_offsets.numpy()[i, mask[i], :]
mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
else:
mesh_posed = mesh_pred
gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
gt_posed = [batch['mesh_posed'].numpy()[i, mask[i], :] for i in range(mask.shape[0])]
body = batch['mesh_body_posed'].numpy()
for i in range(body.shape[0]):
mesh_pred[i, mask[i], :] += shape_offsets.numpy()[i, mask[i], :]
mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
mesh_posed_fixed = fix_collision(mesh_posed, body, mask).numpy()
garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
@@ -236,9 +289,10 @@ def plot_everything(batch: dict, type: str):
garment_posed_fixed = [mesh_posed_fixed[i, mask[i], :] for i in range(body.shape[0])]
collision_fixed = [np.concatenate((body[i], garment_posed_fixed[i]), axis=0) for i in range(len(garment_posed))]
collision_gt = [np.concatenate((body[i], gt_posed[i]), axis=0) for i in range(len(garment_posed))]
figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision[i], collision_fixed[i], collision_faces[i], collision_color[i]) for i in range(4) }
figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0), num_bins=model.latent_dim)
logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0), num_bins=model.latent_dim)
@@ -417,7 +471,7 @@ def test():
if __name__ == '__main__':
# init_wandb()
init_wandb()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
@@ -488,10 +542,10 @@ if __name__ == '__main__':
if step == 0:
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
if train_loss['train/euclid_error'] < 0.019:
if train_loss['train/euclid_error'] < 0.017:
if balancing_activated_step == 0:
balancing_activated_step = step
g_balancing_coeff = g_balancing_coeff + 1 / 10000
g_balancing_coeff = g_balancing_coeff + 1 / 20000
to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } }
cvae_pyramid.py
+153 -96
View File
@@ -1,5 +1,7 @@
import os
from configs.config_pyramid import config
import sys
from configs.config_abl_pyr_64_128_cape_m import config
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
import tensorflow as tf
@@ -22,7 +24,7 @@ import wandb
import random
import numpy as np
import scipy as sp
from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, get_summary_str
from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, get_summary_str, plot_3d_collision, group_inputs_for_levels
import plotly.graph_objects as go
import plotly.express as ex
from plotly.subplots import make_subplots
@@ -33,7 +35,6 @@ if config['is_skirt']:
else:
from losses import edge_loss
from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
from losses import uv_mask_loss, bend_loss, kl_loss, collision_ratio
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
@@ -41,24 +42,26 @@ from custom_utils.custom_utils import batch_to_dict, load_model
from dataset import DatasetFactory, StaticDataset
from custom_utils.smpl import SMPL
from models.pyramid_model_optimized import Pyramid
from postprocessing import fix_collision
def init_wandb():
run = wandb.init(project='cvae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp')
run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/wandb')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('cvae_pyramid.py')
artifact.add_file('models/cvae_model.py')
artifact.add_file('models/pyramid_model.py')
artifact.add_file('models/pyramid_model_optimized.py')
artifact.add_file('models/pyramid_level.py')
artifact.add_file('models/upscaler.py')
artifact.add_file('losses.py')
artifact.add_file('configs/config_base.py')
artifact.add_file('configs/config_pyramid.py')
artifact.add_file('models/base.py')
artifact.add_file('models/cond_encoder.py')
artifact.add_file('dataset.py')
artifact.add_file('convnext/models/convnext_tf.py')
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')]
for config_file_name in config_file_names:
artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
@@ -205,26 +208,18 @@ def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, col
return fig
def group_inputs_for_levels(batch):
levels_input = []
for res in config['pyramid']['resolutions']:
res_modifier = '_' + str(res) if res != 512 else ''
levels_input.append([batch['uv_unposed' + res_modifier], 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
@tf.function
def train_step(batch):
levels_input, conds = group_inputs_for_levels(batch)
levels_input, conds = group_inputs_for_levels(batch, config)
with tf.GradientTape() as tape:
out_levels = model([levels_input, conds], training=True)
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
mesh_orig = batch['mesh_unposed']
if config['use_skinning']:
mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
else:
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
output = dict()
reconstr_loss_total = 0
@@ -237,21 +232,36 @@ def train_step(batch):
pred, mean, logvar = out_levels[i]
mesh_pred = uv2mesh(pred)
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
if config['use_skinning']:
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
else:
mesh_pred_posed = mesh_pred
# is the output already masked?
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
# ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
kl = kl_loss(mean, logvar)
reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
# reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
if config['use_extra_losses']:
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss)
if not config['is_cape']:
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + c_loss) * config['reg_loss_coeff']
else:
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
kl_total += kl
mean = tf.reduce_mean(mean)
@@ -260,9 +270,9 @@ def train_step(batch):
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,'train/euclid_error' + res_modifier: metric,
'train/ms_ssim' + res_modifier: ms_ssim, 'train/mean' + res_modifier: mean, 'train/logvar' + res_modifier: logvar, 'train/kl_loss' + res_modifier: kl,
'train/var' + res_modifier: np.e ** logvar, 'train/l2_loss' + res_modifier: l2_loss,
'train/euclid_error': metric, 'train/l2_mesh_loss': l2_mesh_loss } }
'train/mean' + res_modifier: mean, 'train/logvar' + res_modifier: logvar, 'train/kl_loss' + res_modifier: kl,
'train/var' + res_modifier: np.e ** logvar, 'train/l2_loss' + res_modifier: l2_loss,
'train/euclid_error': metric, 'train/l2_mesh_loss': l2_mesh_loss } } # 'train/ms_ssim' + res_modifier: ms_ssim,
loss = reconstr_loss_total + balancing_coeff * kl_total
@@ -273,13 +283,14 @@ def train_step(batch):
@tf.function
def valid_step(batch):
levels_input, conds = group_inputs_for_levels(batch)
levels_input, conds = group_inputs_for_levels(batch, config)
out_levels = model([levels_input, conds], training=False)
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
mesh_orig = batch['mesh_unposed']
if config['use_skinning']:
mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
else:
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
output = dict()
reconstr_loss_total = 0
@@ -292,21 +303,34 @@ def valid_step(batch):
pred, mean, logvar = out_levels[i]
mesh_pred = uv2mesh(pred)
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
if config['use_skinning']:
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
else:
mesh_pred_posed = mesh_pred
# is the output already masked?
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
# ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
kl = kl_loss(mean, logvar)
reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
# reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
if config['use_extra_losses']:
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss)
if not config['is_cape']:
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + c_loss) * config['reg_loss_coeff']
else:
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
kl_total += kl
mean = tf.reduce_mean(mean)
@@ -315,9 +339,9 @@ def valid_step(batch):
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
output = { **output, **{ 'valid/edge_loss' + res_modifier: e_loss, 'valid/collision_loss' + res_modifier: c_loss,'valid/euclid_error' + res_modifier: metric,
'valid/ms_ssim' + res_modifier: ms_ssim, 'valid/mean' + res_modifier: mean, 'valid/logvar' + res_modifier: logvar, 'valid/kl_loss' + res_modifier: kl,
'valid/mean' + res_modifier: mean, 'valid/logvar' + res_modifier: logvar, 'valid/kl_loss' + res_modifier: kl,
'valid/var' + res_modifier: np.e ** logvar, 'valid/l2_loss' + res_modifier: l2_loss,
'valid/euclid_error': metric, 'valid/l2_mesh_loss': l2_mesh_loss } }
'valid/euclid_error': metric, 'valid/l2_mesh_loss': l2_mesh_loss } } # 'valid/ms_ssim' + res_modifier: ms_ssim,
loss = reconstr_loss_total + balancing_coeff * kl_total
@@ -326,7 +350,7 @@ def valid_step(batch):
@tf.function
def test_step(batch):
levels_input, conds = group_inputs_for_levels(batch)
levels_input, conds = group_inputs_for_levels(batch, config)
out_levels = model([levels_input, conds], training=False)
shape = batch['pose_shape'][:, 72:]
@@ -351,14 +375,17 @@ def test_step(batch):
# is the output already masked?
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
# ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
kl = kl_loss(mean, logvar)
reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss)
if not config['is_cape']:
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + c_loss) * config['reg_loss_coeff']
# reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
kl_total += kl
@@ -369,7 +396,7 @@ def test_step(batch):
euclid_error_per_garment = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1)), axis=1) / tf.reduce_sum(batch['mesh_mask'], axis=1)
output = { **output, **{ 'test/edge_loss' + res_modifier: e_loss, 'test/collision_loss' + res_modifier: c_loss,'test/euclid_error' + res_modifier: metric,
'test/ms_ssim' + res_modifier: ms_ssim, 'test/mean' + res_modifier: mean, 'test/logvar' + res_modifier: logvar, 'test/kl_loss' + res_modifier: kl,
'test/mean' + res_modifier: mean, 'test/logvar' + res_modifier: logvar, 'test/kl_loss' + res_modifier: kl,
'test/var' + res_modifier: np.e ** logvar, 'test/l2_loss' + res_modifier: l2_loss,
'test/euclid_error': metric, 'test/collision_ratio' + res_modifier: c_ratio } }
@@ -380,15 +407,19 @@ def test_step(batch):
def plot_everything(batch: dict, type: str):
levels_input = group_inputs_for_levels(batch)
if config['use_skinning']:
mesh_orig, uv_orig = batch['mesh_unposed'].numpy(), batch['uv_unposed'].numpy()
else:
mesh_orig, uv_orig = batch['mesh_posed'].numpy(), batch['uv_posed'].numpy()
levels_input = group_inputs_for_levels(batch, config)
output_levels = model(levels_input, training=False)
mesh_pred_levels = [uv2mesh(output_levels[i][0]).numpy() for i in range(len(output_levels))]
mesh_orig = batch['mesh_unposed'].numpy()
mesh_orig_posed = batch['mesh_posed'].numpy()
body = batch['mesh_body_posed'].numpy()
orig_img = (batch['uv_unposed'].numpy() + 1) / 2
orig_img = (uv_orig + 1) / 2
uv_static_img = (batch['uv_static'].numpy() + 1) / 2
uv_body_posed_img = (batch['uv_body_posed'].numpy() + 1) / 2
uv_curvature = (np.tile(batch['uv_curvature'].numpy().argmax(-1)[:,:,:,np.newaxis], [1,1,1,3]) / 6)
@@ -397,43 +428,62 @@ def plot_everything(batch: dict, type: str):
img = tf.concat((uv_static_img, uv_body_posed_img, uv_curvature, orig_img, pred_img), axis=2)
imgs = [wandb.Image(img[i]) for i in range(4)]
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]).numpy()
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
mask = batch['mesh_mask'].numpy().astype(bool)
_, faces = maskFacesBatch(np.ones_like(mesh_orig), F, mask)
gt = [mesh_orig[i][mask[i], :] + shape_offsets[i, mask[i], :] for i in range(mask.shape[0])]
if config['use_skinning']:
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]).numpy()
gt = [mesh_orig[i][mask[i], :] + shape_offsets[i, mask[i], :] for i in range(mask.shape[0])]
mesh_unposed_shaped_levels = [[] for _ in range(mask.shape[0])]
for i in range(len(output_levels)):
mesh_unposed_shaped = [mesh_pred_levels[i][j, mask[j], :] + shape_offsets[j, mask[j]] for j in range(mask.shape[0])]
for k in range(mask.shape[0]):
mesh_unposed_shaped_levels[k].append(mesh_unposed_shaped[k])
# full res version
mesh_unposed_shaped_full = np.zeros((mask.shape[0], config['smpl_res'], 3), dtype=np.float32)
for i in range(mask.shape[0]):
mesh_unposed_shaped_full[i, mask[i], :] = mesh_unposed_shaped[i]
mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_unposed_shaped_full).numpy()
else:
mesh_posed = mesh_pred_levels[-1]
gt = [mesh_orig[i][mask[i], :] for i in range(mesh_posed.shape[0])]
mesh_unposed_shaped_levels = [[mesh_pred_levels[i][j, mask[j], :] for i in range(len(output_levels))] for j in range(mask.shape[0])]
gt_posed = [mesh_orig_posed[i, mask[i], :] for i in range(mask.shape[0])]
mesh_posed_fixed = fix_collision(mesh_posed, body, mask).numpy()
mesh_unposed_shaped_levels = [[] for _ in range(mask.shape[0])]
for i in range(len(output_levels)):
mesh_unposed_shaped = [mesh_pred_levels[i][j, mask[j], :] + shape_offsets[j, mask[j]] for j in range(mask.shape[0])]
for k in range(mask.shape[0]):
mesh_unposed_shaped_levels[k].append(mesh_unposed_shaped[k])
# full res version
mesh_unposed_shaped_full = np.zeros((mask.shape[0], config['smpl_res'], 3), dtype=np.float32)
for i in range(mask.shape[0]):
mesh_unposed_shaped_full[i, mask[i], :] = mesh_unposed_shaped[i]
mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_unposed_shaped_full).numpy()
mesh_posed = [mesh_posed[j, mask[j], :] for j in range(mask.shape[0])]
collision = [np.concatenate((body[i], mesh_posed[i]), axis=0) for i in range(len(mesh_posed))]
collision_faces = [np.concatenate((F, faces[i] + config['smpl_res']), axis=0) for i in range(len(mesh_posed))]
collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(mesh_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(mesh_posed))]
# body_2 = [b.copy() for b in body]
# for i in range(len(gt_posed)):
# body_2[i][:, 1] = body[i][:,2].copy()
# body_2[i][:, 2] = body[i][:,1].copy()
# gt_posed_with_body = [np.concatenate((body_2[i], gt_posed[i]), axis=0) for i in range(len(gt_posed))]
if config['is_cape']:
collision = mesh_posed
collision_faces = faces
collision_color = [np.ones_like(mesh_posed[i], dtype=np.int32)*230 for i in range(len(mesh_posed))]
else:
collision = [np.concatenate((body[i], mesh_posed[i]), axis=0) for i in range(len(mesh_posed))]
collision_faces = [np.concatenate((F, faces[i] + config['smpl_res']), axis=0) for i in range(len(mesh_posed))]
collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(mesh_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(mesh_posed))]
figs = { type + "/garment_" + str(i): plot_3d(gt[i], mesh_unposed_shaped_levels[i], faces[i], gt_posed[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
return { **figs, **{ type + '/images': imgs, } }
if config['is_cape']:
figs_coll = dict()
else:
garment_posed_fixed = [mesh_posed_fixed[i, mask[i], :] for i in range(body.shape[0])]
collision_fixed = [np.concatenate((body[i], garment_posed_fixed[i]), axis=0) for i in range(len(mesh_posed))]
gt_posed = [batch['mesh_posed'].numpy()[i, mask[i], :] for i in range(mask.shape[0])]
collision_gt = [np.concatenate((body[i], gt_posed[i]), axis=0) for i in range(len(gt_posed))]
figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i]) for i in range(4) }
return { **figs, **figs_coll, **{ type + '/images': imgs, } }
def sample(folder):
@@ -622,7 +672,6 @@ def test():
if __name__ == '__main__':
init_wandb()
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
@@ -634,10 +683,10 @@ if __name__ == '__main__':
dataset_resolutions.append(512)
factory = DatasetFactory()
dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
model: Pyramid = Pyramid(config)
@@ -654,11 +703,18 @@ if __name__ == '__main__':
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=20)
# load_model(model, config['pyramid']['pretrained_pyramid'])
model.static_enc.trainable = False
model.body_enc.trainable = False
model.curv_enc.trainable = False
# model.cond_encoder.trainable = False
# load_model(model, 'checkpoints/pyr_abl_32_skirt', iteration='180000')
if len(config['pyramid']['resolutions']) == 1:
model.static_enc.trainable = False
model.body_enc.trainable = False
model.curv_enc.trainable = False
else:
load_model(model, config['pyramid']['pretrained_pyramid'])
model.static_enc.trainable = False
model.body_enc.trainable = False
model.curv_enc.trainable = False
model.cond_encoder.trainable = False
# load_model(model.levels[0].upscale_enhancer, config['pyramid']['levels'][0]['upscale_enhancer'])
@@ -669,14 +725,15 @@ if __name__ == '__main__':
for i in range(len(model.levels)-1):
load_model(model.levels[i].upscale_enhancer, config['pyramid']['levels'][i]['upscale_enhancer'])
model.levels[i].upscale_enhancer.trainable = False
load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
model.upscale_enhancer.trainable = False
if config['pyramid']['resolutions'][-1] != 512:
load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
model.upscale_enhancer.trainable = False
for i in range(len(model.levels)-1):
model.levels[i].trainable = False
# load_model(model.cond_encoder, 'checkpoints/cvae_low_4_3_d32')
# load_model(model.levels[0], 'checkpoints/cvae_low_4_3_d32')
# load_model(model.cond_encoder, config['pyramid']['pretrained_pyramid'])
# load_model(model.levels[0], config['pyramid']['pretrained_pyramid'])
# load_model(model.levels[-1], 'checkpoints/cvae_19')
@@ -701,10 +758,10 @@ if __name__ == '__main__':
if step == 0:
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
if train_loss['train/euclid_error'] < 0.01:
if train_loss['train/euclid_error'] < 0.004:
if balancing_activated_step == 0:
balancing_activated_step = step
balancing_coeff = balancing_coeff + 1 / 15000
balancing_coeff = balancing_coeff + 1 / 20000
to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } }
dataset.py
+26 -13
View File
@@ -160,6 +160,8 @@ class DatasetFactory():
for key in signature.keys():
if key == 'id':
tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
elif key == 'frame_nr':
tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
elif key == 'mesh_curvature':
tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.int32))
else:
@@ -209,14 +211,13 @@ class DatasetFactory():
if 'id' in signature.keys():
item['id'] = id
if 'frame_nr' in signature.keys():
item['frame_nr'] = frame_nr
if 'pose_shape' in signature.keys():
pose_shape = np.load(os.path.join(folder, 'pose_shape', id + separator + frame_nr + '.npy'))
if use_no_orient:
pose_shape[:3] = 0
item['pose_shape'] = pose_shape
if 'uv_posed' in signature.keys():
uv_posed = np.load(os.path.join(folder, 'uv_posed', id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']
item['uv_posed'] = uv_posed
if 'mesh_body_posed' in signature.keys():
if use_no_orient:
mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + separator + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
@@ -225,8 +226,12 @@ class DatasetFactory():
mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
item['mesh_body_posed'] = mesh_body_posed
if 'mesh_posed' in signature.keys():
mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + separator + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed']
item['mesh_posed'] = mesh_posed
if use_no_orient:
mesh_posed = np.load(os.path.join(folder, 'mesh_posed_no_orient', id + separator + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed_no_orient'].astype(np.float32)
item['mesh_posed'] = mesh_posed
else:
mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + separator + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed'].astype(np.float32)
item['mesh_posed'] = mesh_posed
# if 'mesh_unposed' in signature.keys():
# mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
# item['mesh_unposed'] = mesh_unposed
@@ -269,7 +274,17 @@ class DatasetFactory():
item['uv_body_posed' + res_string] = uv_posed_body_no_orient
else:
uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']
item['uv_body_posed'] = uv_body_posed
item['uv_body_posed' + res_string] = uv_body_posed
if 'uv_posed' + res_string in signature.keys():
if use_no_orient:
uv_posed = np.load(os.path.join(folder, 'uv_posed_no_orient' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient']
item['uv_posed' + res_string] = uv_posed
else:
uv_posed = np.load(os.path.join(folder, 'uv_posed' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']
item['uv_posed' + res_string] = uv_posed
# print('here')
# item = { **static_info, **item }
as_list = []
@@ -330,7 +345,7 @@ if __name__ == '__main__':
# print(d)
factory = DatasetFactory()
signature_dynamic = { 'id': (),
signature_dynamic = { 'id': (), 'frame_nr': (),
# 'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
# 'uv_body_posed': (512, 256, 3),
@@ -351,16 +366,14 @@ if __name__ == '__main__':
# 'uv_shaped_body': (3, 512, 256,),
}
dataset = factory.get_dataset(folder='/home/workspace/local-data/train',
dataset = factory.get_dataset(folder='/home/workspace/local-data/train_skirt',
batch_size=300,
signature=signature_dynamic,
use_filled=False,
resolutions=[512],
do_repeat=False,
use_one_hot=True)
static_dataset = StaticDataset(folder='/home/workspace/local-data/train',
use_one_hot=False)
static_dataset = StaticDataset(folder='/home/workspace/local-data/train_skirt',
signature=signature_static,
use_filled=False,
resolutions=[512])
@@ -393,7 +406,7 @@ if __name__ == '__main__':
k = F[:,2]
)]).show()
np.save('custom_data/train_no_shape_one_hot_avg_uv.npy', avg_garment)
np.save('custom_data/train_no_shape_skirt_avg_uv.npy', avg_garment)
# print()
deorient.py
+104
View File
@@ -0,0 +1,104 @@
import os
import sys
os.environ["CUDA_VISIBLE_DEVICES"] = '6'
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 glob
import tqdm
import copy
import numpy as np
from custom_utils.smpl import SMPL
# from custom_utils.custom_utils import UV2Mesh, _maskFaces
import plotly.graph_objects as go
# import plotly.express as ex
from scipy.spatial.transform import Rotation as R
if __name__ == '__main__':
folder = '/home/workspace/local-data/test_skirt_ready'
# folder = '/home/workspace/data/dataset/test_ready'
# uv2mesh = UV2Mesh(is_skirt=False)
smpl = SMPL(use_one_hot=False)
F = np.load("custom_utils/smpl_faces.npy")
ids = os.listdir(os.path.join(folder, 'pose_shape'))
ids = [id.split('_')[0] for id in ids]
ids = set(ids)
os.makedirs(os.path.join(folder, 'mesh_posed_no_orient'), exist_ok=True)
for id in tqdm.tqdm(ids):
frames = glob.glob(os.path.join(folder, 'pose_shape', id + '_*'))
frames = len(frames)
mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')).astype(bool)
mesh_posed = [np.load(os.path.join(folder, 'mesh_posed', id + '_' + str(i) + '.npz')) for i in range(frames)]
mesh_posed = [mesh_posed[i]['mesh_posed'] for i in range(frames)]
# mesh_unposed = [np.load(os.path.join(folder, 'mesh_unposed', id + '_' + str(i) + '.npz')) for i in range(frames)]
# mesh_unposed = [mesh_unposed[i]['mesh_unposed'] for i in range(frames)]
pose_shape = [np.load(os.path.join(folder, 'pose_shape', id + '_' + str(i) + '.npy')).astype(np.float32) for i in range(frames)]
# pose_shape_no_orient = copy.deepcopy(pose_shape)
# for i in range(len(pose_shape)):
# pose_shape_no_orient[i][:3] = 0
rodr = smpl.run_rodrigues(np.stack([x[:72] for x in pose_shape]).astype(np.float32)).numpy()
for frame in range(frames):
r = R.from_matrix(rodr[frame, 0, :, :])
# assert r.as_matrix() == rodr[frame, 0, :, :]
out = r.apply(mesh_posed[frame], inverse=True)
# x = smpl.run_garment(pose_shape_no_orient[frame][:72][np.newaxis,:], pose_shape_no_orient[frame][72:][np.newaxis,:], mesh_unposed[frame][np.newaxis,:,:])[0].numpy()[mask, :]
# y = out[mask, :]
# _, 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 = faces[:,0],
# j = faces[:,1],
# k = faces[:,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()
np.savez_compressed(os.path.join(folder, 'mesh_posed_no_orient', id + '_' + str(frame) + '.npz'), mesh_posed_no_orient=out)
evaluate_ae.py
+152
View File
@@ -0,0 +1,152 @@
import os
import sys
from configs.config_abl_skin_skirt import config
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
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 glob
import time
import argparse
import wandb
import tqdm
import random
import numpy as np
import plotly.graph_objects as go
import plotly.express as ex
from PIL import Image
from postprocessing import fix_collision
if config['is_skirt']:
from losses import edge_loss_skirt as edge_loss
else:
from losses import edge_loss
from losses import collision_ratio, kl_loss #, concept_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
from dataset import DatasetFactory, StaticDataset
from custom_utils.smpl import SMPL
from models.cvae_model import CVAE
from custom_utils.supersmpl_io import writeOBJ
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False)
@tf.function
def step(batch):
if config['use_skinning']:
uv_orig = batch['uv_unposed']
else:
uv_orig = batch['uv_posed']
pred, _, _ = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
mesh_pred = uv2mesh(pred)
if config['use_skinning']:
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
else:
mesh_pred_posed = mesh_pred
# mesh_pred_posed_fixed = fix_collision(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
return mesh_pred_posed, None #mesh_pred_posed_fixed
if __name__ == '__main__':
# init_wandb()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
smpl = SMPL(use_one_hot=config['use_one_hot'])
skirt_modifier = '_skirt' if config['is_skirt'] else ''
folder = config['test_dir']
smpl = SMPL(use_one_hot=False)
model = CVAE(config)
# str_split = config['name'].split('_')
# test_idx = str_split.index('test')
# checkpoint_name = '_'.join(str_split[:test_idx])
checkpoint_name = config['name']
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
ids = os.listdir(os.path.join(folder, 'pose_shape'))
ids = [id.split('_')[0] for id in ids]
ids = set(ids)
iteration = '70000'
load_model(model, 'checkpoints/' + checkpoint_name, iteration=iteration)
# load_model(model, 'checkpoints/' + 'abl_all_120k')
output_folder = '/home/workspace/data/results/' + checkpoint_name + '-' + iteration
os.makedirs(output_folder, exist_ok=True)
if config['use_skinning']:
config['signature_dynamic'].pop('mesh_posed')
config['signature_dynamic'].pop('mesh_unposed')
config['signature_dynamic'].pop('mesh_body_posed')
factory = DatasetFactory()
dataset_test = factory.get_dataset(folder=config['test_dir'],
batch_size=config['batch_size'],
signature=config['signature_dynamic'],
resolutions=[config['resolution']],
in_memory=False,
do_repeat=False,
use_one_hot=config['use_one_hot'],
use_no_orient=config['remove_global_orientation'])
static_dataset_test = StaticDataset(folder=config['test_dir'],
signature=config['signature_static'],
resolutions=[config['resolution']],
in_memory=True)
# faces_folder = '/home/workspace/data/results/support/faces' + ('_skirt' if config['is_skirt'] else '')
# ids = os.listdir(faces_folder)
# ids = [id.split('.')[0] for id in ids]
# faces = {id : np.load(faces_folder + '/' + id + '.npy') for id in ids }
for batch in tqdm.tqdm(dataset_test):
batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
mesh_pred_posed, mesh_pred_posed_fixed = step(batch)
# mesh_pred_posed, mesh_pred_posed_fixed = mesh_pred_posed.numpy(), mesh_pred_posed_fixed.numpy()
mesh_pred_posed = mesh_pred_posed.numpy()
mask = batch['mesh_mask'].numpy().astype(bool)
for i in range(mask.shape[0]):
id, frame_nr = batch['id'][i].numpy().decode("utf-8"), batch['frame_nr'][i].numpy().decode("utf-8")
# writeOBJ(os.path.join(output_folder, id + '_' + frame_nr + '.obj'), mesh_pred_posed[i, mask[i], :], faces[id])
np.save(os.path.join(output_folder, id + '_' + frame_nr + '.npy'), mesh_pred_posed[i, mask[i], :])
# np.save(os.path.join(output_folder, id + '_' + frame_nr + '_fixed_coll.npy'), mesh_pred_posed_fixed[i, mask[i], :])
# writeOBJ(os.path.join(output_folder, id + '_' + frame_nr + '_fixed_coll.obj'), mesh_pred_posed_fixed[i, mask[i], :], faces[id])
evaluate_ae_pyramid.py
+186
View File
@@ -0,0 +1,186 @@
import os
import sys
from configs.config_abl_pyr_32 import config
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
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 glob
import time
import argparse
import wandb
import tqdm
import random
import numpy as np
import plotly.graph_objects as go
import plotly.express as ex
from PIL import Image
from models.pyramid_model_optimized import Pyramid
from postprocessing import fix_collision
if config['is_skirt']:
from losses import edge_loss_skirt as edge_loss
else:
from losses import edge_loss
from losses import collision_ratio, kl_loss #, concept_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision, group_inputs_for_levels
from dataset import DatasetFactory, StaticDataset
from custom_utils.smpl import SMPL
from models.cvae_model import CVAE
from custom_utils.supersmpl_io import writeOBJ
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False)
@tf.function
def step(batch):
levels_input, conds = group_inputs_for_levels(batch, config)
out_levels = model([levels_input, conds], training=True)
pred_final = out_levels[-1][0]
mesh_pred = uv2mesh(pred_final)
if config['use_skinning']:
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
else:
mesh_pred_posed = mesh_pred
# mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
# mesh_pred_posed_fixed = fix_collision(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
return mesh_pred_posed, None
if __name__ == '__main__':
# init_wandb()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
smpl = SMPL()
skirt_modifier = '_skirt' if config['is_skirt'] else ''
folder = config['test_dir']
smpl = SMPL(use_one_hot=False)
model: Pyramid = Pyramid(config)
# str_split = config['name'].split('_')
# test_idx = str_split.index('test')
# checkpoint_name = '_'.join(str_split[:test_idx])
checkpoint_name = config['name']
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
ids = os.listdir(os.path.join(folder, 'pose_shape'))
ids = [id.split('_')[0] for id in ids]
ids = set(ids)
iteration = '430000'
load_model(model, 'checkpoints/' + checkpoint_name, iteration=iteration)
# load_model(model, 'checkpoints/' + 'abl_all_120k')
output_folder = '/home/workspace/data/results/' + checkpoint_name + '-' + iteration
os.makedirs(output_folder, exist_ok=True)
dataset_resolutions = config['pyramid']['resolutions'].copy()
if 512 not in dataset_resolutions:
dataset_resolutions.append(512)
if config['use_skinning']:
config['signature_dynamic'].pop('mesh_posed')
config['signature_dynamic'].pop('mesh_unposed')
config['signature_dynamic'].pop('mesh_body_posed')
factory = DatasetFactory()
dataset_test = factory.get_dataset(folder=config['test_dir'],
batch_size=config['batch_size'],
signature=config['signature_dynamic'],
resolutions=dataset_resolutions,
in_memory=False,
do_repeat=False,
use_no_orient=True)
static_dataset_test = StaticDataset(folder=config['test_dir'],
signature=config['signature_static'],
resolutions=dataset_resolutions,
in_memory=True)
faces_folder = '/home/workspace/data/results/support/faces' + ('_skirt' if config['is_skirt'] else '')
ids = os.listdir(faces_folder)
ids = [id.split('.')[0] for id in ids]
faces = {id : np.load(faces_folder + '/' + id + '.npy') for id in ids }
for batch in tqdm.tqdm(dataset_test):
batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
mesh_pred_posed, _ = step(batch)
mesh_pred_posed = mesh_pred_posed.numpy()
mask = batch['mesh_mask'].numpy().astype(bool)
for i in range(mask.shape[0]):
id, frame_nr = batch['id'][i].numpy().decode("utf-8"), batch['frame_nr'][i].numpy().decode("utf-8")
# writeOBJ(os.path.join(output_folder, id + '_' + frame_nr + '.obj'), mesh_pred_posed[i, mask[i], :], faces[id])
np.save(os.path.join(output_folder, id + '_' + frame_nr + '.npy'), mesh_pred_posed[i, mask[i], :])
# np.save(os.path.join(output_folder, id + '_' + frame_nr + '_fixed_coll.npy'), mesh_pred_posed_fixed[i, mask[i], :])
# writeOBJ(os.path.join(output_folder, id + '_' + frame_nr + '_fixed_coll.obj'), mesh_pred_posed_fixed[i, mask[i], :], faces[id])
# x = batch['mesh_posed'].numpy()[i, mask[i], :]
# y = mesh_pred_posed[i, mask[i], :]
# fig = go.Figure(data=[
# go.Mesh3d(
# x=x[:,0],
# y=x[:,1],
# z=x[:,2],
# i = faces[id][:,0],
# j = faces[id][:,1],
# k = faces[id][:,2],
# ),
# go.Mesh3d(
# x=y[:,0],
# y=y[:,1],
# z=y[:,2],
# i = faces[id][:,0],
# j = faces[id][:,1],
# k = faces[id][:,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()
evaluate_ae_support.py
+110
View File
@@ -0,0 +1,110 @@
import os
import sys
from configs.config_abl_all import config
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
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 glob
import time
import argparse
import wandb
import tqdm
import random
import numpy as np
import plotly.graph_objects as go
import plotly.express as ex
from PIL import Image
from postprocessing import fix_collision
if config['is_skirt']:
from losses import edge_loss_skirt as edge_loss
else:
from losses import edge_loss
from losses import collision_ratio, kl_loss #, concept_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
from dataset import DatasetFactory, StaticDataset
from custom_utils.smpl import SMPL
from models.cvae_model import CVAE
from custom_utils.supersmpl_io import writeOBJ
if __name__ == '__main__':
# init_wandb()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
smpl = SMPL(use_one_hot=config['use_one_hot'])
skirt_modifier = '_skirt' if config['is_skirt'] else ''
folder = config['test_dir']
# str_split = config['name'].split('_')
# test_idx = str_split.index('test')
# checkpoint_name = '_'.join(str_split[:test_idx])
checkpoint_name = config['name']
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
output_folder = '/home/workspace/data/results/support'
os.makedirs(output_folder + '/bodies', exist_ok=True)
os.makedirs(output_folder + '/faces', exist_ok=True)
config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
'mesh_body_posed': (14475, 3,),
}
config['signature_static'] = {
'mesh_mask': (14475,),
}
factory = DatasetFactory()
dataset_test = factory.get_dataset(folder=config['test_dir'],
batch_size=config['batch_size'],
signature=config['signature_dynamic'],
resolutions=[config['resolution']],
in_memory=False,
do_repeat=False,
use_one_hot=config['use_one_hot'],
use_no_orient=config['remove_global_orientation'])
static_dataset_test = StaticDataset(folder=config['test_dir'],
signature=config['signature_static'],
resolutions=[config['resolution']],
in_memory=False)
for batch in tqdm.tqdm(dataset_test):
batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
# mask = batch['mesh_mask'].numpy().astype(bool)
body = batch['mesh_body_posed'].numpy()
# _, faces = maskFacesBatch(np.zeros((body.shape[0],14475,3)), F, mask)
for i in range(body.shape[0]):
id, frame_nr = batch['id'][i].numpy().decode("utf-8"), batch['frame_nr'][i].numpy().decode("utf-8")
writeOBJ(os.path.join(output_folder, 'bodies', id + '_' + frame_nr + '_body_no_orient.obj'), body[i], F)
# np.save(os.path.join(output_folder, 'faces', id + '_' + frame_nr + '.npy'), faces[i].astype(np.int32))
evaluate_cape.py
+54 -56
View File
@@ -1,5 +1,5 @@
import os
from configs.config_pyramid import config
from configs.config_abl_pyr_64_128_256_cape_f import config
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
import tensorflow as tf
@@ -44,7 +44,7 @@ from models.pyramid_model_optimized import Pyramid
def init_wandb():
run = wandb.init(project='cvae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp')
run = wandb.init(project='results', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('cvae_pyramid.py')
@@ -220,45 +220,46 @@ def test_step(batch):
levels_input, conds = group_inputs_for_levels(batch)
out_levels = model([levels_input, conds], training=False)
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
mesh_orig = batch['mesh_unposed']
# shape = batch['pose_shape'][:, 72:]
# gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
# shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
# mesh_orig = batch['mesh_unposed']
output = dict()
reconstr_loss_total = 0
kl_total = 0
# output = dict()
# reconstr_loss_total = 0
# kl_total = 0
#for i, res in enumerate(config['pyramid']['resolutions']):
i = -1
res = config['pyramid']['resolutions'][-1]
res_modifier = '_' + str(res) if res != 512 else ''
pred, mean, logvar = out_levels[i]
# #for i, res in enumerate(config['pyramid']['resolutions']):
# i = -1
# res = config['pyramid']['resolutions'][-1]
# res_modifier = '_' + str(res) if res != 512 else ''
# pred, mean, logvar = out_levels[i]
mesh_pred = uv2mesh(pred)
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
# mesh_pred = uv2mesh(pred)
# mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
# mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
# is the output already masked?
# # is the output already masked?
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
# e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
# c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
# c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
# l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
# # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
# ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
# reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
# metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
# # metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
output = { **output, **{ 'test/edge_loss' + res_modifier: e_loss, 'test/collision_loss' + res_modifier: c_loss,
'test/ms_ssim' + res_modifier: ms_ssim, 'test/l2_loss' + res_modifier: l2_loss,
'test/collision_ratio' + res_modifier: c_ratio } }
# output = { **output, **{ 'test/edge_loss' + res_modifier: e_loss, 'test/collision_loss' + res_modifier: c_loss,
# 'test/ms_ssim' + res_modifier: ms_ssim, 'test/l2_loss' + res_modifier: l2_loss,
# 'test/collision_ratio' + res_modifier: c_ratio } }
loss = reconstr_loss_total + balancing_coeff * kl_total
# loss = reconstr_loss_total + balancing_coeff * kl_total
return { **{ 'test/loss': loss, 'test/reconstr_loss_total': reconstr_loss_total }, **output }, out_levels
# return { **{ 'test/loss': loss, 'test/reconstr_loss_total': reconstr_loss_total }, **output }, out_levels
return dict(), out_levels
@@ -322,13 +323,10 @@ def plot_everything(batch: dict, type: str):
def test():
str_split = config['name'].split('_')
test_idx = str_split.index('test')
checkpoint_name = '_'.join(str_split[:test_idx])
load_model(model, 'checkpoints/' + checkpoint_name)
load_model(model, 'checkpoints/' + config['name'])
dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=False, do_repeat=False, is_cape=config['is_cape'])
static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=False, is_cape=config['is_cape'])
dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], resolutions=dataset_resolutions, in_memory=False, do_repeat=False, is_cape=config['is_cape'])
static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], resolutions=dataset_resolutions, in_memory=True, is_cape=config['is_cape'])
step = 0
losses = []
@@ -337,7 +335,7 @@ def test():
for batch in tqdm.tqdm(dataset_test):
batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
test_loss, out_levels = test_step(batch)
_, out_levels = test_step(batch)
all_errors = []
errors = []
@@ -350,15 +348,15 @@ def test():
error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((offsets_orig - offsets_pred) ** 2, axis=-1)))
errors.append(error)
test_loss = { **{ 'test/error': tf.reduce_mean(errors) }, **test_loss }
losses.append(test_loss)
# test_loss = { **{ 'test/error': tf.reduce_mean(errors) }, **test_loss }
# losses.append(test_loss)
all_errors.extend(errors)
to_log = { **{ 'current_step': step, 'test/error': np.asarray(errors).mean() }, **test_loss }
if step % 1000 == 0:
figs = plot_everything(batch, 'test')
to_log = { **to_log, **figs }
# to_log = { **{ 'current_step': step, 'test/error': np.asarray(errors).mean() }, **test_loss }
# if step % 1000 == 0:
# figs = plot_everything(batch, 'test')
# to_log = { **to_log, **figs }
wandb.log(to_log, step = step)
# wandb.log(to_log, step = step)
# print(step)
step += 1
@@ -371,25 +369,25 @@ def test():
print('min: ' + str(errors.min()))
print('max: ' + str(errors.max()))
print('median: ' + str(np.median(errors)))
np.save('temp/cape_f_pyr_64_128_256_error.npy', errors)
np.save('temp/error_'+ config['name'] +'.npy', errors)
histograms = dict()
for key in test_loss.keys():
accumulated = [losses[i][key].numpy() for i in range(len(losses))]
fig = ex.histogram(accumulated, title='Average: ' + str(np.asarray(accumulated).mean()))
fig.write_html(path_to_plotly_html, auto_play = False)
# histograms = dict()
# for key in test_loss.keys():
# accumulated = [losses[i][key].numpy() for i in range(len(losses))]
# fig = ex.histogram(accumulated, title='Average: ' + str(np.asarray(accumulated).mean()))
# fig.write_html(path_to_plotly_html, auto_play = False)
histograms['all/' + key.split('/')[-1] + '_hist'] = wandb.Html(path_to_plotly_html)
# averages['all/' + key.split('/')[-1] + '_avg'] = np.asarray(accumulated).mean()
# histograms['all/' + key.split('/')[-1] + '_hist'] = wandb.Html(path_to_plotly_html)
# # averages['all/' + key.split('/')[-1] + '_avg'] = np.asarray(accumulated).mean()
wandb.log({ **histograms }, step = step)
wandb.finish()
# wandb.log({ **histograms }, step = step)
# wandb.finish()
pass
if __name__ == '__main__':
init_wandb()
# init_wandb()
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
get_error.py
+60
View File
@@ -0,0 +1,60 @@
import os
import csv
import numpy as np
if __name__ == '__main__':
garment_type_table = np.load('custom_data/garment_type_table_test.npy', allow_pickle=True).item()
folder = '/home/workspace/data/results/errors'
checkpoint_name = 'abl_all'
iteration = '70000'
errors = dict()
errors_skirt = dict()
per_garment_errors = { 'Tshirt': [], 'Dress': [], 'Skirt': [], 'Top': [] }
per_garment_errors_skirt = { 'Tshirt': [], 'Jumpsuit': [], 'Trousers': [], 'Top': [] }
with open(os.path.join(folder, checkpoint_name + '-' + iteration + '.csv'), 'r') as f:
csv_reader = csv.reader(f, delimiter=',')
for row in csv_reader:
per_garment_errors[garment_type_table[row[0].split('_')[0]]].append(row[1])
with open(os.path.join(folder, checkpoint_name + '_skirt' + '-' + iteration + '.csv'), 'r') as f:
csv_reader = csv.reader(os.path.join(folder, checkpoint_name + '_skirt' + '-' + iteration + '.csv'), delimiter=',')
for row in csv_reader:
per_garment_errors_skirt[garment_type_table[row[0]]].append(row[1])
print('Non-skirt:')
for key in per_garment_errors.keys():
avg = np.mean(per_garment_errors[key])
errors.append(avg)
print(key + ': ' + str(avg))
print('Skirt:')
for key in per_garment_errors_skirt.keys():
avg = np.mean(per_garment_errors_skirt[key])
errors_skirt.append(avg)
print(key + ': ' + str(avg))
print('Avg non-skirt: ' + str(np.mean(errors)))
print('Avg skirt: ' + str(np.mean(errors_skirt)))
cum_avg_shirt = np.mean(per_garment_errors['Tshirt'] + per_garment_errors_skirt['Tshirt'])
cum_avg_top = np.mean(per_garment_errors['Top'] + per_garment_errors_skirt['Top'])
print('Global tshirt avg: ' + str(cum_avg_shirt))
print('Global tshirt avg: ' + str(cum_avg_top))
print('Globl per garment: ')
global_errors = []
per_garment_errors['Tshirt'].extend(per_garment_errors_skirt['Tshirt'])
per_garment_errors['Top'].extend(per_garment_errors_skirt['Top'])
for key in per_garment_errors.keys():
avg = np.mean(per_garment_errors[key])
global_errors.append(avg)
print(key + ': ' + str(avg))
print('Global avg: ' + str(np.mean(global_errors)))
losses.py
+14 -100
View File
@@ -25,107 +25,14 @@ 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 custom_utils.custom_utils import maskFacesBatch, _maskFaces
from custom_utils.graph_utils import sparse_to_tensor
def d_logistic(real_images, generator, discriminator, z_dim, conditions):
batch_size = tf.shape(real_images)[0]
z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
# forward pass
fake_images = generator([z, conditions], training=True)
real_scores = discriminator(real_images, training=True)
fake_scores = discriminator(fake_images, training=True)
# gan loss
d_loss = tf.math.softplus(fake_scores)
d_loss += tf.math.softplus(-real_scores)
return d_loss
def d_logistic_r1_reg(real_images, generator, discriminator, z_dim, conditions):
batch_size = tf.shape(real_images)[0]
z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
# forward pass
fake_images = generator([z, conditions], training=True)
real_scores = discriminator(real_images, training=True)
fake_scores = discriminator(fake_images, training=True)
# gan loss
d_loss = tf.math.softplus(fake_scores)
d_loss = d_loss + tf.math.softplus(-real_scores)
# gradient penalty
with tf.GradientTape() as r1_tape:
r1_tape.watch([real_images, conditions])
real_scores = discriminator(real_images, training=True)
real_loss = tf.reduce_sum(real_scores)
real_grads = r1_tape.gradient(real_loss, real_images)
r1_penalty = tf.reduce_sum(tf.math.square(real_grads), axis=[1, 2, 3])
r1_penalty = tf.expand_dims(r1_penalty, axis=1)
return d_loss, r1_penalty
def g_logistic_non_saturating(real_images, generator, discriminator, z_dim, conditions):
batch_size = tf.shape(real_images)[0]
z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
# forward pass
fake_images = generator([z, conditions], training=True)
fake_scores = discriminator(fake_images, training=True)
# gan loss
g_loss = tf.math.softplus(-fake_scores)
return fake_images, g_loss
def g_logistic_ns_pathreg(real_images, generator, discriminator, z_dim,
pl_mean, pl_minibatch_shrink, pl_denorm, pl_decay,
conditions):
batch_size = tf.shape(real_images)[0]
z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
pl_minibatch = tf.maximum(1, tf.math.floordiv(batch_size, pl_minibatch_shrink))
pl_z = tf.random.normal(shape=[pl_minibatch, z_dim], dtype=tf.float32)
pl_conditions = { key: conditions[key][:pl_minibatch] for key in conditions.keys() }
# for key in conditions.keys():
# pl_conditions[key] = conditions[key][:pl_minibatch]
# pl_conditions = (conditions[0][:pl_minibatch], conditions[1][:pl_minibatch], conditions[2][:pl_minibatch], conditions[3][:pl_minibatch], conditions[4][:pl_minibatch])
# forward pass
fake_images = generator([z, conditions], ret_w_broadcasted=False, training=True)
fake_scores = discriminator(fake_images, training=True)
g_loss = tf.math.softplus(-fake_scores)
# Evaluate the regularization term using a smaller minibatch to conserve memory.
with tf.GradientTape() as pl_tape:
pl_tape.watch([pl_z, pl_conditions])
pl_fake_images, pl_w_broadcasted = generator([pl_z, pl_conditions], ret_w_broadcasted=True, training=True)
pl_noise = tf.random.normal(tf.shape(pl_fake_images)) * pl_denorm
pl_noise_applied = tf.reduce_sum(pl_fake_images * pl_noise)
pl_grads = pl_tape.gradient(pl_noise_applied, pl_w_broadcasted)
pl_lengths = tf.math.sqrt(tf.reduce_mean(tf.reduce_sum(tf.math.square(pl_grads), axis=2), axis=1))
# Track exponential moving average of |J*y|.
pl_mean_val = pl_mean + pl_decay * (tf.reduce_mean(pl_lengths) - pl_mean)
pl_mean.assign(pl_mean_val)
# Calculate (|J*y|-a)^2.
pl_penalty = tf.square(pl_lengths - pl_mean)
return fake_images, g_loss, pl_penalty
# from hbertiche/PBNS
""" COLLISION FUNCTIONS START """
F = np.load("custom_utils/smpl_faces.npy")
F_coll = np.load("custom_utils/smpl_faces.npy")
F_coll_skirt = np.load("custom_utils/smpl_faces_skirt.npy")
F_norm = np.tile(F_coll[np.newaxis,:,:], [4,1,1])
F_norm_skirt = np.tile(F_coll_skirt[np.newaxis,:,:], [4,1,1])
def _nearest_neighbour(V, B):
tree = cKDTree(B.numpy())
@@ -137,7 +44,7 @@ def _tf_nn_parallel(V, B):
# @tf.function
def collision_loss(V, B, mask, thr=.004):
B_vn = vertex_normals(B, tf.tile(F[None], [B.shape[0], 1, 1]))
B_vn = vertex_normals(B, tf.tile(F_coll[None], [B.shape[0], 1, 1]))
# compute correspondences
idx = _tf_nn_parallel(V, B)
@@ -149,9 +56,11 @@ def collision_loss(V, B, mask, thr=.004):
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.reduce_sum((tf.minimum(dot - thr, 0) * mask) ** 2, axis=-1) # / tf.reduce_sum(mask, axis=-1)
thresholded = tf.minimum(dot - thr, 0) * mask
loss = tf.reduce_sum(thresholded ** 2, axis=-1)
ratio = tf.reduce_sum(tf.cast(thresholded < 0, dtype=tf.float32)) / tf.reduce_sum(mask)
return loss
return loss, ratio
def collision_ratio(V, B, mask, thr=.004):
@@ -327,6 +236,11 @@ def kl_loss(mean, logvar):
return kl_loss
def normal_loss(pred, gt, mask, is_skirt=False):
pred_norm = vertex_normals(pred, F_norm_skirt if is_skirt else F_norm)
gt_norm = vertex_normals(gt, F_norm_skirt if is_skirt else F_norm)
return tf.reduce_mean(tf.reduce_mean((tf.reduce_sum(tf.math.abs(pred_norm - gt_norm), axis=-1) * mask), axis=-1))
# vgg = tf.keras.applications.vgg16.VGG16(include_top=False, input_shape=(512,256,3))
# vgg_block_1 = tf.keras.Sequential(vgg.layers[:4])
# vgg_block_2 = tf.keras.Sequential(vgg.layers[4:7])
models/cvae_model.py
+22 -14
View File
@@ -26,14 +26,16 @@ class CVAE(tf.keras.Model):
if config['body_model'] != '':
load_model(self.body_enc, config['body_model'] + '_enc')
self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
if config['curvature_model'] != '':
load_model(self.curv_enc, config['curvature_model'] + '_enc')
if config['use_curvature']:
self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
if config['curvature_model'] != '':
load_model(self.curv_enc, config['curvature_model'] + '_enc')
if config['freeze_encoders']:
self.static_enc.trainable = False
self.body_enc.trainable = False
self.curv_enc.trainable = False
if config['use_curvature']:
self.curv_enc.trainable = False
self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
tf.keras.layers.Dense(256, activation='relu'),
@@ -50,7 +52,13 @@ class CVAE(tf.keras.Model):
self.add_noise = config['cvae']['add_noise']
self.noise_amount = config['cvae']['noise_amount']
self.use_curvature = config['use_curvature']
self.use_offsets_over_body = not config['use_skinning']
if self.use_offsets_over_body:
assert self.avg_garment is None, 'When we are not using skinning, we can\'t use avg_garment!'
if self.avg_garment is not None:
assert not self.use_offsets_over_body, 'When avg_garment is given, we must use skinning!'
@tf.function
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None):
@@ -84,27 +92,24 @@ class CVAE(tf.keras.Model):
def encode_conds(self, uv_template, uv_body_posed, uv_curv, training=None):
z_body = self.body_enc(uv_body_posed, training=training)
z_static = self.static_enc(uv_template, training=training)
z_curv = self.curv_enc(uv_curv, training=training)
if self.use_curvature:
z_curv = self.curv_enc(uv_curv, training=training)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
else:
z_cond = tf.concat([z_static, z_body], axis=-1)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder(z_cond, training=training)
return z_cond
def call(self, inputs, training=None):
if self.size == 512:
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
uv_template_high, uv_body_posed_high, uv_curv_high = uv_template, uv_body_posed, uv_curv
else:
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_template_high, uv_body_posed_high, uv_curv_high, uv_mask_high = inputs
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
mean, logvar = self.encoder(x, training=training)
z = self.reparameterize(mean, logvar)
z_cond = self.encode_conds(uv_template_high, uv_body_posed_high, uv_curv_high, training=training)
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=training)
if self.add_noise and training == True:
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
@@ -116,6 +121,9 @@ class CVAE(tf.keras.Model):
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
if self.use_offsets_over_body:
x = x + uv_body_posed
x = x * uv_mask[:,:,:,tf.newaxis]
return x, mean, logvar
models/pyramid_model_optimized.py
+75 -31
View File
@@ -43,6 +43,28 @@ class Pyramid(tf.keras.Model):
self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae': config['cond_ae']} }) for i in range(len(config['pyramid']['levels']))]
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
# self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
# self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
if len(config['pyramid']['resolutions']) == 2:
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10)])
if len(config['pyramid']['resolutions']) == 3:
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10)])
if len(config['pyramid']['resolutions']) == 4:
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10)])
if len(config['pyramid']['resolutions']) == 5:
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10)])
self.add_noise = config['pyramid']['add_noise']
self.use_offsets = config['use_offsets']
self.upscale_enhancer = None
if config['pyramid']['final_upscale_enhancer'] != '':
self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1])
@@ -66,42 +88,64 @@ class Pyramid(tf.keras.Model):
def call(self, inputs, training=None):
levels_input, conds = inputs
uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
z_cond = self.encode_conds(uv_static, uv_body, uv_curv)
# out_levels = []
# pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
# if 1 != len(self.levels):
# pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
# pred = tf.image.resize(pred, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
# out_levels.append([pred, mean, logvar])
# pred = pred * levels_input[1][-2][:,:,:,tf.newaxis]
# else:
# pred = self.enhance(pred, levels_input[0][-2], training=training)
# pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
# out_levels.append([pred, mean, logvar])
if self.add_noise:
z_cond = z_cond + tf.random.uniform(z_cond.shape, maxval=10)
# for i in range(1, len(self.levels)):
# pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond, training=training)
# pred = pred + pred_level
# if i != len(self.levels) - 1:
# pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
# out_levels.append([pred, mean, logvar])
# pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.LANCZOS5)
# pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
# else:
# if self.upscale_enhancer is not None:
# pred = self.enhance(pred, levels_input[-1][-2], training=training)
# pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
# out_levels.append([pred, mean, logvar])
out_levels = []
pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
if self.use_offsets:
pred = pred + levels_input[0][2] # uv_body_posed_32
if 1 != len(self.levels):
pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
pred = tf.image.resize(pred, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
out_levels.append([pred, mean, logvar])
pred = pred * levels_input[1][-2][:,:,:,tf.newaxis]
else:
pred = self.enhance(pred, levels_input[0][-2], training=training)
pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
out_levels.append([pred, mean, logvar])
pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
pred_0 = self.enhance(pred_0, levels_input[0][-2], training=training)
pred_0 = tf.image.resize(pred_0, (512, 256), tf.image.ResizeMethod.LANCZOS5)
pred_0 = pred_0 * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
for i in range(1, len(self.levels)):
pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond, training=training)
pred = pred + pred_level
if i != len(self.levels) - 1:
pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
out_levels.append([pred, mean, logvar])
pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.LANCZOS5)
pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
else:
if self.upscale_enhancer is not None:
pred = self.enhance(pred, levels_input[-1][-2], training=training)
pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
out_levels.append([pred, mean, logvar])
return [ [pred_0, mean_0, logvar_0] ]
return out_levels
# # 32
# pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
# pred_0 = self.enhance(pred_0, levels_input[0][-2], training=training)
# pred_0 = tf.image.resize(pred_0, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# pred_0 = pred_0 * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
# return [ [pred_0, mean_0, logvar_0] ]
# # 32 64
# pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
# pred_0 = self.levels[0].enhance(pred_0, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
# pred_0 = tf.image.resize(pred_0, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
# pred = pred_0 * levels_input[1][-2][:,:,:,tf.newaxis]
# pred_1, mean_1, logvar_1 = self.levels[1](levels_input[1], z_cond, training=training)
# pred_1 = pred + pred_1
# pred_1 = self.enhance(pred_1, levels_input[-1][-2], training=training)
# pred_1 = tf.image.resize(pred_1, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# pred_1 = pred_1 * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
# return [ [pred_0, mean_0, logvar_0], [pred_1, mean_1, logvar_1] ]
# class Pyramid(tf.keras.Model):
verify_cape_skinning.py
+1 -1
View File
@@ -38,7 +38,7 @@ if __name__ == '__main__':
F = np.load("custom_utils/smpl_faces_orig.npy")
factory = DatasetFactory()
signature_dynamic = { 'id': (),
signature_dynamic = { 'id': (), 'frame_nr': (),
'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
# 'uv_body_posed': (3, 512, 256,),
verify_poses.py
+82 -83
View File
@@ -1,10 +1,27 @@
import os
import os
os.environ["CUDA_VISIBLE_DEVICES"] = '5'
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 time
import cv2
import random
import numpy as np
import tensorflow as tf
import cmapy as cm
import plotly.graph_objects as go
import plotly.express as ex
@@ -14,7 +31,7 @@ from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
from custom_utils.smpl import SMPL
if __name__ == '__main__':
os.environ["CUDA_VISIBLE_DEVICES"] = "6"
os.environ["CUDA_VISIBLE_DEVICES"] = "4"
uv2mesh = UV2Mesh()
smpl = SMPL()
@@ -22,7 +39,7 @@ if __name__ == '__main__':
F = np.load("custom_utils/smpl_faces.npy")
factory = DatasetFactory()
signature_dynamic = { 'id': (),
signature_dynamic = { 'id': (), 'frame_nr': (),
'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
# 'uv_body_posed': (3, 512, 256,),
@@ -41,11 +58,13 @@ if __name__ == '__main__':
# 'uv_shaped_body': (3, 512, 256,),
}
dataset = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_test/train',
dataset = factory.get_dataset(folder='/home/workspace/local-data/train',
batch_size=1,
signature=signature_dynamic)
static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_test/train',
signature=signature_static)
signature=signature_dynamic,
resolutions=[512])
static_dataset = StaticDataset(folder='/home/workspace/local-data/train',
signature=signature_static,
resolutions=[512])
uv_res = 256
@@ -59,102 +78,82 @@ if __name__ == '__main__':
mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
pose, shape = batch['pose_shape'][:, :72].numpy(), batch['pose_shape'][:, 72:].numpy()
mesh_reposed = smpl.run_garment(pose, shape, batch['mesh_unposed']).numpy()[0]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]).numpy()
mesh_unposed_shaped = batch['mesh_unposed'] + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_reposed = smpl.run_garment(pose, shape, mesh_unposed_shaped).numpy()[0]
mesh_posed = np.squeeze(batch['mesh_posed'].numpy())
mesh_body_posed = np.squeeze(batch['mesh_body_posed'].numpy())
mesh_posed_body_no_orient = np.squeeze(batch['mesh_posed_body_no_orient'].numpy())
# x = mesh_posed[mask, :]
R = smpl.run_rodrigues(pose)
# shape[:, -1] = 1
x = smpl.run(R, shape)[0][0]
# x = smpl.run(smpl.run_rodrigues(pose), shape)[0][0]
x = mesh_reposed[mask, :]
y = mesh_body_posed
z = mesh_posed[mask, :]
_, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=x[:,0],
y=x[:,1],
z=x[:,2],
# i, j and k give the vertices of triangles
i = F[:,0],
j = F[:,1],
k = F[:,2]
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),
go.Mesh3d(
# 8 vertices of a cube
x=y[:,0],
y=y[:,1],
z=y[:,2],
# i, j and k give the vertices of triangles
i = F[:,0],
j = F[:,1],
k = F[:,2]
),
# go.Mesh3d(
# # 8 vertices of a cube
# x=z[:,0],
# y=z[:,1],
# z=z[:,2],
# # i, j and k give the vertices of triangles
# i = faces[:,0],
# j = faces[:,1],
# k = faces[:,2]
# ),
)
]).show()
pose_2 = pose.copy()
pose_2[:, :3] = 0
R = smpl.run_rodrigues(pose_2)
x = smpl.run(R, shape)[0][0]
y = mesh_posed_body_no_orient
_, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=x[:,0],
y=x[:,1],
z=x[:,2],
# i, j and k give the vertices of triangles
i = F[:,0],
j = F[:,1],
k = F[:,2]
),
go.Mesh3d(
# 8 vertices of a cube
x=y[:,0],
y=y[:,1],
z=y[:,2],
# i, j and k give the vertices of triangles
i = F[:,0],
j = F[:,1],
k = F[:,2]
)]).show()
# pose_2 = pose.copy()
# pose_2[:, :3] = 0
# # x = smpl.run(smpl.run_rodrigues(pose_2), shape)[0][0]
# x = mesh_unposed
# y = mesh_posed_body_no_orient
# _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
# go.Figure(data=[
# go.Mesh3d(
# # 8 vertices of a cube
# x=x[:,0],
# y=x[:,1],
# z=x[:,2],
# # i, j and k give the vertices of triangles
# i = F[:,0],
# j = F[:,1],
# k = F[:,2]
# ),
# go.Mesh3d(
# # 8 vertices of a cube
# x=y[:,0],
# y=y[:,1],
# z=y[:,2],
# i = F[:,0],
# j = F[:,1],
# k = F[:,2]
# )]).show()
x = mesh_posed[mask, :]
y = mesh_reposed[mask, :]
_, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=x[:,0],
y=x[:,1],
z=x[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),
go.Mesh3d(
# 8 vertices of a cube
x=y[:,0],
y=y[:,1],
z=y[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
)]).show()
# x = mesh_posed[mask, :]
# y = mesh_reposed[mask, :]
# _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
# 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 = faces[:,0],
# j = faces[:,1],
# k = faces[:,2]
# )]).show()
print()