diff --git a/convnext/models/convnext_tf.py b/convnext/models/convnext_tf.py index 6a36bb1..e2781c9 100644 --- a/convnext/models/convnext_tf.py +++ b/convnext/models/convnext_tf.py @@ -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) diff --git a/cvae_curv.py b/cvae_curv.py index 146b885..3e23054 100644 --- a/cvae_curv.py +++ b/cvae_curv.py @@ -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: diff --git a/models/cvae_curv_model.py b/models/cvae_curv_model.py index 0a0e211..07fc82b 100644 --- a/models/cvae_curv_model.py +++ b/models/cvae_curv_model.py @@ -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)