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

update curvature CVAE

Browse Source
This commit is contained in:
luni
2022-09-22 11:26:07 +00:00
parent aa1e7be0e6
commit c5ae37aacc
3 changed files with 46 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
convnext/models/convnext_tf.py
+1 -1
View File
@@ -369,7 +369,7 @@ class ConvNeXtReverseLabels(tf.keras.Model):
self.last = Dense(nr_of_classes, activation='softmax')
def call(self, x):
for i in reversed(len(depths)):
for i in reversed(range(len(self.downsample_layers))):
x = self.downsample_layers[i](x)
x = self.stages[i](x)
x = self.last(x)
cvae_curv.py
+17 -20
View File
@@ -1,5 +1,5 @@
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -26,10 +26,7 @@ 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 tf_utils import check_tf_version, allow_memory_growth
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from dataset import DatasetFactory, StaticDataset
@@ -38,14 +35,15 @@ from custom_utils.smpl import SMPL
from models.cvae_curv_model import CVAE
name='cvae_curv_2_2_2'
name='cvae_curv_3'
def init_wandb():
run = wandb.init(project='vae', name=name,
notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. new implementation for kl loss. \
no lr schedule. 25 jumpsuits, average reconst_loss before calculation, offsets over average unposed garment. gamma balancing. remove global \
orientation (in posed body input and in pose_shape parameters.. retrained body encoder with deoriented bodies. lower kl coeff to 0.1. fixed plotting colors\
fixed deoriented bodies and need retraining (ae_body_6 with 50k iteration), addent MLP for lowering conditional Z diemnsion\
higher starting kl coeff (0.1)')
orientation (in posed body input and in pose_shape parameters.. r lower kl coeff to 0.1. fixed plotting colors\
addent MLP for lowering conditional Z diemnsion\
higher starting kl coeff (0.1)\
train on 100 garments')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('cvae_curv.py')
artifact.add_file('models/cvae_curv_model.py')
@@ -161,7 +159,7 @@ 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_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
pred, mean, logvar, z = model([batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']])
ce_uv_loss = scce(batch['uv_curvature'], pred)
ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1))
@@ -197,7 +195,7 @@ def train_step(batch):
def valid_step(batch):
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
pred, mean, logvar, z = model([batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
pred, mean, logvar, z = model([batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']])
ce_uv_loss = scce(batch['uv_curvature'], pred)
ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1))
@@ -220,7 +218,7 @@ def valid_step(batch):
def plot_everything(batch: dict, type: str):
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
pred, mean, logvar, z = model([batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
pred, mean, logvar, z = model([batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']])
pred = np.argmax(pred.numpy(), axis=-1)
mesh_pred = uv2mesh(pred)
@@ -336,7 +334,7 @@ def sample(folder):
print()
if __name__ == '__main__':
# init_wandb()
init_wandb()
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
@@ -363,17 +361,17 @@ if __name__ == '__main__':
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
'uv_shaped_body': (3, 512, 256,),
# 'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', batch_size=batch_size, signature=signature_dynamic)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static)
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=False, resolutions=[512])
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=[512])
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 = CVAE(latent_dim=128, static_model='checkpoints/ae_static_3',
body_model='checpoints/ae_body_6',)
model = CVAE(latent_dim=128, static_model='checkpoints/ae_static_4',
body_model='checpoints/ae_body_8',)
# lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)
@@ -384,14 +382,13 @@ if __name__ == '__main__':
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_curv_2')
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=5)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3)
step = 0
for batch in dataset_train:
models/cvae_curv_model.py
+28 -13
View File
@@ -122,21 +122,33 @@ class CVAE(tf.keras.Model):
return x
def sample(self, uv_template, uv_body_posed, uv_mask, eps=None):
eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
uv_template = tf.transpose(uv_template, [0,2,3,1])
uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1])
def encode_conds(self, uv_template, uv_body_posed):
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])
z_static = self.static_enc(uv_template)
z_body = self.body_enc(uv_body_posed)
z_cond = tf.concat([z_static, z_body], axis=-1)
z_cond = self.cond_encoder(z_cond)
return z_cond
def sample(self, uv_template, uv_body_posed, 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)
z = tf.concat([eps, z_cond], axis=-1)
x = self.decoder(z)
x = x * uv_mask[:,:,:,tf.newaxis]
if uv_mask is not None:
x = x * uv_mask[:,:,:,tf.newaxis]
return x
@@ -168,17 +180,20 @@ class CVAE(tf.keras.Model):
def call(self, inputs, training=None):
uv_template, uv_body_posed, uv_curv, uv_mask, uv_shaped_body = inputs
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)
uv_template, uv_body_posed, uv_curv, uv_mask = inputs
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_cond = self.encode_conds(uv_template, uv_body_posed)
x = tf.concat([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], axis=-1)
z_cond = self.cond_encoder(z_cond)
z = tf.concat([z, z_cond], axis=-1)
x = self.decoder(z)