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

working pyramid pipeline

Browse Source
This commit is contained in:
luni
2022-09-27 21:26:15 +00:00
parent 1da83294ee
commit fa8f24f0c5
10 changed files with 997 additions and 213 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cvae_pyramid.py
+158 -198
View File
@@ -1,5 +1,6 @@
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
from configs.config_pyramid import config
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -23,178 +24,180 @@ import scipy as sp
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces
import plotly.graph_objects as go
import plotly.express as ex
from plotly.subplots import make_subplots
from PIL import Image
from utils import str_to_bool
from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
from load_models import load_generator, load_discriminator
from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss, bend_loss
from losses import edge_loss, uv_mask_loss, bend_loss, kl_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
from dataset import DatasetFactory, StaticDataset
from custom_utils.smpl import SMPL
from models.cvae_pyramid_model import Pyramid
from models.pyramid_model import Pyramid
name='cvae_pyr_4_d32'
def init_wandb():
run = wandb.init(project='vae', name=name,
notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. fixed balancing of reconst and kl loss for baseline. \
new implementation for kl loss. no lr schedule. 25 jumpsuits, average reconst_loss before calculation, linearly increasing balancing, up to 1 from 0.1. \
offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \
retrained body encoder with deoriented bodies \
use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.01\
fixed gender issue and retrained body encoder\
remove shape from unposed garment\
training on 100 garments with offsets over the average of these\
l2 loss on both 128 and full res, bottom predicts base, high res predicts offsets over this\
low and high res both use same full res condition encoders, but low res also gets low res versions for main encoder input\
not using filled in versions, but using custom MLP upscaling')
run = wandb.init(project='vae', name=config['name'], notes=config['notes'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('cvae_pyramid.py')
artifact.add_file('models/cvae_model.py')
artifact.add_file('models/cvae_pyramid_model.py')
artifact.add_file('models/pyramid_model.py')
artifact.add_file('models/pyramid_level.py')
artifact.add_file('models/upscaler.py')
artifact.add_file('dataset.py')
artifact.add_file('losses.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')
run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256)
def plot_3d(mesh_gt, mesh_pred_unposed, mesh_pred_unposed_low, faces, collision, collision_faces, collision_color):
mesh_pred_unposed[:,2] += 1
mesh_pred_unposed_low[:,2] += 2
collision[:,2] += 3
fig = go.Figure(data=[
go.Mesh3d(
def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, collision_faces, collision_color):
fig = make_subplots(rows=1, cols=1+len(mesh_pred_unposed_levels)+1+1, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(1+len(mesh_pred_unposed_levels)+1+1)] ])
fig.add_trace(go.Mesh3d(
x=mesh_gt[:,0],
y=mesh_gt[:,2],
z=mesh_gt[:,1],
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
k = faces[:,2],
text="Ground Truth"
),
go.Mesh3d(
x=mesh_pred_unposed[:,0],
y=mesh_pred_unposed[:,2],
z=mesh_pred_unposed[:,1],
row=1, col=1)
for i, mesh in enumerate(mesh_pred_unposed_levels):
fig.add_trace(go.Mesh3d(
x=mesh[:,0],
y=mesh[:,2],
z=mesh[:,1],
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
text='Unposed level_' + str(i)
),
go.Mesh3d(
x=mesh_pred_unposed_low[:,0],
y=mesh_pred_unposed_low[:,2],
z=mesh_pred_unposed_low[:,1],
row=1, col=i+2)
fig.add_trace(go.Mesh3d(
x=mesh_posed[:,0],
y=mesh_posed[:,1],
z=mesh_posed[:,2],
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
text='Posed GT'
),
go.Mesh3d(
row=1, col=1+len(mesh_pred_unposed_levels)+1)
fig.add_trace(go.Mesh3d(
x=collision[:,0],
y=collision[:,2],
z=collision[:,1],
i = collision_faces[:,0],
j = collision_faces[:,1],
k = collision_faces[:,2],
vertexcolor=collision_color
)
])
vertexcolor=collision_color,
text='Posed output on body'
),
row=1, col=1+len(mesh_pred_unposed_levels)+1+1)
fig.update_layout(scene_aspectmode='manual',
scene_aspectratio=dict(x=1, y=3, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,5],),
zaxis = dict(nticks=4, range=[-1,1],),)
if len(mesh_pred_unposed_levels) == 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),
scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1),
scene6_aspectmode='manual', scene6_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,
scene4 = helper,
scene5 = helper,
scene6 = helper,
)
if len(mesh_pred_unposed_levels) == 2:
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),
scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
scene5_aspectmode='manual', scene5_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,
scene4 = helper,
scene5 = helper,
)
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']])
return levels_input
def log_normal_pdf(sample, mean, logvar, raxis=1):
log2pi = tf.math.log(2. * np.pi)
res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis)
return res
def kl_loss(mean, logvar):
kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar)
kl_loss = tf.reduce_mean(kl_loss)
kl_loss *= -0.5
return kl_loss
# @tf.function
@tf.function
def train_step(batch):
with tf.GradientTape() as tape:
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
uv_curvature_low = tf.one_hot(batch['uv_curvature' + resolution_modifier], 7)
pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'],
batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low, batch['uv_mask' + resolution_modifier],])
mesh_pred_full = uv2mesh(pred_full)
mesh_pred_posed_full = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_full)
mesh_pred_low = uv2mesh(pred_low)
mesh_pred_posed_low = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_low)
mesh_orig = batch['mesh_unposed']
levels_input = group_inputs_for_levels(batch)
out_levels = model(levels_input, 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_pred_posed_full = mesh_pred_posed_full + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed_low = mesh_pred_posed_low + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_orig = batch['mesh_unposed']
e_loss = tf.reduce_mean(edge_loss(mesh_pred_full, mesh_orig, batch['mesh_mask']))
c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed_full, batch['mesh_body_posed'], batch['mesh_mask']))
l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed'] - pred_full)), axis=-1), axis=[1, 2])
l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred_full) ** 2), axis=-1), axis=[1, 2])
ms_ssim = tf.image.ssim_multiscale(pred_full + 1, batch['uv_unposed'] + 1, max_val=3)
kl = kl_loss(mean_full, logvar_full)
output = dict()
reconstr_loss_total = 0
kl_total = 0
for i, res in enumerate(config['pyramid']['resolutions']):
res_modifier = '_' + str(res) if res != 512 else ''
pred, mean, logvar = out_levels[i]
pred_low = pred_low * batch['uv_mask' + resolution_modifier][:,:,:,tf.newaxis]
e_loss_low = tf.reduce_mean(edge_loss(mesh_pred_low, mesh_orig, batch['mesh_mask']))
c_loss_low = tf.reduce_mean(collision_penalty(mesh_pred_posed_low, batch['mesh_body_posed'], batch['mesh_mask']))
l1_loss_low = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed' + resolution_modifier] - pred_low)), axis=-1), axis=[1, 2])
l2_loss_low = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed' + resolution_modifier] - pred_low) ** 2), axis=-1), axis=[1, 2])
ms_ssim_low = tf.image.ssim(pred_low + 1, batch['uv_unposed' + resolution_modifier] + 1, max_val=3)
kl_low = kl_loss(mean_full, logvar_full)
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)
reconstr_loss = l2_loss + l2_loss_low
reconstr_loss = tf.reduce_mean(reconstr_loss)
loss = reconstr_loss + balancing_coeff * (kl + kl_low)
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]))
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_unposed'] + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
kl = kl_loss(mean, logvar)
ms_ssim = tf.reduce_mean(ms_ssim)
l1_loss = tf.reduce_mean(l1_loss)
l2_loss = tf.reduce_mean(l2_loss)
mean = tf.reduce_mean(mean_full)
logvar = tf.reduce_mean(logvar_full)
reconstr_loss_total += l2_loss #+ e_loss + c_loss
kl_total += kl
ms_ssim_low = tf.reduce_mean(ms_ssim_low)
l1_loss_low = tf.reduce_mean(l1_loss_low)
l2_loss_low = tf.reduce_mean(l2_loss_low)
mean_low = tf.reduce_mean(mean_low)
logvar_low = tf.reduce_mean(logvar_low)
mean = tf.reduce_mean(mean)
logvar = tf.reduce_mean(logvar)
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 } }
loss = reconstr_loss_total + balancing_coeff * kl_total
mesh_orig_2 = uv2mesh(batch['uv_unposed'])
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred_full) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
metric_low = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred_low) ** 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_full) ** 2, axis=-1))) / tf.reduce_sum(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/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
'train/euclid_error_reconst': metric_2, 'train/reconstr_loss': reconstr_loss, 'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl,
'train/var': np.e ** logvar, 'train/reg_coeff': reg_coeff, 'train/l2_loss': l2_loss,
'train/edge_loss' + resolution_modifier: e_loss_low, 'train/collision_loss' + resolution_modifier: c_loss_low, 'train/l1_loss' + resolution_modifier: l1_loss_low, 'train/euclid_error' + resolution_modifier: metric_low, 'train/ms_ssim' + resolution_modifier: ms_ssim_low,
'train/mean' + resolution_modifier: mean_low, 'train/logvar' + resolution_modifier: logvar_low, 'train/kl_loss' + resolution_modifier: kl_low,
'train/var' + resolution_modifier: np.e ** logvar_low, 'train/l2_loss' + resolution_modifier: l2_loss_low }
return { **{ 'train/loss': loss, 'train/reconstr_loss_total': reconstr_loss_total, 'train/kl_total': kl_total }, **output }
@tf.function
@@ -237,51 +240,53 @@ def valid_step(batch):
def plot_everything(batch: dict, type: str):
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
uv_curvature_low = tf.one_hot(batch['uv_curvature' + resolution_modifier], 7)
pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'],
batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low, batch['uv_mask' + resolution_modifier]])
levels_input = group_inputs_for_levels(batch)
output_levels = model(levels_input, training=False)
mesh_pred_full = uv2mesh(pred_full)
mesh_pred_low = uv2mesh(pred_low)
mesh_orig = batch['mesh_unposed']
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']
uv_static_img = batch['uv_static']
uv_body_posed_img = batch['uv_body_posed']
pred_img = pred_full
pred_low_up = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img, pred_low_up), axis=2)
uv_curvature = np.tile(batch['uv_curvature'].numpy().argmax(-1)[:,:,:,np.newaxis], [1,1,1,3]) / 2 - 1
pred_img = np.concatenate([output_levels[i][0] for i in range(len(output_levels))], axis=2)
img = tf.concat((uv_static_img, uv_body_posed_img, uv_curvature, 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])])
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.npy")
mask = batch['mesh_mask'].numpy().astype(bool)
garment_unposed, faces = maskFacesBatch(mesh_pred_full.numpy(), F, mask)
garment_unposed_low, faces = maskFacesBatch(mesh_pred_low.numpy(), F, mask)
garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
garment_unposed_low = [garment_unposed_low[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(mask.shape[0])]
_, 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])]
gt_posed = [mesh_orig_posed[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], 14475, 3), dtype=np.float32)
for i in range(mask.shape[0]):
mesh_pred_full[i, mask[i], :] += shape_offsets.numpy()[i, mask[i], :]
body = batch['mesh_body_posed'].numpy()
mesh_posed_full = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_full).numpy()
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])]
garment_posed = [mesh_posed_full[i, mask[i], :] for i in range(body.shape[0])]
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_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))]
collision = [np.concatenate((body[i], mesh_posed[i]), axis=0) for i in range(len(mesh_posed))]
collision_faces = [np.concatenate((F, faces[i] + 14475), 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], garment_unposed[i], garment_unposed_low[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
mean_hist = wandb.Histogram(sequence=mean_full.numpy().mean(axis=0))
logvar_hist = wandb.Histogram(sequence=logvar_full.numpy().mean(axis=0))
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, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
return { **figs, **{ type + '/images': imgs, } }
def sample(folder):
@@ -289,7 +294,7 @@ def sample(folder):
load_model(model,folder)
count = 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']))
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
nr_of_samples=2
@@ -407,75 +412,41 @@ def sample(folder):
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())
batch_size = 4
low_resolution = 32
resolution_modifier = '_' + str(low_resolution)
smpl = SMPL()
factory = DatasetFactory()
signature_dynamic = { 'id': (),
'pose_shape': (72 + 11,),
# 'uv_posed': (512, 256, 3),
'uv_body_posed': (512, 256, 3),
'uv_body_posed' + resolution_modifier: (low_resolution, low_resolution//2, 3),
'mesh_body_posed': (14475, 3,),
'mesh_posed': (14475, 3,),
'mesh_unposed': (14475, 3,),
'uv_unposed': (512, 256, 3),
'uv_unposed' + resolution_modifier: (low_resolution, low_resolution//2, 3),
'uv_curvature': (512, 256),
'uv_curvature' + resolution_modifier: (low_resolution, low_resolution//2),
}
signature_static = { 'uv_static': (512, 256, 3),
'uv_static' + resolution_modifier: (low_resolution, low_resolution//2, 3),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
'uv_mask' + resolution_modifier: (low_resolution, low_resolution//2,),
# 'uv_shaped_body': (512, 256, 3),
}
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['pyramid']['resolutions'])
static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=config['pyramid']['resolutions'])
dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=config['pyramid']['resolutions'])
static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'])
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=False, resolutions=[512, low_resolution])
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=[512, low_resolution])
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=batch_size, signature=signature_dynamic)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
model: Pyramid = Pyramid(config)
model = Pyramid(latent_dim=128,
low_res=low_resolution,
freeze_encoders=True,
static_model='checkpoints/ae_static_6',
body_model='checkpoints/ae_body_8',
curvature_model='checkpoints/ae_curvature_7',
upscaler_model='checkpoints/up_2_3')
# lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)
optimizer = tf.optimizers.Adam(learning_rate=0.0001)
balancing_coeff = 0.0001
optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
balancing_coeff = config['starting_balancing_coeff']
reg_coeff = 1
folder = 'models/ae_logs/'
valid_iterator = iter(dataset_valid)
gamma = 1
balancing_activated_step = 0
# avg_jumpsuit = np.load('custom_data/jumpsuit_20_avg_uv.npy')
# sample('checkpoints/cvae_13')
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3)
load_model(model.cond_encoder, 'checkpoints/cvae_19')
load_model(model.levels[0], 'checkpoints/cvae_low_4_3_d32')
load_model(model.levels[-1], 'checkpoints/cvae_19')
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)
@@ -498,18 +469,7 @@ if __name__ == '__main__':
if balancing_activated_step == 0:
balancing_activated_step = step
balancing_coeff = balancing_coeff + 1 / 15000
# if train_loss['train/euclid_error'] < 0.015:
# reg_coeff = reg_coeff + 100 / 15000
# if step > 15000 and balancing_coeff <= 1e-3:
# balancing_coeff = 1e-2
# if step > 20000 and balancing_coeff <= 1e-3:
# balancing_coeff = 1e-2
# if step > 25000 and balancing_coeff <= 1e-2:
# balancing_coeff = 1e-1
# if step > 30000 and balancing_coeff <= 1e-1:
# balancing_coeff = 1
# gamma = 0.99 * gamma + 0.01 * train_loss['train/reconstr_loss']
to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } }