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

large refactor

Browse Source
This commit is contained in:
luni
2022-09-25 08:31:13 +00:00
parent 84ec5a056e
commit 92891f5b26
33 changed files with 1661 additions and 1635 deletions
Download Patch File Download Diff File Expand all files Collapse all files
auto_encoder_body.py
+34 -125
View File
@@ -1,5 +1,6 @@
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
from configs.config_ae_body import config
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
import time
import argparse
@@ -8,26 +9,20 @@ import random
import cv2
import numpy as np
import tensorflow as tf
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
from custom_utils.custom_utils import UV2Mesh, load_model, maskFacesBatch
import plotly.graph_objects as go
import plotly.express as ex
from tf_utils import check_tf_version, allow_memory_growth
from losses import edge_loss, uv_mask_loss
from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL
from models.ae_body_model import AutoEncoder, AutoEncoder128, AutoEncoder32
from custom_utils.custom_utils import batch_to_dict, plot_3d
from models.cond_encoder import AutoEncoder
name='ae_body_8'
def init_wandb():
run = wandb.init(project='vae', name=name,
notes='convnext + convnext, add dense with tanh at the end of reverse convnext, larger features. add lr scheduler. fft3d instead of rfft3d. \
fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \
. lowered learning rate because it goes to nan. shallow third stage. slower lr decay. with validation. higher lr rate.\
two more downscaling than up, so decoder output is resized wto final resolution\
larger batch_size.\
run = wandb.init(project='vae', name=config['name'],
notes='convnext + convnext, add dense with tanh at the end of reverse convnext,\
flattened bottleneck with 128 latent_dim\ remove global orientation\
fixed gender issue with deoriented bodies, again, aaand again\
trained on 100 garments\
@@ -36,7 +31,10 @@ def init_wandb():
artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_body.py')
artifact.add_file('losses.py')
artifact.add_file('models/ae_body_model.py')
artifact.add_file('configs/config.py')
artifact.add_file('configs/config_ae_body.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')
run.log_artifact(artifact)
@@ -44,90 +42,24 @@ def init_wandb():
uv2mesh = UV2Mesh(res=256)
def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
mesh_pred_unposed[:,1] += 1
collision[:,1] += 2
fig = go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=mesh_gt[:,0],
y=mesh_gt[:,1],
z=mesh_gt[:,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=mesh_pred_unposed[:,0],
y=mesh_pred_unposed[:,1],
z=mesh_pred_unposed[:,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=collision[:,0],
y=collision[:,1],
z=collision[:,2],
# i, j and k give the vertices of triangles
i = collision_faces[:,0],
j = collision_faces[:,1],
k = collision_faces[:,2],
vertexcolor=collision_color
)
])
fig.update_layout(scene_aspectmode='manual',
scene_aspectratio=dict(x=1, y=2, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,3],),
zaxis = dict(nticks=4, range=[-1,1],),)
)
return fig
@tf.function
def train_step(batch):
with tf.GradientTape() as tape:
pred = model(batch['uv_body_posed' + resolution_modifier])
# pred = (pred + zero_body) * batch['mesh_mask'][:,:,tf.newaxis]
pred = model([batch['uv_body_posed' + resolution_modifier], uv_mask_body])
mesh_pred = uv2mesh(pred)
# mesh_pred = pred
mesh_orig = batch['mesh_body_posed']
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, tf.ones((pred.shape[0], 14475))))
l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_body_posed' + resolution_modifier] - pred))), axis=1))
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_body_posed' + resolution_modifier] + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, tf.ones_like(batch['mesh_mask'])))
c_loss = 0 # 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
# l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['mesh_unposed'] - pred)), axis=-1) * batch['mesh_mask'])
l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_body_posed' + resolution_modifier] - pred) * uv_mask_body[tf.newaxis,tf.newaxis,:,:])), axis=1))
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, tf.transpose(batch['uv_body_posed' + resolution_modifier] * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, max_val=2))
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, tf.transpose(batch['uv_body_posed' + resolution_modifier] * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
fft_loss = 0 # tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_body_posed'], tf.zeros_like(batch['uv_body_posed'])))), axis=[1,2,3]))
# loss = fft_loss_pc # + e_loss + c_loss
loss = 1e3 * (1 - ms_ssim)
mesh_orig_2 = uv2mesh(batch['uv_body_posed' + resolution_modifier])
# mesh_skinned = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred)
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_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
gradients = tape.gradient(loss, model.trainable_variables)
# gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]
# gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/euclid_error_posed': euclid_posed }
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, }
@@ -155,35 +87,27 @@ def valid_step(batch):
def plot_everything(batch: dict, type: str):
pred = model(batch['uv_body_posed' + resolution_modifier])
pred = model([batch['uv_body_posed' + resolution_modifier], uv_mask_body])
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_body_posed']
orig_img = tf.transpose(batch['uv_body_posed' + resolution_modifier], [0,2,3,1])
pred_img = tf.transpose(pred, [0,2,3,1])
orig_img = batch['uv_body_posed' + resolution_modifier]
pred_img = pred
img = tf.concat((orig_img, pred_img), axis=2)
imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
F = np.load("custom_utils/smpl_faces.npy")
mask = batch['mesh_mask'].numpy().astype(bool)
# garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
garment_unposed, faces = mesh_pred.numpy(), F
# gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
gt = [mesh_orig[i].numpy()[:, :] for i in range(pred.shape[0])]
body = uv2mesh(batch['uv_body_posed' + resolution_modifier]).numpy()
# body = batch['mesh_body_posed'].numpy()
# mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
# garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
garment_posed = garment_unposed
collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
# collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment_posed))]
collision_faces = [np.concatenate((F, faces + 14475), axis=0) for i in range(len(garment_posed))]
collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces, collision[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]) for i in range(4) }
return { **figs, **{ type + '/images': imgs } }
@@ -194,11 +118,11 @@ def train():
step = 0
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
train_loss = train_step(batch)
# data_valid = valid_iterator.get_next()
# data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
# data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
valid_loss = {}#valid_step(data_valid)
if step % 2500 == 0:
@@ -211,52 +135,37 @@ def train():
if step % 10000 == 0:
manager.save(checkpoint_number=step)
print(step)
step += 1
if __name__ == '__main__':
init_wandb()
# init_wandb()
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
smpl = SMPL()
uv_mask_body = np.load('custom_utils/mask_256x512.npy')
uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
# uv_mask_body = tf.ones((128, 64))
# TODO: different resolutions not working yet, model should be ready though
resolution_modifier = ''
resolutions = [512]
if config['resolution'] != 512:
resolution_modifier = '_' + str(config['resolution'])
factory = DatasetFactory()
signature_dynamic = { 'id': (),
'pose_shape': (72 + 11,),
'uv_body_posed': (3, 512, 256,),
'mesh_body_posed': (14475, 3,),
}
signature_static = {
'mesh_mask': (14475,),
# 'uv_mask': (512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=32, signature=signature_dynamic, use_filled=False, resolutions=resolutions)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=resolutions)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=32, signature=signature_dynamic)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']])
static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']])
dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'])
static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'])
model = AutoEncoder()
# model = AutoEncoder32()
# model = AutoEncoder128()
optimizer = tf.optimizers.Adam(learning_rate=0.0001)
folder = 'models/ae_logs/'
model = AutoEncoder(config['cond_ae'])
optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3)
ckpt.restore(manager.latest_checkpoint)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3)
train()