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.py
+1 -48
View File
@@ -21,7 +21,7 @@ from stylegan2_ae.generator import Generator
from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder
from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from custom_utils.custom_utils import batch_to_dict, plot_3d
from custom_utils.smpl import SMPL
name='ae_body_1'
@@ -113,53 +113,6 @@ def filter_resolutions_featuremaps(resolutions, featuremaps, res):
# return fig
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
class Bottleneck(tf.keras.Model):
auto_encoder_body.py
+33 -124
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()
auto_encoder_curv.py
+26 -59
View File
@@ -5,7 +5,7 @@ import wandb
import random
import numpy as np
import tensorflow as tf
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
@@ -15,20 +15,22 @@ import plotly.express as ex
from tf_utils import check_tf_version, allow_memory_growth
from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL
from models.ae_curv_model import AutoEncoder
from models.cond_encoder import AutoEncoder
from configs.config_ae_curv import config
name='ae_curvature_x_2'
testing = False
def init_wandb():
if not testing:
run = wandb.init(project='vae', name=name, notes='convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128. trained on 100 garments. same as 7')
run = wandb.init(project='vae', name=config['name'], notes='convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128. trained on 100 garments. same as 7')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_curv.py')
artifact.add_file('dataset.py')
artifact.add_file('losses.py')
artifact.add_file('stylegan2_ae/layers/mapping_block.py')
artifact.add_file('configs/config.py')
artifact.add_file('configs/config_ae_curv.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)
@@ -110,7 +112,7 @@ def dice_dense(prediction, ground_truth, weight_map=None):
return tf.reduce_mean(dice_score)
@tf.function
# @tf.function
def train_step(batch):
with tf.GradientTape() as tape:
@@ -121,22 +123,14 @@ def train_step(batch):
ce_uv_loss = scce(batch['uv_curvature'], pred)
ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1))
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)[:,:,:,tf.newaxis], batch['uv_curvature'][:,:,:,tf.newaxis], max_val=7*2))
# loss = 1e3 * (1 - ms_ssim)
loss = ce_uv_loss
# mesh_pred = uv2mesh(tf.tile(tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)[:,:,:,tf.newaxis], [1,1,1,3]))[:,:,0]
# mesh_orig = tf.one_hot(batch['mesh_curvature'] + tf.cast(batch['mesh_mask'], tf.int32), 7)
# ce_mesh_loss = scce(mesh_orig, mesh_pred)
mean_iou = get_mean_iou(batch['uv_curvature'], tf.math.argmax(pred, axis=-1))
dice = dice_dense(pred, uv_curvature)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return { 'train/loss': loss, 'train/ce_uv_loss': ce_uv_loss, 'train/ssim': ms_ssim, 'train/mean_iou': mean_iou, 'train/dice': dice #'ce_mesh_loss': ce_mesh_loss
}
return { 'train/loss': loss, 'train/ce_uv_loss': ce_uv_loss, 'train/ssim': ms_ssim, 'train/mean_iou': mean_iou, 'train/dice': dice }
@@ -157,11 +151,10 @@ def valid_step(batch):
}
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)
@@ -194,7 +187,7 @@ def train():
def test():
ckpt.restore(manager.latest_checkpoint)
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)
pred = model([uv_curvature, batch['uv_mask']])
@@ -208,57 +201,31 @@ def test():
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()
scce = tf.keras.losses.SparseCategoricalCrossentropy()
resolution_modifier = ''
if config['resolution'] != 512:
resolution_modifier = '_' + str(config['resolution'])
factory = DatasetFactory()
signature_dynamic = { 'id': (),
# 'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
# 'uv_body_posed': (3, 512, 256,),
# 'mesh_body_posed': (14475, 3,),
# 'mesh_posed': (14475, 3,),
# 'mesh_unposed': (14475, 3,),
# 'uv_unposed': (3, 512, 256,)
'uv_curvature': (512, 256),
'mesh_curvature': (14475)
}
signature_static = { #'uv_static': (3, 512, 256,),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
# 'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100',
batch_size=8,
signature=signature_dynamic,
use_filled=False)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100',
signature=signature_static,
use_filled=False)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
batch_size=8,
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'])
valid_iterator = iter(dataset_valid)
model = AutoEncoder()
# optimizer = tf.optimizers.Adam(learning_rate=0.002, beta_1=0.0, beta_2=0.99, epsilon=1e-8)
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.001, decay_steps=10000, decay_rate=0.1, staircase=False)
optimizer = tf.optimizers.Adam(learning_rate=lr_schedule)
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=5)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3)
if testing:
test()
auto_encoder_static.py
+38 -114
View File
@@ -1,5 +1,7 @@
import imp
import os
from configs.config_ae_static import config
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
import time
import argparse
import wandb
@@ -13,24 +15,22 @@ import plotly.express as ex
from tf_utils import check_tf_version, allow_memory_growth
from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, 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_static_model import AutoEncoder
from custom_utils.custom_utils import batch_to_dict, plot_3d
from models.cond_encoder import AutoEncoder
name='ae_static_5'
def init_wandb():
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
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.\
convnext first stage down/upscale by factor fo 2 instead of 4, and 5 stages instead of 4 so bottlneck same size\
run = wandb.init(project='vae', name=config['name'],
notes='convnext + convnext, add dense with tanh at the end of reverse convnext,\
flattened latent dim to 128\
trai on 100 garments')
train on 100 garments')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_static.py')
artifact.add_file('losses.py')
artifact.add_file('configs/config_base.py')
artifact.add_file('configs/config_ae_static.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)
@@ -39,86 +39,24 @@ 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_static'])
uv_mask = batch['uv_mask']
pred = model([batch['uv_static'], batch['uv_mask']])
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_static']
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, 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_static'] - pred)), axis=1) * uv_mask)
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) * uv_mask[:,:,:,tf.newaxis] + 1, tf.transpose(batch['uv_static'], [0,2,3,1]) * uv_mask[:,:,:,tf.newaxis] + 1, max_val=2))
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_static'], tf.zeros_like(batch['uv_static'])))), axis=[1,2,3]))
l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_static'] - pred)), axis=1))
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_static'] + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
loss = 1e3 * (1 - ms_ssim)
mesh_orig_2 = uv2mesh(batch['uv_static'])
# 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, }
@@ -145,14 +83,13 @@ def valid_step(batch):
def plot_everything(batch: dict, type: str):
pred = model(batch['uv_static'])
pred = model([batch['uv_static'], batch['uv_mask']])
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_static']
orig_img = tf.transpose(batch['uv_static'], [0,2,3,1])
pred_img = tf.transpose(pred, [0,2,3,1])
orig_img = batch['uv_static']
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)]
@@ -175,20 +112,23 @@ def plot_everything(batch: dict, type: str):
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))
valid_loss = {}#valid_step(data_valid)
if step % 1000 == 0:
if step % 2500 == 0:
train_figs = plot_everything(batch, 'train')
valid_figs = {}#plot_everything(data_valid, 'valid')
wandb.log({ **train_loss, **valid_loss, **train_figs, **valid_figs }, step=step)
else:
wandb.log({ **train_loss, **valid_loss }, step=step)
if step % 10000 == 0:
manager_enc.save(checkpoint_number=step)
manager_dec.save(checkpoint_number=step)
print(step)
@@ -201,41 +141,25 @@ if __name__ == '__main__':
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
smpl = SMPL()
# TODO: different resolutions not working yet, model should be ready though
resolution_modifier = '_32'
resolutions = [32]
resolution_modifier = ''
if config['resolution'] != 512:
resolution_modifier = '_' + str(config['resolution'])
factory = DatasetFactory()
signature_dynamic = { 'id': (),
# 'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
# 'uv_body_posed': (3, 512, 256,),
# 'mesh_body_posed': (14475, 3,),
# 'mesh_posed': (14475, 3,),
# 'mesh_unposed': (14475, 3,),
# 'uv_unposed': (3, 512, 256,)
}
signature_static = { 'uv_static': (3, 512, 256,),
'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
# 'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=8, signature=signature_dynamic)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=8, 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 = AutoEncoder(config['cond_ae'])
optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
model = AutoEncoder()
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0005, decay_steps=30000, decay_rate=0.1, staircase=False)
optimizer = tf.optimizers.Adam(learning_rate=lr_schedule)
folder = 'models/ae_logs/'
valid_iterator = iter(dataset_valid)
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=5)
ckpt_enc = tf.train.Checkpoint(model=model.encoder)
manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=3)
ckpt_dec = tf.train.Checkpoint(model=model.decoder)
manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=3)
train()
configs/config_ae_body.py
+14
View File
@@ -0,0 +1,14 @@
from configs.config_base import config
config['gpu'] = '1'
config['lr'] = 0.0001
config['name'] = 'ae_body_9'
config['signature_dynamic'] = { 'id': (),
'pose_shape': (72 + 11,),
'uv_body_posed': (3, 512, 256,),
'mesh_body_posed': (14475, 3,),
}
config['signature_static'] = {
'mesh_mask': (14475,),
}
configs/config_ae_curv.py
+15
View File
@@ -0,0 +1,15 @@
from configs.config_base import config
config['lr'] = 0.001
config['name'] = 'ae_body_9'
config['cond_ae']['is_label'] = True
config['signature_dynamic'] = { 'id': (),
'uv_curvature': (512, 256),
'mesh_curvature': (14475)
}
config['signature_static'] = { 'mesh_mask': (14475,),
'uv_mask': (512, 256,),
}
configs/config_ae_static.py
+12
View File
@@ -0,0 +1,12 @@
from configs.config_base import config
config['gpu'] = '0'
config['lr'] = 0.0005
config['name'] = 'ae_static_7'
config['signature_dynamic'] = { 'id': (), }
config['signature_static'] = { 'uv_static': (512, 256, 3),
'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
}
configs/config_base.py
+21
View File
@@ -0,0 +1,21 @@
config = dict()
config['batch_size'] = 4
config['train_dir'] = '/home/workspace/local-data/train_100'
config['valid_dir'] = '/home/workspace/local-data/train_100'
config['use_filled'] = False
config['resolution'] = 512
config['latent_dim'] = 128
config['cond_ae'] = dict()
config['cond_ae']['is_label'] = False
config['cond_ae']['latent_dim'] = config['latent_dim']
config['cond_ae']['kernel_size'] = 7
config['cond_ae']['bottleneck_size'] = 64
config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
configs/config_cvae.py
+53
View File
@@ -0,0 +1,53 @@
from configs.config_base import config
config['lr'] = 0.0001
config['name'] = 'cvae_19'
config['starting_balancing_coeff'] = 0.0001
config['signature_dynamic'] = { 'id': (),
'pose_shape': (72 + 11,),
'uv_body_posed': (512, 256, 3),
'mesh_body_posed': (14475, 3,),
'mesh_posed': (14475, 3,),
'mesh_unposed': (14475, 3,),
'uv_unposed': (512, 256, 3),
'uv_curvature': (512, 256),
}
config['signature_static'] = { 'uv_static': (512, 256, 3),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
}
config['avg_garment_path'] = 'custom_data/train_100_no_shape_avg_uv.npy'
config['static_model'] = 'checkpoints/ae_static_6_copy'
config['body_model'] = 'checkpoints/ae_body_8_copy'
config['curvature_model'] = 'checkpoints/ae_curvature_7_copy'
config['freeze_encoders'] = True
config['cvae'] = dict()
config['cvae']['is_label'] = False
config['cvae']['latent_dim'] = config['latent_dim']
if config['resolution'] == 512:
config['cvae']['kernel_size'] = 7
config['cvae']['encoder_depths'] = [3, 3, 9, 3]
config['cvae']['encoder_dims'] = [32, 64, 128, 256]
config['cvae']['decoder_input'] = (16, 8, 80)
config['cvae']['decoder_depths'] = [3, 3, 9, 3]
config['cvae']['decoder_dims'] = [256, 128, 64, 16]
elif config['resolution'] == 128:
config['cvae']['kernel_size'] = 5
config['cvae']['encoder_depths'] = [16, 16, 9, 9]
config['cvae']['encoder_dims'] = [64, 128, 256, 256]
config['cvae']['decoder_input'] = (4, 2, 128)
config['cvae']['decoder_depths'] = [16, 16, 9, 9]
config['cvae']['decoder_dims'] = [256, 256, 128, 64]
elif config['resolution'] == 32:
config['cvae']['kernel_size'] = 3
config['cvae']['encoder_depths'] = [27, 27]
config['cvae']['encoder_dims'] = [256, 256]
config['cvae']['decoder_input'] = (4, 2, 256)
config['cvae']['decoder_depths'] = [27, 27]
config['cvae']['decoder_dims'] = [256, 256]
else:
raise NotImplementedError
convnext/models/convnext_tf.py
+12 -9
View File
@@ -104,16 +104,18 @@ class ConvNeXt(tf.keras.Model):
super().__init__()
self.include_top = include_top
self.downsample_layers = [] # stem and 3 intermediate downsampling conv layers
stem = tf.keras.Sequential([
Conv2D(dims[0], kernel_size=4, strides=4, padding='same'),
LayerNormalization(epsilon=1e-6)]
)
# stem = tf.keras.Sequential([
# Conv2D(dims[0], kernel_size=4, strides=4, padding='same'),
# LayerNormalization(epsilon=1e-6)]
# )
stem = [Conv2D(dims[0], kernel_size=4, strides=4, padding='same'), LayerNormalization(epsilon=1e-6)]
self.downsample_layers.append(stem)
for i in range(len(depths) - 1):
downsample_layer = tf.keras.Sequential([
LayerNormalization(epsilon=1e-6),
Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
)
# downsample_layer = tf.keras.Sequential([
# LayerNormalization(epsilon=1e-6),
# Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
# )
downsample_layer = [LayerNormalization(epsilon=1e-6), Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
self.downsample_layers.append(downsample_layer)
self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks
@@ -138,7 +140,8 @@ class ConvNeXt(tf.keras.Model):
def forward_features(self, x):
for i in range(len(self.downsample_layers)):
x = self.downsample_layers[i](x)
x = self.downsample_layers[i][0](x)
x = self.downsample_layers[i][1](x)
x = self.stages[i](x)
return x
custom_utils/custom_utils.py
+14 -10
View File
@@ -1,3 +1,4 @@
from faulthandler import disable
import numpy as np
import tensorflow as tf
import plotly.graph_objects as go
@@ -135,14 +136,14 @@ def maskFacesBatch(V, F, mask):
def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
mesh_pred_unposed[:,1] += 1
collision[:,1] += 2
mesh_pred_unposed[:,2] += 1
collision[:,2] += 2
fig = go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=mesh_gt[:,0],
y=mesh_gt[:,1],
z=mesh_gt[:,2],
y=mesh_gt[:,2],
z=mesh_gt[:,1],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
@@ -151,8 +152,8 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli
go.Mesh3d(
# 8 vertices of a cube
x=mesh_pred_unposed[:,0],
y=mesh_pred_unposed[:,1],
z=mesh_pred_unposed[:,2],
y=mesh_pred_unposed[:,2],
z=mesh_pred_unposed[:,1],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
@@ -161,8 +162,8 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli
go.Mesh3d(
# 8 vertices of a cube
x=collision[:,0],
y=collision[:,1],
z=collision[:,2],
y=collision[:,2],
z=collision[:,1],
# i, j and k give the vertices of triangles
i = collision_faces[:,0],
j = collision_faces[:,1],
@@ -202,7 +203,10 @@ def sparse_to_tensor(L):
return tf.SparseTensor(idx, data, shape)
def load_model(model, checkpoint):
def load_model(model, checkpoint, disable_assert=False):
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
ckpt.restore(manager.latest_checkpoint)
if not disable_assert:
assert manager.latest_checkpoint is not None
assert manager.latest_checkpoint.split('/')[-1] != 'ckpt-0'
ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed()
custom_utils/fill_uv.py
+23 -3
View File
@@ -5,12 +5,13 @@ import numpy as np
import plotly.graph_objects as go
import plotly.express as ex
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
import tensorflow as tf
import tensorflow_addons as tfa
def fill(img, mask):
coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0))
img_idx = np.transpose(np.stack(np.nonzero((mask))), (1,0))
img_data = img[img_idx[:,0], img_idx[:,1], :]
interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data)
img_2 = interpolator(coords)
res = np.zeros((512,256,3), dtype=np.float32)
@@ -19,6 +20,19 @@ def fill(img, mask):
return res
# @tf.function
def fill_tf(img, mask):
img = img[tf.newaxis,:,:,:]
coords_int = tf.transpose(tf.stack(tf.experimental.numpy.nonzero(1-mask)), (1,0))
coords = tf.cast(coords_int, dtype=tf.float32)
img_2 = tfa.image.interpolate_bilinear(img, coords[tf.newaxis,:,:])
# coords_2 = list(zip(tf.ones(coords_int.shape[0]), coords_int[:,0], coords_int[:,1]))
# coords = tf.transpose(tf.stack([coords_int[:,0], coords_int[:,1]]))
out = tf.tensor_scatter_nd_update(img[0], coords_int, img_2[0])
return out
def fill_label(img, mask):
coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0))
@@ -44,7 +58,13 @@ def fill_uv_maps_posed_body(folder, type):
id, _ = frame.split('_')
# mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask']
img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_body_no_orient']
filled = fill(img, mask)
# filled = fill(img, mask)
img = np.zeros((100, 100, 3), dtype=np.float32)
mask = np.zeros((100, 100), dtype=np.int32)
img[25:75, 25:75, :] = 1
mask[25:75, 25:75] = 1
filled = fill_tf(img, mask)
np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_posed_body_no_orient_filled = filled)
cvae.py
+37 -137
View File
@@ -1,5 +1,5 @@
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -31,83 +31,37 @@ 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 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 custom_utils.custom_utils import batch_to_dict, load_model, plot_3d
from dataset import DatasetFactory, StaticDataset
from custom_utils.smpl import SMPL
from models.cvae_model import CVAE
from configs.config_cvae import config
name='cvae_17'
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. \
run = wandb.init(project='vae', name=config['name'],
notes='\
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, FINAL\
remove shape from unposed garment\
training on 100 garments with offsets over the average of these')
training on 100 garments with offsets over the average of these\
fixed autoencoders!!!!!!!!!')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('cvae.py')
artifact.add_file('models/cvae_model.py')
artifact.add_file('dataset.py')
artifact.add_file('losses.py')
artifact.add_file('configs/config.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, faces, collision, collision_faces, collision_color):
mesh_pred_unposed[:,2] += 1
collision[:,2] += 2
fig = go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=mesh_gt[:,0],
y=mesh_gt[:,2],
z=mesh_gt[:,1],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),
go.Mesh3d(
# 8 vertices of a cube
x=mesh_pred_unposed[:,0],
y=mesh_pred_unposed[:,2],
z=mesh_pred_unposed[:,1],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
),
go.Mesh3d(
# 8 vertices of a cube
x=collision[:,0],
y=collision[:,2],
z=collision[:,1],
# i, j and k give the vertices of triangles
i = collision_faces[:,0],
j = collision_faces[:,1],
k = collision_faces[:,2],
vertexcolor=collision_color
)
])
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
def log_normal_pdf(sample, mean, logvar, raxis=1):
log2pi = tf.math.log(2. * np.pi)
@@ -122,13 +76,14 @@ def kl_loss(mean, logvar):
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)
pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']], training=True, use_mask=True)
pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']],
training=True, use_mask=True)
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed']
@@ -141,38 +96,27 @@ def train_step(batch):
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']))
l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2])
l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred) ** 2), axis=-1) * batch['uv_mask'], axis=[1, 2])
l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed'] - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2])
l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1) * batch['uv_mask'], axis=[1, 2])
# ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2)
ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))
ms_ssim = 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))
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale((tf.transpose(pred, [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], (tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], max_val=2))
fft_loss = 0 #tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed'], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed'], [0,2,3,1]))))), axis=[1,2,3])
tv_loss = tf.image.total_variation(pred)
reconstr_loss = l2_loss
reconstr_loss = tf.reduce_mean(reconstr_loss)
kl = kl_loss(mean, logvar)
loss = reconstr_loss + balancing_coeff * kl
# loss = reconstr_loss + balancing_coeff * kl
ms_ssim = tf.reduce_mean(ms_ssim)
l1_loss = tf.reduce_mean(l1_loss)
l2_loss = tf.reduce_mean(l2_loss)
fft_loss = tf.reduce_mean(fft_loss)
tv_loss = tf.reduce_mean(tv_loss)
mean = tf.reduce_mean(mean)
logvar = tf.reduce_mean(logvar)
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) ** 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, 'train/reconstr_loss': reconstr_loss, 'train/tv_loss': tv_loss,
'train/reconstr_loss': reconstr_loss,
'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/reconst_loss_norm': reconstr_loss / (2 * (gamma ** 2)), 'train/var': np.e ** logvar,
'train/reg_coeff': reg_coeff, 'train/l2_loss': l2_loss }
@@ -183,7 +127,7 @@ def valid_step(batch):
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed']
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
@@ -197,10 +141,6 @@ def valid_step(batch):
metric_2 = 0 #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'])
reconstr_loss = l1_loss
# logpx_z = -reconstr_loss
# logpz = log_normal_pdf(z[:,:model.latent_dim], 0., 0.)
# logqz_x = log_normal_pdf(z[:,:model.latent_dim], mean, logvar)
# loss = -tf.reduce_mean(logpx_z + balancing_coeff * (logpz - logqz_x))
kl = kl_loss(mean, logvar)
# loss = reconstr_loss / (2 * (gamma ** 2)) + kl
loss = reconstr_loss + balancing_coeff * kl
@@ -224,9 +164,9 @@ def plot_everything(batch: dict, type: str):
mesh_pred = uv2mesh(pred).numpy()
mesh_orig = batch['mesh_unposed']
orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
uv_static_img = tf.transpose(batch['uv_static'], [0,2,3,1])
uv_body_posed_img = tf.transpose(batch['uv_body_posed'], [0,2,3,1])
orig_img = batch['uv_unposed']
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)]
@@ -259,8 +199,6 @@ def plot_everything(batch: dict, type: str):
return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
def sample(folder):
# run = wandb.init(project='vae_plots', name=folder.split('/')[-1])
load_model(model,folder)
@@ -274,7 +212,7 @@ def sample(folder):
pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1])
mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
mesh_pred = uv2mesh(pred)
mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy()
mesh_orig = batch['mesh_posed'][0].numpy()
mesh_orig_reposed = smpl.run_garment(pose[:1], shape[:1], batch['mesh_unposed'][:1]).numpy()[0]
@@ -286,9 +224,8 @@ def sample(folder):
mask = batch['mesh_mask'].numpy().astype(bool)[0]
_, faces = _maskFaces(np.ones((14475, 3)), F, mask)
pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
mesh_pred_zero = uv2mesh(tf.transpose(pred_zero, [0,3,1,2]))
mesh_pred_zero = uv2mesh(pred_zero)
mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
@@ -382,67 +319,42 @@ def sample(folder):
print()
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
batch_size = config['batch_size']
smpl = SMPL()
factory = DatasetFactory()
signature_dynamic = { 'id': (),
'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
'uv_body_posed': (3, 512, 256,),
'mesh_body_posed': (14475, 3,),
'mesh_posed': (14475, 3,),
'mesh_unposed': (14475, 3,),
'uv_unposed': (3, 512, 256,),
'uv_curvature': (512, 256),
}
signature_static = { 'uv_static': (3, 512, 256,),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
# 'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=False)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False)
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)
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 = CVAE(config)
model = CVAE(avg_garment_path='custom_data/train_100_no_shape_avg_uv.npy',
latent_dim=128, static_model='checkpoints/ae_static_4',
body_model='checpoints/ae_body_8',
curvature_model='checpoints/ae_curvature_7')
# 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)
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)
@@ -465,18 +377,6 @@ 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 } }
cvae_curv.py
+1 -1
View File
@@ -370,7 +370,7 @@ if __name__ == '__main__':
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
model = CVAE(latent_dim=128, static_model='checkpoints/ae_static_4',
model = CVAE(latent_dim=128, static_model='checkpoints/ae_static_5',
body_model='checpoints/ae_body_8',)
cvae_low.py
+92 -65
View File
@@ -21,24 +21,22 @@ import random
import numpy as np
import scipy as sp
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces
from tensorboardX import SummaryWriter
import plotly.graph_objects as go
import plotly.express as ex
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 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_model import CVAE, CVAE128, CVAE32
from models.cvae_model import CVAE#, CVAE128, CVAE32
from models.upscaler import UpscalerBasic as Upscaler
name='cvae_low_1_d128'
name='cvae_low_4_2_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. \
@@ -49,7 +47,8 @@ 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 , absolute, not offsets\
using filled versions')
using custom enhancer for upscaling, using not filled versions\
calculating loss at full res')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('cvae.py')
artifact.add_file('models/cvae_model.py')
@@ -130,62 +129,57 @@ 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' + res_modifier], 7)
pred, mean, logvar, z = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], uv_curvature_low, batch['uv_static'], batch['uv_body_posed'], uv_curvature])
pred, mean, logvar, z = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], uv_curvature_low,
batch['uv_mask' + res_modifier], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']], use_mask=True)
# pred = pred * batch['uv_mask' + res_modifier][:,:,:,tf.newaxis]
mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred)
pred_enhanced = enhance(upscale_enhancer, pred, batch['uv_mask' + res_modifier], training=False)
pred_full_res = tf.image.resize(pred_enhanced, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
mesh_pred = uv2mesh(pred_enhanced)
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_pred_posed = mesh_pred_posed + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
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)
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']))
# l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask' + res_modifier], axis=[1, 2])
l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]) - pred)), axis=-1), axis=[1, 2])
# l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]) - pred) ** 2), axis=-1) * batch['uv_mask' + res_modifier], axis=[1, 2])
l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]) - pred) ** 2), axis=-1), axis=[1, 2])
# ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2)
# ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]) + 1, max_val=3, power_factors=(0.6, 0.4)) #(0.0448, 0.2856, 0.3001, 0.2363, 0.1333)
ms_ssim = tf.image.ssim(pred + 1, tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]) + 1, max_val=3) #(0.0448, 0.2856, 0.3001, 0.2363, 0.1333)
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale((tf.transpose(pred, [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], (tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], max_val=2))
rfft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.rfft3d(pred) - tf.signal.rfft3d(batch['uv_unposed'])), axis=[1,2,3]))
fft_loss = tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1]))))), axis=[1,2,3])
fft_loss_pc, fft_loss_pc_low, fft_loss_pc_high = 0,0,0 #fft_loss_per_channel(pred, batch['uv_unposed'])
tv_loss = tf.image.total_variation(pred)
l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed' + res_modifier] - pred)), axis=-1), axis=[1, 2])
l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed' + res_modifier] - pred) ** 2), axis=-1), axis=[1, 2])
l2_loss_full = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred_full_res) ** 2), axis=-1), axis=[1, 2])
if resolution == 32:
ms_ssim = tf.image.ssim(pred + 1, batch['uv_unposed' + res_modifier] + 1, max_val=3)
ms_ssim_full = tf.image.ssim(pred_full_res + 1, batch['uv_unposed'] + 1, max_val=3)
else:
ms_ssim = tf.image.ssim_multiscale(pred + 1, batch['uv_unposed' + res_modifier] + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))
reconstr_loss = l2_loss
reconstr_loss = l2_loss_full
reconstr_loss = tf.reduce_mean(reconstr_loss)
# logpx_z = -reconstr_loss
# logpz = log_normal_pdf(z[:,:model.latent_dim], 0., 0.)
# logqz_x = log_normal_pdf(z[:,:model.latent_dim], mean, logvar)
# loss = -tf.reduce_mean(logpx_z + balancing_coeff * (logpz - logqz_x))
kl = kl_loss(mean, logvar)
loss = reconstr_loss + balancing_coeff * kl
# loss = reconstr_loss + balancing_coeff * kl
ms_ssim = tf.reduce_mean(ms_ssim)
ms_ssim_full = tf.reduce_mean(ms_ssim_full)
l1_loss = tf.reduce_mean(l1_loss)
l2_loss = tf.reduce_mean(l2_loss)
fft_loss = tf.reduce_mean(fft_loss)
tv_loss = tf.reduce_mean(tv_loss)
l2_loss_full = tf.reduce_mean(l2_loss_full)
mean = tf.reduce_mean(mean)
logvar = tf.reduce_mean(logvar)
mesh_orig_2 = uv2mesh(batch['uv_unposed' + res_modifier])
# mesh_orig_2 = uv2mesh(batch['uv_unposed' + res_modifier])
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'])
# 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/rfft_loss': rfft_loss, 'train/fft_oc_low': fft_loss_pc_low, 'train/fft_pc_high': fft_loss_pc_high,
'train/fft_loss_per_channel': fft_loss_pc, 'train/euclid_error_posed': euclid_posed, 'train/reconstr_loss': reconstr_loss, 'train/tv_loss': tv_loss,
'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/reconst_loss_norm': reconstr_loss / (2 * (gamma ** 2)), 'train/var': np.e ** logvar,
#'train/euclid_error_reconst': metric_2,
'train/l2_loss_full_res': l2_loss_full, 'train/ms_ssim_full_res': ms_ssim_full,
'train/euclid_error_posed': euclid_posed, '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 }
@@ -235,29 +229,36 @@ def plot_everything(batch: dict, type: str):
pred, mean, logvar, z = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], uv_curvature_low, batch['uv_static'], batch['uv_body_posed'], uv_curvature])
# pred = pred * batch['uv_mask' + res_modifier][:,:,:,tf.newaxis]
mesh_pred = uv2mesh(pred)
pred = enhance(upscale_enhancer, pred, batch['uv_mask' + res_modifier], training=False)
mesh_pred = uv2mesh(pred).numpy()
mesh_orig = batch['mesh_unposed']
orig_img = tf.transpose(batch['uv_unposed' + res_modifier], [0,2,3,1])
uv_static_img = tf.transpose(batch['uv_static' + res_modifier], [0,2,3,1])
uv_body_posed_img = tf.transpose(batch['uv_body_posed' + res_modifier], [0,2,3,1])
pred_img = pred
orig_img = batch['uv_unposed' + res_modifier]
uv_static_img = batch['uv_static' + res_modifier]
uv_body_posed_img = batch['uv_body_posed' + res_modifier]
pred_img = pred.numpy()
pred_img[pred_img > 1] = 0
pred_img[pred_img < -1] = 0
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)]
F = np.load("custom_utils/smpl_faces.npy")
mask = batch['mesh_mask'].numpy().astype(bool)
garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
body = batch['mesh_body_posed'].numpy()
mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
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_posed = [mesh_posed[i, mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(body.shape[0])]
F = np.load("custom_utils/smpl_faces.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])]
for i in range(pred.shape[0]):
mesh_pred[i, mask[i], :] += shape_offsets.numpy()[i, mask[i], :]
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])]
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))]
@@ -393,6 +394,28 @@ def sample(folder):
print()
def get_upscale_enhancer(checkpoint, res):
upscale_enhancer = Upscaler(512, res)
load_model(upscale_enhancer, checkpoint)
upscale_enhancer = tf.keras.Sequential([
upscale_enhancer.flatten,
upscale_enhancer.seq,
upscale_enhancer.reshape
])
upscale_enhancer.trainable = False
return upscale_enhancer
def enhance(upscale_enhancer, input, mask, training=None):
x = input * mask[:,:,:,tf.newaxis]
x = upscale_enhancer(x, training=training)
x = x * (1 - mask[:,:,:,tf.newaxis])
x = x + input
return x
if __name__ == '__main__':
init_wandb()
allow_memory_growth()
@@ -404,7 +427,7 @@ if __name__ == '__main__':
smpl = SMPL()
resolution = 128
resolution = 32
res_modifier = ''
if resolution != 512:
res_modifier = '_' + str(resolution)
@@ -412,37 +435,41 @@ if __name__ == '__main__':
factory = DatasetFactory()
signature_dynamic = { 'id': (),
'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
'uv_body_posed': (3, 512, 256,),
'uv_body_posed' + res_modifier: (3, resolution, resolution//2,),
# 'uv_posed': (512, 256, 3),
'uv_body_posed': (512, 256, 3),
'uv_body_posed' + res_modifier: (resolution, resolution//2, 3),
'mesh_body_posed': (14475, 3,),
'mesh_posed': (14475, 3,),
'mesh_unposed': (14475, 3,),
'uv_unposed': (3, 512, 256,),
'uv_unposed' + res_modifier: (3, resolution, resolution//2,),
'uv_unposed': (512, 256, 3),
'uv_unposed' + res_modifier: (resolution, resolution//2, 3),
'uv_curvature': (512, 256,),
'uv_curvature' + res_modifier: (resolution, resolution//2),
}
signature_static = { 'uv_static': (3, 512, 256,),
'uv_static' + res_modifier: (3, resolution, resolution//2,),
signature_static = { 'uv_static': (512, 256, 3),
'uv_static' + res_modifier: (resolution, resolution//2, 3),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask' + res_modifier: (resolution, resolution//2,),
# 'uv_shaped_body': (3, 512, 256,),
'uv_mask': (512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=True, resolutions=[512, resolution])
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=True, resolutions=[512, resolution])
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=False, resolutions=[512, resolution])
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=[512, 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 = CVAE128(#avg_garment_path='custom_data/train_100_no_shape_avg_uv.npy',
latent_dim=128, static_model='checkpoints/ae_static_4',
body_model='checpoints/ae_body_8',
curvature_model='checpoints/ae_curvature_7',
model = CVAE(#avg_garment_path='custom_data/train_100_no_shape_avg_uv.npy',
size=resolution,
latent_dim=128, static_model='checkpoints/ae_static_6_copy',
body_model='checkpoints/ae_body_8_copy',
curvature_model='checkpoints/ae_curvature_7_copy',
freeze_encoders=True,
use_full_res_encoders=True)
encoder_res=512)
upscale_enhancer = get_upscale_enhancer('checkpoints/up_2_3', res=resolution)
# 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)
cvae_pyramid.py
+28 -26
View File
@@ -1,5 +1,5 @@
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -37,7 +37,7 @@ from custom_utils.smpl import SMPL
from models.cvae_pyramid_model import Pyramid
name='cvae_pyr_3_d32'
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. \
@@ -130,7 +130,7 @@ def kl_loss(mean, logvar):
return kl_loss
@tf.function
# @tf.function
def train_step(batch):
with tf.GradientTape() as tape:
@@ -154,17 +154,17 @@ def train_step(batch):
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(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred_full)), axis=-1), axis=[1, 2])
l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred_full) ** 2), axis=-1), axis=[1, 2])
ms_ssim = tf.image.ssim_multiscale(pred_full + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3)
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)
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(tf.transpose(batch['uv_unposed' + resolution_modifier], [0,2,3,1]) - pred_low)), axis=-1), axis=[1, 2])
l2_loss_low = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed' + resolution_modifier], [0,2,3,1]) - pred_low) ** 2), axis=-1), axis=[1, 2])
ms_ssim_low = tf.image.ssim(pred_low + 1, tf.transpose(batch['uv_unposed' + resolution_modifier], [0,2,3,1]) + 1, max_val=3)
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)
reconstr_loss = l2_loss + l2_loss_low
@@ -246,9 +246,9 @@ def plot_everything(batch: dict, type: str):
mesh_pred_low = uv2mesh(pred_low)
mesh_orig = batch['mesh_unposed']
orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
uv_static_img = tf.transpose(batch['uv_static'], [0,2,3,1])
uv_body_posed_img = tf.transpose(batch['uv_body_posed'], [0,2,3,1])
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)
@@ -264,12 +264,14 @@ def plot_everything(batch: dict, type: str):
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], :] for i in range(pred_full.shape[0])]
gt = [mesh_orig[i].numpy()[mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
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()
garment_posed = [mesh_posed_full[i, mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(body.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))]
@@ -405,7 +407,7 @@ def sample(folder):
if __name__ == '__main__':
init_wandb()
# init_wandb()
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
@@ -420,24 +422,24 @@ if __name__ == '__main__':
factory = DatasetFactory()
signature_dynamic = { 'id': (),
'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
'uv_body_posed': (3, 512, 256,),
'uv_body_posed' + resolution_modifier: (3, low_resolution, low_resolution//2,),
# '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': (3, 512, 256,),
'uv_unposed' + resolution_modifier: (3, low_resolution, low_resolution//2,),
'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': (3, 512, 256,),
'uv_static' + resolution_modifier: (3, 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': (3, 512, 256,),
# 'uv_shaped_body': (512, 256, 3),
}
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])
@@ -449,9 +451,9 @@ if __name__ == '__main__':
model = Pyramid(latent_dim=128,
low_res=low_resolution,
freeze_encoders=True,
static_model='checkpoints/ae_static_4',
body_model='checpoints/ae_body_7',
curvature_model='checpoints/ae_curvature_7',
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)
cvae_vae.py
+1 -48
View File
@@ -36,7 +36,7 @@ from stylegan2_ae.generator import Generator
from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder, sparse_to_tensor
from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from custom_utils.custom_utils import batch_to_dict, plot_3d
from custom_utils.smpl import SMPL
from auto_encoder_body import AutoEncoder as BodyAE
from auto_encoder_curv import AutoEncoder as CurvatureAE
@@ -61,53 +61,6 @@ rebalanced l1 loss. retrained body encoder with deoriented bodies \
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
class Bottleneck(tf.keras.Model):
dataset.py
+42 -45
View File
@@ -10,46 +10,46 @@ from custom_utils.custom_utils import nest_conditions, batch_to_dict
class DatasetSimple(tf.data.Dataset):
# class DatasetSimple(tf.data.Dataset):
def _generator(folder):
# Opening the file
folder = folder.decode("utf-8")
uv_folder = os.path.join(folder, 'uv_posed')
pose_hape_folder = os.path.join(folder, 'poses_shapes')
frames = os.listdir(uv_folder)
random.shuffle(frames)
# def _generator(folder):
# # Opening the file
# folder = folder.decode("utf-8")
# uv_folder = os.path.join(folder, 'uv_posed')
# pose_hape_folder = os.path.join(folder, 'poses_shapes')
# frames = os.listdir(uv_folder)
# random.shuffle(frames)
for frame in frames:
id, frame_nr = frame.split('.')[0].split('_')
uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1])
uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']; uv_body_posed = np.transpose(uv_body_posed, [2, 0, 1])
pose_shape = np.load(os.path.join(folder, 'pose_shape', id + '_' + frame_nr + '.npy'))
uv_static = np.load(os.path.join(folder, 'uv_static', id + '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1])
mesh_static = np.load(os.path.join(folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']
mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32)
uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32)
mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed']
mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
# for frame in frames:
# id, frame_nr = frame.split('.')[0].split('_')
# uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1])
# uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']; uv_body_posed = np.transpose(uv_body_posed, [2, 0, 1])
# pose_shape = np.load(os.path.join(folder, 'pose_shape', id + '_' + frame_nr + '.npy'))
# uv_static = np.load(os.path.join(folder, 'uv_static', id + '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1])
# mesh_static = np.load(os.path.join(folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']
# mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
# mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32)
# uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32)
# mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed']
# mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
# uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
# uv_posed, (pose, shape, uv_static), mesh_static, mesh_body_posed, mask
yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static, uv_mask, uv_body_posed), mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed
# # uv_posed, (pose, shape, uv_static), mesh_static, mesh_body_posed, mask
# yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static, uv_mask, uv_body_posed), mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed
def __new__(cls, folder, res=256):
return tf.data.Dataset.from_generator(
cls._generator,
output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32),
(tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ),
tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
tf.TensorSpec(shape = (14475,), dtype = tf.float32),
tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), ),
args=(folder, )
)
# def __new__(cls, folder, res=256):
# return tf.data.Dataset.from_generator(
# cls._generator,
# output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32),
# (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ),
# tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
# tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
# tf.TensorSpec(shape = (14475,), dtype = tf.float32),
# tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
# tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
# tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), ),
# args=(folder, )
# )
@@ -59,7 +59,7 @@ class StaticDataset():
self.signature = signature
self.folder = folder
assert set(resolutions).issubset(set([32, 128, 512]))
self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_128' 'mesh_static', 'mesh_mask', 'uv_mask', 'uv_mask_32', 'uv_mask_128', 'uv_shaped_body']
self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_128', 'mesh_static', 'mesh_mask', 'uv_mask', 'uv_mask_32', 'uv_mask_128', 'uv_shaped_body']
self.keys = [value for value in signature.keys() if value in self.possible_keys]
if 'uv_static_128' in signature.keys():
self.keys.append('uv_static_128')
@@ -78,7 +78,7 @@ class StaticDataset():
res_string = '_' + str(res)
if 'uv_static' + res_string in self.signature.keys():
uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; uv_static = np.transpose(uv_static, [2, 0, 1]); data['uv_static' + res_string] = uv_static
uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; data['uv_static' + res_string] = uv_static
if 'uv_mask' + res_string in self.signature.keys():
uv_mask = np.load(os.path.join(self.folder, 'uv_mask' + res_string, id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask
@@ -88,14 +88,11 @@ class StaticDataset():
if 'mesh_mask' in self.signature.keys():
mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask
if 'uv_shaped_body' in self.signature.keys():
uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; uv_shaped_body = np.transpose(uv_shaped_body, [2, 0, 1]); data['uv_shaped_body'] = uv_shaped_body
uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; data['uv_shaped_body'] = uv_shaped_body
return data
def __getitem__(self, id) -> dict:
# id = tf.strings.unicode_decode(id, 'UTF-8')
# id = tf.strings.as_string(id)
# id = id.numpy()
if id not in self.store:
self.store[id] = self.read(id)
@@ -170,7 +167,7 @@ class DatasetFactory():
item['pose_shape'] = pose_shape
if 'uv_posed' in signature.keys():
uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1])
uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']
item['uv_posed'] = uv_posed
# if 'uv_body_posed' in signature.keys():
# uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']; uv_body_posed = np.transpose(uv_body_posed, [2, 0, 1])
@@ -218,11 +215,11 @@ class DatasetFactory():
item['uv_curvature' + res_string] = uv_curvature.astype(np.int32)
if 'uv_unposed' + res_string in signature.keys():
uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + filled + res_string, id + '_' + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape' + filled]; uv_unposed_no_shape = np.transpose(uv_unposed_no_shape, [2, 0, 1])
uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + filled + res_string, id + '_' + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape' + filled]
item['uv_unposed' + res_string] = uv_unposed_no_shape
if 'uv_body_posed' + res_string in signature.keys():
uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient' + filled + res_string, id + '_' + frame_nr + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient' + filled]; uv_posed_body_no_orient = np.transpose(uv_posed_body_no_orient, [2, 0, 1])
uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient' + filled + res_string, id + '_' + frame_nr + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient' + filled]
item['uv_body_posed' + res_string] = uv_posed_body_no_orient
downscale_unpose_test.py
+1 -1
View File
@@ -49,7 +49,7 @@ if __name__ == '__main__':
mask_idx = np.array(range(14475))[mesh_mask]
uv_posed_filled = fill(uv_posed, uv_mask)
down = tf.image.resize(uv_posed_filled, size=(32, 16), method=tf.image.ResizeMethod.LANCZOS5)
down = tf.image.resize(uv_posed_filled, size=(32, 16), method=tf.image.ResizeMethod.BILINEAR)
mesh_posed_low = uv2mesh(down[tf.newaxis,:,:,:])
unposed = unpose(mesh_static[mesh_mask,:], mesh_posed[np.newaxis,mesh_mask,:], pose[:,np.newaxis], shape, gender, mesh_mask, mask_idx, smpl_2, id)
downscaling_test.py
+19 -15
View File
@@ -44,11 +44,11 @@ def downscale_1(img):
def downscale_2(img):
results = [img]
for i in range(5):
# print(np.array(img.shape[:2]) / (2 ** (i + 1)))
# results.append(tf.image.resize(img, size=np.array(img.shape[:2]) / (2 ** (i + 1)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR, antialias=True, preserve_aspect_ratio=True))
size = (np.array(img.shape[:2]) / (2 ** (i + 1))).astype(np.int32)
size = (size[1], size[0])
results.append(cv2.resize(img, size, interpolation=cv2.INTER_NEAREST))
print(np.array(img.shape[:2]) / (2 ** (i + 1)))
results.append(tf.image.resize(img, size=np.array(img.shape[:2]) / (2 ** (i + 1)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR, antialias=True, preserve_aspect_ratio=True))
# size = (np.array(img.shape[:2]) / (2 ** (i + 1))).astype(np.int32)
# size = (size[1], size[0])
# results.append(cv2.resize(img, size, interpolation=cv2.INTER_NEAREST))
return results
def upscale_1(imgs):
@@ -127,27 +127,30 @@ if __name__ == '__main__':
factory = DatasetFactory()
signature_dynamic = { 'id': (),
'pose_shape': (72 + 11,),
'uv_posed': (3, 512, 256,),
# 'uv_posed': (3, 512, 256,),
# 'uv_body_posed': (3, 512, 256,),
'mesh_body_posed': (14475, 3,),
'mesh_posed': (14475, 3,),
# 'mesh_unposed': (14475, 3,),
# 'uv_unposed': (3, 512, 256,)
}
signature_static = { #'uv_static': (3, 512, 256,),
signature_static = { 'uv_static': (3, 512, 256,),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
'uv_mask_32': (32, 16,),
# 'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100',
batch_size=1,
signature=signature_dynamic,
use_filled=True)
use_filled=True,
resolutions=[512, 32])
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100',
signature=signature_static,
use_filled=True)
use_filled=True,
resolutions=[512, 32])
step = 0
@@ -155,7 +158,7 @@ if __name__ == '__main__':
# batch = batch_to_dict(batch, signature_dynamic)
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
img = tf.transpose(tf.squeeze(batch['uv_posed']), [1,2,0]).numpy()
img = tf.transpose(tf.squeeze(batch['uv_static']), [1,2,0]).numpy()
coords = np.transpose(np.stack(np.nonzero((1 - batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
img_idx = np.transpose(np.stack(np.nonzero((batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
img_data = img[img_idx[:,0], img_idx[:,1], :]
@@ -169,7 +172,8 @@ if __name__ == '__main__':
img_2 = res
uv_mask = batch['uv_mask'].numpy()[0]
uv_mask = downscale_2(uv_mask)
uv_mask_32 = batch['uv_mask_32'].numpy()[0]
# uv_mask = downscale_2(uv_mask)
# for i in range(0,3):
# uv_mask[i] = cv2.dilate(uv_mask[i].astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)#.astype(bool)[:, :, np.newaxis]
# for i in range(3,6):
@@ -178,9 +182,9 @@ if __name__ == '__main__':
down_1 = downscale_1(img)
down_2 = downscale_1(img_2)
up_1 = upscale_1(down_1)
up_1 = upscale_1([down_1[-2]])
# down_2 = [down_2[i] * uv_mask[i][:,:,np.newaxis] for i in range(len(down_2))]
up_2 = upscale_1(down_2)
up_2 = upscale_1([down_1[-2] * uv_mask_32[:,:,tf.newaxis]])
mask = batch['mesh_mask'].numpy().astype(bool)[0]
mesh_2 = [uv2mesh(tf.transpose(up_2[i], [2,0,1])[tf.newaxis,:,:,:])[0] for i in range(len(up_2))]
@@ -230,7 +234,7 @@ if __name__ == '__main__':
fig = make_subplots(rows=2, cols=6, start_cell="top-left")
for i, d2 in enumerate(uv_mask):
for i, d2 in enumerate(up_1):
fig.add_trace(ex.imshow((d2 + 1) * 128).data[0], row=1, col=i+1)
for i, u2 in enumerate(up_2):
@@ -240,7 +244,7 @@ if __name__ == '__main__':
fig = make_subplots(rows=2, cols=6, start_cell="top-left", specs = [ [{ 'type': 'xy' } for _ in range(6)], [{ 'type': 'scene' } for _ in range(6)] ])
for i, u2 in enumerate(up_2):
for i, u2 in enumerate(up_1):
fig.add_trace(ex.imshow((u2 + 1) * 128).data[0], row=1, col=i+1)
# for i, m1 in enumerate(mesh_1):
downscaling_test_curv.py
+2 -2
View File
@@ -114,10 +114,10 @@ if __name__ == '__main__':
# 'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/stylegan_new/train',
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100',
batch_size=1,
signature=signature_dynamic)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/stylegan_new/train',
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100',
signature=signature_static)
models/ae_body_model.py
+154 -175
View File
@@ -1,223 +1,202 @@
import tensorflow as tf
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
class Bottleneck(tf.keras.Model):
def __init__(self, dim, latent_dim):
super(Bottleneck, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
)
self.flatten = tf.keras.layers.Flatten()
self.dense = tf.keras.layers.Dense(latent_dim)
def call(self, inputs, training=None):
x = inputs
x = self.bottleneck(x)
x = self.flatten(x)
x = self.dense(x)
return x
class Decoder(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder, self).__init__()
self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
from models.base import AEEncoder, Bottleneck, Decoder
class AutoEncoder(tf.keras.Model):
def __init__(self):
def __init__(self, size):
super(AutoEncoder, self).__init__()
if size == 512:
bottleneck_size = 64
latent_dim = 128
kernel_size = 7
encoder_depths = [3, 3, 9, 9, 3]
encoder_dims = [32, 64, 128, 256, bottleneck_size]
decoder_input = (8, 4, bottleneck_size)
decoder_depths = [3, 3, 9, 3, 3]
decoder_dims = [bottleneck_size, 256, 128, 64, 16]
elif size == 128:
raise NotImplementedError
bottleneck_size = 512
latent_dim = 128
self.encoder = ConvNeXt(depths=[3, 3, 9, 9, 3, 3], dims=[32, 64, 128, 256, 256, bottleneck_size], include_top=False)
self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
self.decoder = Decoder(bottleneck_size, latent_dim)
def call(self, inputs, training=None):
x = inputs
x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
# x = self.pooling(x)
# x = self.reshape(x)
# x = self.flatten(x)
x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,:, tf.newaxis]
# x = x * uv_mask_body[tf.newaxis,tf.newaxis,:,:]
return x
class Bottleneck32(tf.keras.Model):
def __init__(self, dim, latent_dim):
super(Bottleneck32, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
)
self.flatten = tf.keras.layers.Flatten()
self.dense = tf.keras.layers.Dense(latent_dim)
def call(self, inputs, training=None):
x = inputs
x = self.bottleneck(x)
x = self.flatten(x)
x = self.dense(x)
return x
class Decoder32(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder32, self).__init__()
self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
self.decoder = ConvNeXtReverse(depths=list(reversed([27, 27])), dims=list(reversed([bottleneck_size, 256])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class AutoEncoder32(tf.keras.Model):
def __init__(self):
super(AutoEncoder32, self).__init__()
kernel_size = 5
encoder_depths =[16, 16, 9, 9]
encoder_dims = [64, 128, 256, bottleneck_size]
decoder_input = (4, 2, bottleneck_size)
decoder_depths = [16, 16, 9, 9]
decoder_dims = [bottleneck_size, 256, 128, 64]
elif size == 32:
raise NotImplementedError
bottleneck_size = 512
latent_dim = 128
kernel_size = 3
encoder_depths = [27, 27]
encoder_dims = [256, bottleneck_size]
decoder_input = (4, 2, bottleneck_size)
decoder_depths = [27, 27]
decoder_dims = [bottleneck_size, 256]
else:
raise NotImplementedError
self.encoder = ConvNeXt(depths=[27, 27], dims=[256, bottleneck_size], include_top=False)
self.bottleneck = Bottleneck32(bottleneck_size, latent_dim)
self.decoder = Decoder32(bottleneck_size, latent_dim)
self.encoder : AEEncoder = AEEncoder(latent_dim, encoder_depths, encoder_dims, bottleneck_size, kernel_size)
self.decoder = Decoder(decoder_input, decoder_depths, decoder_dims, kernel_size=kernel_size)
def call(self, inputs, training=None):
def call(self, inputs, training=None, use_mask=True):
if use_mask:
x, mask = inputs
else:
x = inputs
x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
# x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
x = tf.transpose(x, [0,3,1,2])
if use_mask:
x = x * mask[tf.newaxis,:,:,tf.newaxis]
return x
# class Bottleneck32(tf.keras.Model):
class Bottleneck128(tf.keras.Model):
# def __init__(self, dim, latent_dim):
# super(Bottleneck32, self).__init__()
def __init__(self, dim, latent_dim):
super(Bottleneck128, self).__init__()
# self.bottleneck = tf.keras.Sequential(
# [Block(dim=dim, drop_path=0,
# layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
# )
# self.flatten = tf.keras.layers.Flatten()
# self.dense = tf.keras.layers.Dense(latent_dim)
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
)
self.flatten = tf.keras.layers.Flatten()
self.dense = tf.keras.layers.Dense(latent_dim)
# def call(self, inputs, training=None):
# x = inputs
def call(self, inputs, training=None):
x = inputs
# x = self.bottleneck(x)
# x = self.flatten(x)
# x = self.dense(x)
x = self.bottleneck(x)
x = self.flatten(x)
x = self.dense(x)
return x
# return x
class Decoder128(tf.keras.Model):
# class Decoder32(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder128, self).__init__()
# def __init__(self, bottleneck_size, latent_dim):
# super(Decoder32, self).__init__()
self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
self.decoder = ConvNeXtReverse(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([bottleneck_size, 256, 128, 64])))
# self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
# self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
# self.decoder = ConvNeXtReverse(depths=list(reversed([27, 27])), dims=list(reversed([bottleneck_size, 256])))
def call(self, inputs, training=None):
x = inputs
# def call(self, inputs, training=None):
# x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
# x = self.dense(x)
# x = self.reshape(x)
# x = self.decoder(x)
return x
# return x
class AutoEncoder128(tf.keras.Model):
# class AutoEncoder32(tf.keras.Model):
def __init__(self):
super(AutoEncoder128, self).__init__()
# def __init__(self):
# super(AutoEncoder32, self).__init__()
bottleneck_size = 512
latent_dim = 128
# bottleneck_size = 512
# latent_dim = 128
self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[64, 128, 256, bottleneck_size], include_top=False)
self.bottleneck = Bottleneck128(bottleneck_size, latent_dim)
self.decoder = Decoder128(bottleneck_size, latent_dim)
# self.encoder = ConvNeXt(depths=[27, 27], dims=[256, bottleneck_size], include_top=False)
# self.bottleneck = Bottleneck32(bottleneck_size, latent_dim)
# self.decoder = Decoder32(bottleneck_size, latent_dim)
def call(self, inputs, training=None):
x = inputs
# def call(self, inputs, training=None):
# x = inputs
x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
# x = tf.transpose(x, [0,2,3,1])
# x = self.encoder(x, training=training)
x = self.bottleneck(x, training=training)
# x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
# x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
x = tf.transpose(x, [0,3,1,2])
# x = self.decoder(x, training=training)
# # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# x = tf.transpose(x, [0,3,1,2])
return x
# return x
# class Bottleneck128(tf.keras.Model):
# def __init__(self, dim, latent_dim):
# super(Bottleneck128, self).__init__()
# self.bottleneck = tf.keras.Sequential(
# [Block(dim=dim, drop_path=0,
# layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
# )
# self.flatten = tf.keras.layers.Flatten()
# self.dense = tf.keras.layers.Dense(latent_dim)
# def call(self, inputs, training=None):
# x = inputs
# x = self.bottleneck(x)
# x = self.flatten(x)
# x = self.dense(x)
# return x
# class Decoder128(tf.keras.Model):
# def __init__(self, bottleneck_size, latent_dim):
# super(Decoder128, self).__init__()
# self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
# self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
# self.decoder = ConvNeXtReverse(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([bottleneck_size, 256, 128, 64])))
# def call(self, inputs, training=None):
# x = inputs
# x = self.dense(x)
# x = self.reshape(x)
# x = self.decoder(x)
# return x
# class AutoEncoder128(tf.keras.Model):
# def __init__(self):
# super(AutoEncoder128, self).__init__()
# bottleneck_size = 512
# latent_dim = 128
# self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[64, 128, 256, bottleneck_size], include_top=False)
# self.bottleneck = Bottleneck128(bottleneck_size, latent_dim)
# self.decoder = Decoder128(bottleneck_size, latent_dim)
# def call(self, inputs, training=None):
# x = inputs
# x = tf.transpose(x, [0,2,3,1])
# x = self.encoder(x, training=training)
# x = self.bottleneck(x, training=training)
# x = self.decoder(x, training=training)
# # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# x = tf.transpose(x, [0,3,1,2])
# return x
models/ae_curv_model.py
+108 -102
View File
@@ -1,63 +1,73 @@
import tensorflow as tf
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverseLabels, Block
from models.base import AEEncoder, Bottleneck, Decoder
class Bottleneck(tf.keras.Model):
# class Bottleneck(tf.keras.Model):
def __init__(self, dim, latent_dim):
super(Bottleneck, self).__init__()
# def __init__(self, dim, latent_dim):
# super(Bottleneck, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
)
self.flatten = tf.keras.layers.Flatten()
self.dense = tf.keras.layers.Dense(latent_dim)
# self.bottleneck = tf.keras.Sequential(
# [Block(dim=dim, drop_path=0,
# layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
# )
# self.flatten = tf.keras.layers.Flatten()
# self.dense = tf.keras.layers.Dense(latent_dim)
def call(self, inputs, training=None):
x = inputs
# def call(self, inputs, training=None):
# x = inputs
x = self.bottleneck(x)
x = self.flatten(x)
x = self.dense(x)
# x = self.bottleneck(x)
# x = self.flatten(x)
# x = self.dense(x)
return x
# return x
class Decoder(tf.keras.Model):
# class Decoder(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder, self).__init__()
# def __init__(self, bottleneck_size, latent_dim):
# super(Decoder, self).__init__()
self.dense = tf.keras.layers.Dense(16 * 8 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((16, 8, bottleneck_size))
self.decoder = ConvNeXtReverseLabels(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])), nr_of_classes=7)
# self.dense = tf.keras.layers.Dense(16 * 8 * bottleneck_size)
# self.reshape = tf.keras.layers.Reshape((16, 8, bottleneck_size))
# self.decoder = ConvNeXtReverseLabels(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])), nr_of_classes=7)
def call(self, inputs, training=None):
x = inputs
# def call(self, inputs, training=None):
# x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
# x = self.dense(x)
# x = self.reshape(x)
# x = self.decoder(x)
return x
# return x
class AutoEncoder(tf.keras.Model):
def __init__(self):
def __init__(self, size=512):
super(AutoEncoder, self).__init__()
bottleneck_size = 32
if size == 512:
bottleneck_size = 64
latent_dim = 128
self.decoder = Decoder(bottleneck_size, latent_dim)
self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, bottleneck_size], include_top=False)
kernel_size = 7
encoder_depths = [3, 3, 9, 3, 3]
encoder_dims = [32, 64, 128, 256, bottleneck_size]
decoder_input = (16, 8, bottleneck_size)
decoder_depths = [3, 3, 9, 3, 3]
decoder_dims = [bottleneck_size, 256, 128, 64, 16]
elif size == 128:
raise NotImplementedError
elif size == 32:
raise NotImplementedError
else:
raise NotImplementedError
self.encoder : AEEncoder = AEEncoder(latent_dim, encoder_depths, encoder_dims, bottleneck_size, kernel_size)
self.decoder = Decoder(decoder_input, decoder_depths, decoder_dims, kernel_size, is_label=True)
def call(self, inputs, training=None):
@@ -65,117 +75,113 @@ class AutoEncoder(tf.keras.Model):
# x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
# x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,:, tf.newaxis]
# x = x * uv_mask[:,tf.newaxis,:,:]
# x = x * uv_mask[:,:,:,tf.newaxis]
x = x * uv_mask[:,:,:,tf.newaxis]
return x
class Decoder32(tf.keras.Model):
# class Decoder32(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder32, self).__init__()
# def __init__(self, bottleneck_size, latent_dim):
# super(Decoder32, self).__init__()
self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
self.decoder = ConvNeXtReverseLabels(depths=list(reversed([27, 27])), dims=list(reversed([bottleneck_size, 256])), nr_of_classes=7)
# self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
# self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
# self.decoder = ConvNeXtReverseLabels(depths=list(reversed([27, 27])), dims=list(reversed([bottleneck_size, 256])), nr_of_classes=7)
def call(self, inputs, training=None):
x = inputs
# def call(self, inputs, training=None):
# x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
# x = self.dense(x)
# x = self.reshape(x)
# x = self.decoder(x)
return x
# return x
class AutoEncoder32(tf.keras.Model):
# class AutoEncoder32(tf.keras.Model):
def __init__(self):
super(AutoEncoder32, self).__init__()
# def __init__(self):
# super(AutoEncoder32, self).__init__()
bottleneck_size = 512
latent_dim = 128
# bottleneck_size = 512
# latent_dim = 128
self.encoder = ConvNeXt(depths=[27, 27], dims=[256, bottleneck_size], include_top=False)
self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
self.decoder = Decoder32(bottleneck_size, latent_dim)
# self.encoder = ConvNeXt(depths=[27, 27], dims=[256, bottleneck_size], include_top=False)
# self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
# self.decoder = Decoder32(bottleneck_size, latent_dim)
def call(self, inputs, training=None):
x, uv_mask = inputs
# def call(self, inputs, training=None):
# x, uv_mask = inputs
# x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
# # x = tf.transpose(x, [0,2,3,1])
# x = self.encoder(x, training=training)
x = self.bottleneck(x, training=training)
# x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
# x = tf.transpose(x, [0,3,1,2])
# x = self.decoder(x, training=training)
# # x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,:, tf.newaxis]
# x = x * uv_mask[:,tf.newaxis,:,:]
# x = x * uv_mask[:,:,:,tf.newaxis]
# # x = x * uv_mask[:,:, tf.newaxis]
# # x = x * uv_mask[:,tf.newaxis,:,:]
# # x = x * uv_mask[:,:,:,tf.newaxis]
return x
# return x
class Decoder128(tf.keras.Model):
# class Decoder128(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder128, self).__init__()
# def __init__(self, bottleneck_size, latent_dim):
# super(Decoder128, self).__init__()
self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
self.decoder = ConvNeXtReverseLabels(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([bottleneck_size, 256, 128, 64])), nr_of_classes=7)
# self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
# self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
# self.decoder = ConvNeXtReverseLabels(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([bottleneck_size, 256, 128, 64])), nr_of_classes=7)
def call(self, inputs, training=None):
x = inputs
# def call(self, inputs, training=None):
# x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
# x = self.dense(x)
# x = self.reshape(x)
# x = self.decoder(x)
return x
# return x
class AutoEncoder128(tf.keras.Model):
# class AutoEncoder128(tf.keras.Model):
def __init__(self):
super(AutoEncoder128, self).__init__()
# def __init__(self):
# super(AutoEncoder128, self).__init__()
bottleneck_size = 512
latent_dim = 128
# bottleneck_size = 512
# latent_dim = 128
self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[64, 128, 256, bottleneck_size], include_top=False)
self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
self.decoder = Decoder128(bottleneck_size, latent_dim)
# self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[64, 128, 256, bottleneck_size], include_top=False)
# self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
# self.decoder = Decoder128(bottleneck_size, latent_dim)
def call(self, inputs, training=None):
x, uv_mask = inputs
# def call(self, inputs, training=None):
# x, uv_mask = inputs
# x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
# # x = tf.transpose(x, [0,2,3,1])
# x = self.encoder(x, training=training)
x = self.bottleneck(x, training=training)
# x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
# x = tf.transpose(x, [0,3,1,2])
# x = self.decoder(x, training=training)
# # x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,:, tf.newaxis]
# x = x * uv_mask[:,tf.newaxis,:,:]
# x = x * uv_mask[:,:,:,tf.newaxis]
# # x = x * uv_mask[:,:, tf.newaxis]
# # x = x * uv_mask[:,tf.newaxis,:,:]
# # x = x * uv_mask[:,:,:,tf.newaxis]
return x
# return x
models/ae_static_model.py
+133 -149
View File
@@ -1,172 +1,156 @@
import tensorflow as tf
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
from models.base import AEEncoder, Bottleneck, Decoder
class Bottleneck(tf.keras.Model):
def __init__(self, dim, latent_dim):
super(Bottleneck, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
)
self.flatten = tf.keras.layers.Flatten()
self.dense = tf.keras.layers.Dense(latent_dim)
def call(self, inputs, training=None):
x = inputs
x = self.bottleneck(x)
x = self.flatten(x)
x = self.dense(x)
return x
class Decoder(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder, self).__init__()
self.dense = tf.keras.layers.Dense(8 * 4 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((8, 4, bottleneck_size))
self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3, 3])), dims=list(reversed([bottleneck_size, 128, 128, 64, 16])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class AutoEncoder(tf.keras.Model):
def __init__(self):
def __init__(self, size=512):
super(AutoEncoder, self).__init__()
bottleneck_size = 32
if size == 512:
bottleneck_size = 64
latent_dim = 128
self.encoder = ConvNeXt(depths=[3, 3, 9, 3, 3], dims=[32, 64, 128, 128, bottleneck_size], include_top=False)
self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
self.decoder = Decoder(bottleneck_size, latent_dim)
def call(self, inputs, training=None):
x = inputs
x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,:, tf.newaxis]
# x = x * uv_mask[:,tf.newaxis,:,:]
return x
class Decoder32(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder32, self).__init__()
self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
self.decoder = ConvNeXtReverse(depths=list(reversed([27, 27])), dims=list(reversed([bottleneck_size, 256])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class AutoEncoder32(tf.keras.Model):
def __init__(self):
super(AutoEncoder32, self).__init__()
kernel_size = 7
encoder_depths = [3, 3, 9, 3, 3]
encoder_dims = [32, 64, 128, 256, bottleneck_size]
decoder_input = (8, 4, bottleneck_size)
decoder_depths = [3, 3, 9, 3, 3]
decoder_dims = [bottleneck_size, 256, 128, 64, 16]
elif size == 128:
bottleneck_size = 512
latent_dim = 128
self.encoder = ConvNeXt(depths=[27, 27], dims=[256, bottleneck_size], include_top=False)
self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
self.decoder = Decoder32(bottleneck_size, latent_dim)
def call(self, inputs, training=None):
x = inputs
x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,:, tf.newaxis]
# x = x * uv_mask[:,tf.newaxis,:,:]
return x
class Decoder128(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder128, self).__init__()
self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
self.decoder = ConvNeXtReverse(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([bottleneck_size, 256, 128, 64])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class AutoEncoder128(tf.keras.Model):
def __init__(self):
super(AutoEncoder128, self).__init__()
kernel_size = 5
encoder_depths =[16, 16, 9, 9]
encoder_dims = [64, 128, 256, bottleneck_size]
decoder_input = (4, 2, bottleneck_size)
decoder_depths = [16, 16, 9, 9]
decoder_dims = [bottleneck_size, 256, 128, 64]
elif size == 32:
bottleneck_size = 512
latent_dim = 128
kernel_size = 3
encoder_depths = [27, 27]
encoder_dims = [256, bottleneck_size]
decoder_input = (4, 2, bottleneck_size)
decoder_depths = [27, 27]
decoder_dims = [bottleneck_size, 256]
else:
raise NotImplementedError
self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[64, 128, 256, bottleneck_size], include_top=False)
self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
self.decoder = Decoder128(bottleneck_size, latent_dim)
self.encoder : AEEncoder = AEEncoder(latent_dim, encoder_depths, encoder_dims, bottleneck_size, kernel_size)
self.decoder : Decoder = Decoder(decoder_input, decoder_depths, decoder_dims, kernel_size)
def call(self, inputs, training=None):
def call(self, inputs, training=None, use_mask=True):
if use_mask:
x, mask = inputs
else:
x = inputs
x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,:, tf.newaxis]
# x = x * uv_mask[:,tf.newaxis,:,:]
if use_mask:
x = x * mask[:,:,:,tf.newaxis]
return x
# class Decoder32(tf.keras.Model):
# def __init__(self, bottleneck_size, latent_dim):
# super(Decoder32, self).__init__()
# self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
# self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
# self.decoder = ConvNeXtReverse(depths=list(reversed([27, 27])), dims=list(reversed([bottleneck_size, 256])))
# def call(self, inputs, training=None):
# x = inputs
# x = self.dense(x)
# x = self.reshape(x)
# x = self.decoder(x)
# return x
# class AutoEncoder32(tf.keras.Model):
# def __init__(self):
# super(AutoEncoder32, self).__init__()
# bottleneck_size = 512
# latent_dim = 128
# self.encoder = ConvNeXt(depths=[27, 27], dims=[256, bottleneck_size], include_top=False)
# self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
# self.decoder = Decoder32(bottleneck_size, latent_dim)
# def call(self, inputs, training=None):
# x = inputs
# x = tf.transpose(x, [0,2,3,1])
# x = self.encoder(x, training=training)
# x = self.bottleneck(x, training=training)
# x = self.decoder(x, training=training)
# x = tf.transpose(x, [0,3,1,2])
# # x = x * uv_mask[:,:, tf.newaxis]
# # x = x * uv_mask[:,tf.newaxis,:,:]
# return x
# class Decoder128(tf.keras.Model):
# def __init__(self, bottleneck_size, latent_dim):
# super(Decoder128, self).__init__()
# self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
# self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
# self.decoder = ConvNeXtReverse(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([bottleneck_size, 256, 128, 64])))
# def call(self, inputs, training=None):
# x = inputs
# x = self.dense(x)
# x = self.reshape(x)
# x = self.decoder(x)
# return x
# class AutoEncoder128(tf.keras.Model):
# def __init__(self):
# super(AutoEncoder128, self).__init__()
# bottleneck_size = 512
# latent_dim = 128
# self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[64, 128, 256, bottleneck_size], include_top=False)
# self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
# self.decoder = Decoder128(bottleneck_size, latent_dim)
# def call(self, inputs, training=None):
# x = inputs
# x = tf.transpose(x, [0,2,3,1])
# x = self.encoder(x, training=training)
# x = self.bottleneck(x, training=training)
# x = self.decoder(x, training=training)
# x = tf.transpose(x, [0,3,1,2])
# # x = x * uv_mask[:,:, tf.newaxis]
# # x = x * uv_mask[:,tf.newaxis,:,:]
# return x
models/base.py
+106
View File
@@ -0,0 +1,106 @@
import os
import tensorflow as tf
from custom_utils.custom_utils import load_model
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels
class Bottleneck(tf.keras.Model):
def __init__(self, dim, latent_dim, kernel_size=7):
super(Bottleneck, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0, kernel_size=kernel_size,
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
)
# self.bottleneck = [Block(dim=dim, drop_path=0, kernel_size=kernel_size, layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
self.flatten = tf.keras.layers.Flatten()
self.dense = tf.keras.layers.Dense(latent_dim)
def call(self, inputs, training=None):
x = inputs
# for bt in self.bottleneck:
# x = bt(x, training=training)
x = self.bottleneck(x)
x = self.flatten(x)
x = self.dense(x)
return x
class VariationalEncoder(tf.keras.Model):
def __init__(self, latent_dim, depths, dims, use_bottleneck=True, kernel_size=7):
super(VariationalEncoder, self).__init__()
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
self.flatten = tf.keras.layers.Flatten()
self.mean = tf.keras.layers.Dense(latent_dim)
self.log_var = tf.keras.layers.Dense(latent_dim)
def call(self, inputs, training=None):
x = inputs
x = self.encoder(x, training=training)
x = self.flatten(x)
mean = self.mean(x)
log_var = self.log_var(x)
return mean, log_var
class AEEncoder(tf.keras.Model):
def __init__(self, latent_dim, depths, dims, bottleneck_size=32, kernel_size=7):
super(AEEncoder, self).__init__()
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
self.bottleneck = Bottleneck(bottleneck_size, latent_dim, kernel_size)
def call(self, inputs, training=None):
if len(inputs) == 2:
x, _ = inputs
else:
x = inputs
x = self.encoder(x, training=training)
x = self.bottleneck(x, training=training)
return x
class Decoder(tf.keras.Model):
def __init__(self, input_size, depths, dims, kernel_size=7, is_label=False):
super(Decoder, self).__init__()
self.dense = tf.keras.layers.Dense(tf.reduce_prod(input_size))
self.reshape = tf.keras.layers.Reshape(input_size)
if not is_label:
self.decoder = ConvNeXtReverse(depths=list(reversed(depths)), dims=list(reversed(dims)), kernel_size=kernel_size)
else:
self.decoder = ConvNeXtReverseLabels(depths=list(reversed(depths)), dims=list(reversed(dims)), nr_of_classes=7)
def call(self, inputs, training=None):
if len(inputs) == 2:
x, mask = inputs
else:
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x, training=training)
if len(inputs) == 2:
if len(mask.shape) == 2:
x = x * mask[tf.newaxis,:,:,tf.newaxis]
else:
x = x * mask[:,:,:,tf.newaxis]
return x
models/cond_encoder.py
+20
View File
@@ -0,0 +1,20 @@
import tensorflow as tf
from models.base import AEEncoder, Decoder
class AutoEncoder(tf.keras.Model):
def __init__(self, config):
super(AutoEncoder, self).__init__()
self.encoder : AEEncoder = AEEncoder(config['latent_dim'], config['encoder_depths'], config['encoder_dims'], config['bottleneck_size'], config['kernel_size'])
self.decoder : Decoder= Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'])
def call(self, inputs, training=None):
x = inputs
x = self.encoder(x, training=training)
x = self.decoder(x, training=training)
return x
models/cvae_curv_model.py
+2 -2
View File
@@ -80,12 +80,12 @@ class CVAE(tf.keras.Model):
super(CVAE, self).__init__()
self.latent_dim = latent_dim
model : StaticAE = StaticAE()
model : StaticAE = StaticAE(size=512)
load_model(model, static_model)
self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
self.static_enc.trainable = False
model : BodyAE = BodyAE()
model : BodyAE = BodyAE(size=512)
load_model(model, body_model)
self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
self.body_enc.trainable = False
models/cvae_model.py
+46 -467
View File
@@ -9,119 +9,61 @@ import scipy as sp
from custom_utils.custom_utils import load_model
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
from models.ae_body_model import AutoEncoder as BodyAE
from models.ae_curv_model import AutoEncoder as CurvatureAE
from models.ae_static_model import AutoEncoder as StaticAE
from models.ae_body_model import AutoEncoder32 as BodyAE32
from models.ae_curv_model import AutoEncoder32 as CurvatureAE32
from models.ae_static_model import AutoEncoder32 as StaticAE32
from models.ae_body_model import AutoEncoder128 as BodyAE128
from models.ae_curv_model import AutoEncoder128 as CurvatureAE128
from models.ae_static_model import AutoEncoder128 as StaticAE128
from models.base import Bottleneck, VariationalEncoder, Decoder
from models.cond_encoder import AutoEncoder
# from models.ae_body_model import AutoEncoder as BodyAE
# from models.ae_curv_model import AutoEncoder as CurvatureAE
# from models.ae_static_model import AutoEncoder as StaticAE
# from models.ae_body_model import AutoEncoder32 as BodyAE32
# from models.ae_curv_model import AutoEncoder32 as CurvatureAE32
# from models.ae_static_model import AutoEncoder32 as StaticAE32
# from models.ae_body_model import AutoEncoder128 as BodyAE128
# from models.ae_curv_model import AutoEncoder128 as CurvatureAE128
# from models.ae_static_model import AutoEncoder128 as StaticAE128
class Bottleneck(tf.keras.Model):
def __init__(self, dim):
super(Bottleneck, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
)
def call(self, inputs, training=None):
x = inputs
x = self.bottleneck(x, training=training)
return x
class Encoder(tf.keras.Model):
def __init__(self, latent_dim):
super(Encoder, self).__init__()
self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, 256], include_top=False)
self.flatten = tf.keras.layers.Flatten()
self.mean = tf.keras.layers.Dense(latent_dim)
self.log_var = tf.keras.layers.Dense(latent_dim)
def call(self, inputs, training=None):
x = inputs
x = self.encoder(x)
x = self.flatten(x)
mean = self.mean(x)
log_var = self.log_var(x)
return mean, log_var
class Decoder(tf.keras.Model):
def __init__(self):
super(Decoder, self).__init__()
self.dense = tf.keras.layers.Dense(16 * 8 * 80)
self.reshape = tf.keras.layers.Reshape((16, 8, 80))
self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([256, 128, 64, 16])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class CVAE(tf.keras.Model):
def __init__(self, latent_dim=128,
avg_garment_path='',
static_model='',
body_model='',
curvature_model='',
freeze_encoders=False):
def __init__(self, config):
super(CVAE, self).__init__()
self.latent_dim = latent_dim
self.latent_dim = config['latent_dim']
self.size = config['resolution']
model : StaticAE = StaticAE()
if static_model != '':
load_model(model, static_model)
self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
model : AutoEncoder = AutoEncoder(config['cond_ae'])
if config['static_model'] != '':
load_model(model, config['static_model'])
# self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
self.static_enc = model.encoder
model : BodyAE = BodyAE()
if body_model != '':
load_model(model, body_model)
self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
model : AutoEncoder = AutoEncoder(config['cond_ae'])
if config['body_model'] != '':
load_model(model, config['body_model'])
self.body_enc = model.encoder
model : CurvatureAE = CurvatureAE()
if curvature_model != '':
load_model(model, curvature_model)
self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
model : AutoEncoder = AutoEncoder(config['cond_ae'])
if config['curvature_model'] != '':
load_model(model, config['curvature_model'])
self.curv_enc = model.encoder
if freeze_encoders:
if config['freeze_encoders']:
self.static_enc.trainable = False
self.body_enc.trainable = False
self.curv_enc.trainable = False
self.encoder = Encoder(latent_dim)
self.decoder = Decoder()
if avg_garment_path != '':
self.avg_garment = np.load(avg_garment_path)
self.encoder = VariationalEncoder(config['cvae']['latent_dim'], config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'])
self.decoder = Decoder(config['cvae']['decoder_input'], config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'])
if config['avg_garment_path'] != '':
self.avg_garment = np.load(config['avg_garment_path'])
else:
self.avg_garment = None
self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
tf.keras.layers.Dense(256, activation='relu'),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(latent_dim),])
tf.keras.layers.Dense(config['latent_dim']),])
# self.shuffle = tf.keras.layers.Lambda(lambda x: tf.stack([tf.gather(x[i], tf.random.shuffle(tf.range(x[i].shape[0]))) for i in range(x.shape[0])]))
@@ -212,30 +154,28 @@ class CVAE(tf.keras.Model):
def call(self, inputs, training=None, use_mask=False):
if use_mask:
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
else:
if self.size == 512:
uv_unposed, uv_template, uv_body_posed, uv_curv = inputs
uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
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_template_high, uv_body_posed_high, uv_curv_high = inputs
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
mean, logvar = self.encoder(x)
mean, logvar = self.encoder(x, training=training)
z = self.reparameterize(mean, logvar)
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
z_cond = self.encode_conds(uv_template_high, uv_body_posed_high, uv_curv_high)
z = tf.concat([z, z_cond], axis=-1)
x = self.decoder(z, training=training)
# z = self.shuffle(z)
# z = tf.concat([z, z_curv], axis=-1)
x = self.decoder(z)
# x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,tf.newaxis,:,:]
# x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# x = x + tf.transpose(uv_shaped_body, [0,2,3,1])
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
if use_mask:
@@ -243,364 +183,3 @@ class CVAE(tf.keras.Model):
return x, mean, logvar, z
class Encoder32(tf.keras.Model):
def __init__(self, latent_dim):
super(Encoder32, self).__init__()
self.encoder = ConvNeXt(depths=[27, 27], dims=[256, 256], include_top=False)
self.flatten = tf.keras.layers.Flatten()
self.mean = tf.keras.layers.Dense(latent_dim)
self.log_var = tf.keras.layers.Dense(latent_dim)
def call(self, inputs, training=None):
x = inputs
x = self.encoder(x)
x = self.flatten(x)
mean = self.mean(x)
log_var = self.log_var(x)
return mean, log_var
class Decoder32(tf.keras.Model):
def __init__(self):
super(Decoder32, self).__init__()
self.dense = tf.keras.layers.Dense(4 * 2 * 256)
self.reshape = tf.keras.layers.Reshape((4, 2, 256))
self.decoder = ConvNeXtReverse(depths=list(reversed([27, 27])), dims=list(reversed([256, 256])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class CVAE32(tf.keras.Model):
def __init__(self, latent_dim=128,
avg_garment_path='',
static_model='',
body_model='',
curvature_model='',
freeze_encoders=False,
use_full_res_encoders=True):
super(CVAE32, self).__init__()
self.latent_dim = latent_dim
if use_full_res_encoders:
model : StaticAE = StaticAE()
else:
model : StaticAE32 = StaticAE32()
if static_model != '':
load_model(model, static_model)
self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
if use_full_res_encoders:
model : BodyAE = BodyAE()
else:
model : BodyAE32 = BodyAE32()
if body_model != '':
load_model(model, body_model)
self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
if use_full_res_encoders:
model : CurvatureAE = CurvatureAE()
else:
model : CurvatureAE32 = CurvatureAE32()
if curvature_model != '':
load_model(model, curvature_model)
self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
if freeze_encoders:
self.static_enc.trainable = False
self.body_enc.trainable = False
self.curv_enc.trainable = False
self.encoder = Encoder32(latent_dim)
self.decoder = Decoder32()
if avg_garment_path != '':
self.avg_garment = np.load(avg_garment_path)
else:
self.avg_garment = None
self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
tf.keras.layers.Dense(256, activation='relu'),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(latent_dim),])
# @tf.function
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False, zero_latent=False):
if not fix_noise:
eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim))
else:
eps = tf.random.normal(shape=(1, self.latent_dim))
eps = tf.tile(eps, [nr_of_samples,1])
uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
if len(uv_curv.shape) == 3:
uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [nr_of_samples,1,1])
# uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
z_static = self.static_enc(uv_template)
z_body = self.body_enc(uv_body_posed)
z_curv = self.curv_enc(uv_curv)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder(z_cond)
if zero_latent:
eps = tf.zeros_like(eps)
z = tf.concat([eps, z_cond], axis=-1)
x = self.decoder(z)
x = x + self.avg_garment[tf.newaxis,:,:,:]
x = x * uv_mask[:,:,:,tf.newaxis]
return x
def reparameterize(self, mean, logvar):
eps = tf.random.normal(shape=mean.shape)
return eps * tf.exp(logvar * .5) + mean
def decode(self, z, apply_sigmoid=False):
logits = self.decoder(z)
if apply_sigmoid:
probs = tf.sigmoid(logits)
return probs
return logits
def call(self, inputs, training=None, zero_latent=False):
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_template_high, uv_body_posed_high, uv_curv_high = inputs
uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
z_static = self.static_enc(uv_template_high)
z_body = self.body_enc(uv_body_posed_high)
z_curv = self.curv_enc(uv_curv_high)
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
mean, logvar = self.encoder(x)
z = self.reparameterize(mean, logvar)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder(z_cond)
if zero_latent:
z = tf.zeros_like(z)
z = tf.concat([z, z_cond], axis=-1)
# z = self.shuffle(z)
# z = tf.concat([z, z_curv], axis=-1)
x = self.decoder(z)
# x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,tf.newaxis,:,:]
# x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# x = x + tf.transpose(uv_shaped_body, [0,2,3,1])
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
# x = x * uv_mask[:,:,:,tf.newaxis]
return x, mean, logvar, z
class Encoder128(tf.keras.Model):
def __init__(self, latent_dim):
super(Encoder128, self).__init__()
self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[256, 256, 128, 64], include_top=False)
self.flatten = tf.keras.layers.Flatten()
self.mean = tf.keras.layers.Dense(latent_dim)
self.log_var = tf.keras.layers.Dense(latent_dim)
def call(self, inputs, training=None):
x = inputs
x = self.encoder(x)
x = self.flatten(x)
mean = self.mean(x)
log_var = self.log_var(x)
return mean, log_var
class Decoder128(tf.keras.Model):
def __init__(self):
super(Decoder128, self).__init__()
self.dense = tf.keras.layers.Dense(4 * 2 * 128)
self.reshape = tf.keras.layers.Reshape((4, 2, 128))
self.decoder = ConvNeXtReverse(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([256, 256, 128, 64])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class CVAE128(tf.keras.Model):
def __init__(self, latent_dim=128,
avg_garment_path='',
static_model='',
body_model='',
curvature_model='',
freeze_encoders=False,
use_full_res_encoders=True):
super(CVAE128, self).__init__()
self.latent_dim = latent_dim
if use_full_res_encoders:
model : StaticAE = StaticAE()
else:
model : StaticAE128 = StaticAE128()
if static_model != '':
load_model(model, static_model)
self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
if use_full_res_encoders:
model : BodyAE = BodyAE()
else:
model : BodyAE128 = BodyAE128()
if body_model != '':
load_model(model, body_model)
self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
if use_full_res_encoders:
model : CurvatureAE = CurvatureAE()
else:
model : CurvatureAE128 = CurvatureAE128()
if curvature_model != '':
load_model(model, curvature_model)
self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
if freeze_encoders:
self.static_enc.trainable = False
self.body_enc.trainable = False
self.curv_enc.trainable = False
self.encoder = Encoder128(latent_dim)
self.decoder = Decoder128()
if avg_garment_path != '':
self.avg_garment = np.load(avg_garment_path)
else:
self.avg_garment = None
self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
tf.keras.layers.Dense(256, activation='relu'),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(latent_dim),])
# @tf.function
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False, zero_latent=False):
if not fix_noise:
eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim))
else:
eps = tf.random.normal(shape=(1, self.latent_dim))
eps = tf.tile(eps, [nr_of_samples,1])
uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
if len(uv_curv.shape) == 3:
uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [nr_of_samples,1,1])
# uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
z_static = self.static_enc(uv_template)
z_body = self.body_enc(uv_body_posed)
z_curv = self.curv_enc(uv_curv)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder(z_cond)
if zero_latent:
eps = tf.zeros_like(eps)
z = tf.concat([eps, z_cond], axis=-1)
x = self.decoder(z)
x = x + self.avg_garment[tf.newaxis,:,:,:]
x = x * uv_mask[:,:,:,tf.newaxis]
return x
def reparameterize(self, mean, logvar):
eps = tf.random.normal(shape=mean.shape)
return eps * tf.exp(logvar * .5) + mean
def decode(self, z, apply_sigmoid=False):
logits = self.decoder(z)
if apply_sigmoid:
probs = tf.sigmoid(logits)
return probs
return logits
def call(self, inputs, training=None, zero_latent=False):
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_template_high, uv_body_posed_high, uv_curv_high = inputs
uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
z_static = self.static_enc(uv_template_high)
z_body = self.body_enc(uv_body_posed_high)
z_curv = self.curv_enc(uv_curv_high)
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
mean, logvar = self.encoder(x)
z = self.reparameterize(mean, logvar)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder(z_cond)
if zero_latent:
z = tf.zeros_like(z)
z = tf.concat([z, z_cond], axis=-1)
# z = self.shuffle(z)
# z = tf.concat([z, z_curv], axis=-1)
x = self.decoder(z)
# x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,tf.newaxis,:,:]
# x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# x = x + tf.transpose(uv_shaped_body, [0,2,3,1])
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
# x = x * uv_mask[:,:,:,tf.newaxis]
return x, mean, logvar, z
models/cvae_model_old.py
+621
View File
@@ -0,0 +1,621 @@
import os
import time
import random
import tensorflow as tf
import numpy as np
import scipy as sp
# import plotly.graph_objects as go
# import plotly.express as ex
from custom_utils.custom_utils import load_model
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
from models.ae_body_model import AutoEncoder as BodyAE
from models.ae_curv_model import AutoEncoder as CurvatureAE
from models.ae_static_model import AutoEncoder as StaticAE
from models.ae_body_model import AutoEncoder32 as BodyAE32
from models.ae_curv_model import AutoEncoder32 as CurvatureAE32
from models.ae_static_model import AutoEncoder32 as StaticAE32
from models.ae_body_model import AutoEncoder128 as BodyAE128
from models.ae_curv_model import AutoEncoder128 as CurvatureAE128
from models.ae_static_model import AutoEncoder128 as StaticAE128
class Bottleneck(tf.keras.Model):
def __init__(self, dim):
super(Bottleneck, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
)
def call(self, inputs, training=None):
x = inputs
x = self.bottleneck(x, training=training)
return x
class Encoder(tf.keras.Model):
def __init__(self, latent_dim):
super(Encoder, self).__init__()
self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, 256], include_top=False)
self.flatten = tf.keras.layers.Flatten()
self.mean = tf.keras.layers.Dense(latent_dim)
self.log_var = tf.keras.layers.Dense(latent_dim)
def call(self, inputs, training=None):
x = inputs
x = self.encoder(x)
x = self.flatten(x)
mean = self.mean(x)
log_var = self.log_var(x)
return mean, log_var
class Decoder(tf.keras.Model):
def __init__(self):
super(Decoder, self).__init__()
self.dense = tf.keras.layers.Dense(16 * 8 * 80)
self.reshape = tf.keras.layers.Reshape((16, 8, 80))
self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([256, 128, 64, 16])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class CVAE(tf.keras.Model):
def __init__(self, latent_dim=128,
avg_garment_path='',
static_model='',
body_model='',
curvature_model='',
freeze_encoders=False):
super(CVAE, self).__init__()
self.latent_dim = latent_dim
model : StaticAE = StaticAE()
if static_model != '':
load_model(model, static_model)
self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
model : BodyAE = BodyAE()
if body_model != '':
load_model(model, body_model)
self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
model : CurvatureAE = CurvatureAE()
if curvature_model != '':
load_model(model, curvature_model)
self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
if freeze_encoders:
self.static_enc.trainable = False
self.body_enc.trainable = False
self.curv_enc.trainable = False
self.encoder = Encoder(latent_dim)
self.decoder = Decoder()
if avg_garment_path != '':
self.avg_garment = np.load(avg_garment_path)
else:
self.avg_garment = None
self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
tf.keras.layers.Dense(256, activation='relu'),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(latent_dim),])
# self.shuffle = tf.keras.layers.Lambda(lambda x: tf.stack([tf.gather(x[i], tf.random.shuffle(tf.range(x[i].shape[0]))) for i in range(x.shape[0])]))
# @tf.function
def sample_single(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False):
if not fix_noise:
eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim))
else:
eps = tf.random.normal(shape=(1, self.latent_dim))
eps = tf.tile(eps, [nr_of_samples,1])
uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
if len(uv_curv.shape) == 3:
uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [nr_of_samples,1,1])
# uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
z_static = self.static_enc(uv_template)
z_body = self.body_enc(uv_body_posed)
z_curv = self.curv_enc(uv_curv)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder(z_cond)
z = tf.concat([eps, z_cond], axis=-1)
x = self.decoder(z)
x = x + self.avg_garment[tf.newaxis,:,:,:]
x = x * uv_mask[:,:,:,tf.newaxis]
return x
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask=None, noise=None):
if noise is None:
eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
else:
eps = noise
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
z = tf.concat([eps, z_cond], axis=-1)
x = self.decoder(z)
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
if uv_mask is not None:
x = x * uv_mask[:,:,:,tf.newaxis]
return x
def reparameterize(self, mean, logvar, fix_noise=False):
if not fix_noise:
eps = tf.random.normal(shape=mean.shape)
else:
eps = tf.random.normal(shape=(1, mean.shape[1]))
eps = tf.tile(eps, [mean.shape[0],1])
return eps * tf.exp(logvar * .5) + mean
def decode(self, z, apply_sigmoid=False):
logits = self.decoder(z)
if apply_sigmoid:
probs = tf.sigmoid(logits)
return probs
return logits
def encode_conds(self, uv_template, uv_body_posed, uv_curv):
if uv_template.shape[1] == 3:
uv_template = tf.transpose(uv_template, [0,2,3,1])
if uv_body_posed.shape[1] == 3:
uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1])
if uv_curv.shape[1] == 3:
uv_curv = tf.transpose(uv_curv, [0,2,3,1])
z_static = self.static_enc(uv_template)
z_body = self.body_enc(uv_body_posed)
z_curv = self.curv_enc(uv_curv)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder(z_cond)
return z_cond
def call(self, inputs, training=None, use_mask=False):
if use_mask:
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
else:
uv_unposed, uv_template, uv_body_posed, uv_curv = inputs
uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
mean, logvar = self.encoder(x)
z = self.reparameterize(mean, logvar)
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
z = tf.concat([z, z_cond], axis=-1)
# z = self.shuffle(z)
# z = tf.concat([z, z_curv], axis=-1)
x = self.decoder(z)
# x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,tf.newaxis,:,:]
# x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# x = x + tf.transpose(uv_shaped_body, [0,2,3,1])
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
if use_mask:
x = x * uv_mask[:,:,:,tf.newaxis]
return x, mean, logvar, z
class Encoder32(tf.keras.Model):
def __init__(self, latent_dim):
super(Encoder32, self).__init__()
self.encoder = ConvNeXt(depths=[27, 27], dims=[256, 256], include_top=False)
self.flatten = tf.keras.layers.Flatten()
self.mean = tf.keras.layers.Dense(latent_dim)
self.log_var = tf.keras.layers.Dense(latent_dim)
def call(self, inputs, training=None):
x = inputs
x = self.encoder(x)
x = self.flatten(x)
mean = self.mean(x)
log_var = self.log_var(x)
return mean, log_var
class Decoder32(tf.keras.Model):
def __init__(self):
super(Decoder32, self).__init__()
self.dense = tf.keras.layers.Dense(4 * 2 * 256)
self.reshape = tf.keras.layers.Reshape((4, 2, 256))
self.decoder = ConvNeXtReverse(depths=list(reversed([27, 27])), dims=list(reversed([256, 256])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class CVAE32(tf.keras.Model):
def __init__(self, latent_dim=128,
avg_garment_path='',
static_model='',
body_model='',
curvature_model='',
freeze_encoders=False,
use_full_res_encoders=True):
super(CVAE32, self).__init__()
self.latent_dim = latent_dim
if use_full_res_encoders:
model : StaticAE = StaticAE()
else:
model : StaticAE32 = StaticAE32()
if static_model != '':
load_model(model, static_model)
self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
if use_full_res_encoders:
model : BodyAE = BodyAE()
else:
model : BodyAE32 = BodyAE32()
if body_model != '':
load_model(model, body_model)
self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
if use_full_res_encoders:
model : CurvatureAE = CurvatureAE()
else:
model : CurvatureAE32 = CurvatureAE32()
if curvature_model != '':
load_model(model, curvature_model)
self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
if freeze_encoders:
self.static_enc.trainable = False
self.body_enc.trainable = False
self.curv_enc.trainable = False
self.encoder = Encoder32(latent_dim)
self.decoder = Decoder32()
if avg_garment_path != '':
self.avg_garment = np.load(avg_garment_path)
else:
self.avg_garment = None
self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
tf.keras.layers.Dense(256, activation='relu'),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(latent_dim),])
# @tf.function
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False, zero_latent=False):
if not fix_noise:
eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim))
else:
eps = tf.random.normal(shape=(1, self.latent_dim))
eps = tf.tile(eps, [nr_of_samples,1])
uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
if len(uv_curv.shape) == 3:
uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [nr_of_samples,1,1])
# uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
z_static = self.static_enc(uv_template)
z_body = self.body_enc(uv_body_posed)
z_curv = self.curv_enc(uv_curv)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder(z_cond)
if zero_latent:
eps = tf.zeros_like(eps)
z = tf.concat([eps, z_cond], axis=-1)
x = self.decoder(z)
x = x + self.avg_garment[tf.newaxis,:,:,:]
x = x * uv_mask[:,:,:,tf.newaxis]
return x
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask=None, noise=None):
if noise is None:
eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
else:
eps = noise
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
z = tf.concat([eps, z_cond], axis=-1)
x = self.decoder(z)
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
if uv_mask is not None:
x = x * uv_mask[:,:,:,tf.newaxis]
return x
def reparameterize(self, mean, logvar):
eps = tf.random.normal(shape=mean.shape)
return eps * tf.exp(logvar * .5) + mean
def encode_conds(self, uv_template, uv_body_posed, uv_curv):
if uv_template.shape[1] == 3:
uv_template = tf.transpose(uv_template, [0,2,3,1])
if uv_body_posed.shape[1] == 3:
uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1])
if uv_curv.shape[1] == 3:
uv_curv = tf.transpose(uv_curv, [0,2,3,1])
z_static = self.static_enc(uv_template)
z_body = self.body_enc(uv_body_posed)
z_curv = self.curv_enc(uv_curv)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder(z_cond)
return z_cond
def call(self, inputs, training=None, use_mask=False):
if use_mask:
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
else:
uv_unposed, uv_template, uv_body_posed, uv_curv = inputs
uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
mean, logvar = self.encoder(x)
z = self.reparameterize(mean, logvar)
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
z = tf.concat([z, z_cond], axis=-1)
x = self.decoder(z)
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
if use_mask:
x = x * uv_mask[:,:,:,tf.newaxis]
return x, mean, logvar, z
class Encoder128(tf.keras.Model):
def __init__(self, latent_dim):
super(Encoder128, self).__init__()
self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[256, 256, 128, 64], include_top=False)
self.flatten = tf.keras.layers.Flatten()
self.mean = tf.keras.layers.Dense(latent_dim)
self.log_var = tf.keras.layers.Dense(latent_dim)
def call(self, inputs, training=None):
x = inputs
x = self.encoder(x)
x = self.flatten(x)
mean = self.mean(x)
log_var = self.log_var(x)
return mean, log_var
class Decoder128(tf.keras.Model):
def __init__(self):
super(Decoder128, self).__init__()
self.dense = tf.keras.layers.Dense(4 * 2 * 128)
self.reshape = tf.keras.layers.Reshape((4, 2, 128))
self.decoder = ConvNeXtReverse(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([256, 256, 128, 64])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class CVAE128(tf.keras.Model):
def __init__(self, latent_dim=128,
avg_garment_path='',
static_model='',
body_model='',
curvature_model='',
freeze_encoders=False,
use_full_res_encoders=True):
super(CVAE128, self).__init__()
self.latent_dim = latent_dim
if use_full_res_encoders:
model : StaticAE = StaticAE()
else:
model : StaticAE128 = StaticAE128()
if static_model != '':
load_model(model, static_model)
self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
if use_full_res_encoders:
model : BodyAE = BodyAE()
else:
model : BodyAE128 = BodyAE128()
if body_model != '':
load_model(model, body_model)
self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
if use_full_res_encoders:
model : CurvatureAE = CurvatureAE()
else:
model : CurvatureAE128 = CurvatureAE128()
if curvature_model != '':
load_model(model, curvature_model)
self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
if freeze_encoders:
self.static_enc.trainable = False
self.body_enc.trainable = False
self.curv_enc.trainable = False
self.encoder = Encoder128(latent_dim)
self.decoder = Decoder128()
if avg_garment_path != '':
self.avg_garment = np.load(avg_garment_path)
else:
self.avg_garment = None
self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
tf.keras.layers.Dense(256, activation='relu'),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(latent_dim),])
# @tf.function
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False, zero_latent=False):
if not fix_noise:
eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim))
else:
eps = tf.random.normal(shape=(1, self.latent_dim))
eps = tf.tile(eps, [nr_of_samples,1])
uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
if len(uv_curv.shape) == 3:
uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [nr_of_samples,1,1])
# uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
z_static = self.static_enc(uv_template)
z_body = self.body_enc(uv_body_posed)
z_curv = self.curv_enc(uv_curv)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder(z_cond)
if zero_latent:
eps = tf.zeros_like(eps)
z = tf.concat([eps, z_cond], axis=-1)
x = self.decoder(z)
x = x + self.avg_garment[tf.newaxis,:,:,:]
x = x * uv_mask[:,:,:,tf.newaxis]
return x
def reparameterize(self, mean, logvar):
eps = tf.random.normal(shape=mean.shape)
return eps * tf.exp(logvar * .5) + mean
def decode(self, z, apply_sigmoid=False):
logits = self.decoder(z)
if apply_sigmoid:
probs = tf.sigmoid(logits)
return probs
return logits
def call(self, inputs, training=None, zero_latent=False):
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_template_high, uv_body_posed_high, uv_curv_high = inputs
uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
z_static = self.static_enc(uv_template_high)
z_body = self.body_enc(uv_body_posed_high)
z_curv = self.curv_enc(uv_curv_high)
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
mean, logvar = self.encoder(x)
z = self.reparameterize(mean, logvar)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder(z_cond)
if zero_latent:
z = tf.zeros_like(z)
z = tf.concat([z, z_cond], axis=-1)
# z = self.shuffle(z)
# z = tf.concat([z, z_curv], axis=-1)
x = self.decoder(z)
# x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,tf.newaxis,:,:]
# x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# x = x + tf.transpose(uv_shaped_body, [0,2,3,1])
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
# x = x * uv_mask[:,:,:,tf.newaxis]
return x, mean, logvar, z
models/cvae_pyramid_model.py
+16 -16
View File
@@ -7,7 +7,7 @@ import numpy as np
# import plotly.express as ex
from custom_utils.custom_utils import load_model
from models.cvae_model import CVAE, CVAE128, CVAE32
from models.cvae_model import CVAE
from models.upscaler import UpscalerBasic as Upscaler
@@ -24,26 +24,22 @@ class Pyramid(tf.keras.Model):
use_full_res_encoders=True):
super(Pyramid, self).__init__()
if low_res == 32:
self.model_low = CVAE32(latent_dim=latent_dim,
self.model_low = CVAE(latent_dim=latent_dim,
size=low_res,
static_model=static_model,
body_model=body_model,
curvature_model=curvature_model,
freeze_encoders=freeze_encoders,
use_full_res_encoders=use_full_res_encoders)
if low_res == 128:
self.model_low = CVAE128(latent_dim=latent_dim,
static_model=static_model,
body_model=body_model,
curvature_model=curvature_model,
freeze_encoders=freeze_encoders,
use_full_res_encoders=use_full_res_encoders)
encoder_res=512 if use_full_res_encoders else 32)
self.model_full = CVAE(latent_dim=latent_dim,
size=512,
static_model=static_model,
body_model=body_model,
curvature_model=curvature_model,
freeze_encoders=freeze_encoders)
freeze_encoders=freeze_encoders,
encoder_res=512)
self.upscale_enhancer = None
if upscaler_model != '':
@@ -64,6 +60,12 @@ class Pyramid(tf.keras.Model):
else:
self.avg_garment = None
def enhance(self, low, mask, training=None):
x = self.upscale_enhancer(low, training=training)
x = x * (1 - mask[:,:,:,tf.newaxis])
low = x + low
return low
def call(self, inputs, training=None):
uv_unposed, uv_static, uv_body_posed, uv_curv, uv_mask, uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low, uv_mask_low = inputs
@@ -75,11 +77,9 @@ class Pyramid(tf.keras.Model):
pred_low = pred_low * uv_mask_low[:,:,:,tf.newaxis]
if self.upscale_enhancer is not None:
x = self.upscale_enhancer(pred_low, training=training)
x = x * (1 - uv_mask_low[:,:,:,tf.newaxis])
pred_low = x + pred_low
pred_low = self.enhance(pred_low, uv_mask_low, training=training)
x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
x = x * uv_mask[:,:,:,tf.newaxis]
pred_full, mean_full, logvar_full, _ = self.model_full([uv_unposed, uv_static, uv_body_posed, uv_curv], training=training)
sample_pro.py
+12 -12
View File
@@ -1,5 +1,5 @@
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -27,6 +27,7 @@ from PIL import Image
from plotly.subplots import make_subplots
from custom_utils.smpl import SMPL
from custom_utils.custom_utils import load_model
from cvae_curv import CVAE as CVAECurv
from models.cvae_model import CVAE as CVAEMain
@@ -136,12 +137,6 @@ def plot_everything(batch: dict, type: str):
return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
def load_model(model, checkpoint):
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
ckpt.restore(manager.latest_checkpoint)
def get_frame(folder, nr):
pose_folder = os.path.join(folder, 'pose_shape')
frames = sorted(os.listdir(pose_folder))
@@ -391,6 +386,7 @@ def sample(folder, variable_condition, fix_noise, nr=4):
assert False
plot_variable_condition(variable_cond_data, mesh_mask, is_curv=random_curv, is_body=random_body)
plot_variable_condition(curv, mesh_mask, is_curv=True, is_body=False)
plot_results(mesh_pred_shaped, mesh_posed, body, mesh_mask)
@@ -407,7 +403,7 @@ if __name__ == '__main__':
args = vars(parser.parse_args())
folder = '/home/workspace/local-data/train_100'
folder_main = 'checkpoints/cvae_13'
folder_main = 'checkpoints/cvae_18_copy'
folder_curv = 'checkpoints/cvae_curv_3'
variable_condition = 'static'
batch_size = 4
@@ -415,13 +411,17 @@ if __name__ == '__main__':
smpl = SMPL()
model = CVAEMain(avg_garment_path='custom_data/train_100_no_shape_avg_uv.npy',
latent_dim=128, static_model='checkpoints/ae_static_4',
body_model='checpoints/ae_body_8',
curvature_model='checpoints/ae_curvature_7')
curv_model = CVAECurv()
latent_dim=128, static_model='checkpoints/ae_static_6_copy',
body_model='checkpoints/ae_body_8_copy',
curvature_model='checkpoints/ae_curvature_7_copy',
freeze_encoders=True)
curv_model = CVAECurv(static_model='checkpoints/ae_static_6_copy',
body_model='checkpoints/ae_body_8_copy',)
load_model(model, folder_main)
load_model(curv_model, folder_curv)
sample(folder, variable_condition, fix_noise=True)
print()
vae.py
+1 -48
View File
@@ -22,7 +22,7 @@ from stylegan2_ae.generator import Generator
from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder, sparse_to_tensor
from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from custom_utils.custom_utils import batch_to_dict, plot_3d
from custom_utils.smpl import SMPL
name='vae_11'
@@ -38,53 +38,6 @@ run.log_artifact(artifact)
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
class Bottleneck(tf.keras.Model):
verify_upscaler.py
+1 -1
View File
@@ -2,7 +2,7 @@
import os
from models.upscaler import UpscalerBasic as Upscaler
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
import tqdm
import numpy as np