working pyramid update

convnext/models/convnext_tf.py

@@ -45,10 +45,10 @@ class Block(Layer):
         layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
     """
 
-    def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6, prefix=''):
+    def __init__(self, dim, kernel_size=7, drop_path=0., layer_scale_init_value=1e-6, prefix=''):
         super().__init__()
         self.dwconv = DepthwiseConv2D(
-            kernel_size=7, padding='same')  # depthwise conv
+            kernel_size=kernel_size, padding='same')  # depthwise conv
         self.norm = LayerNormalization(epsilon=1e-6)
         # pointwise/1x1 convs, implemented with linear layers
         self.pwconv1 = Dense(4 * dim)
@@ -97,7 +97,7 @@ class ConvNeXt(tf.keras.Model):
         head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
     """
 
-    def __init__(self, num_classes=1000,
+    def __init__(self, num_classes=1000, kernel_size=7,
                  depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], include_top=True,
                  drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
                  ):
@@ -121,7 +121,7 @@ class ConvNeXt(tf.keras.Model):
         cur = 0
         for i in range(len(depths)):
             stage = tf.keras.Sequential(
-                [Block(dim=dims[i], drop_path=dp_rates[cur + j],
+                [Block(dim=dims[i], drop_path=dp_rates[cur + j], kernel_size=kernel_size,
                        layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
             )
             self.stages.append(stage)
@@ -167,7 +167,7 @@ class ConvNeXtWithSkip(tf.keras.Model):
     """
 
     def __init__(self, 
-                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], 
+                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], kernel_size=7, 
                  drop_path_rate=0., layer_scale_init_value=1e-6, 
                  ):
         super().__init__()
@@ -189,7 +189,7 @@ class ConvNeXtWithSkip(tf.keras.Model):
         cur = 0
         for i in range(len(depths)):
             stage = tf.keras.Sequential(
-                [Block(dim=dims[i], drop_path=dp_rates[cur + j],
+                [Block(dim=dims[i], drop_path=dp_rates[cur + j], kernel_size=kernel_size,
                        layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
             )
             self.stages.append(stage)
@@ -221,7 +221,7 @@ class ConvNeXtReverse(tf.keras.Model):
     """
 
     def __init__(self, 
-                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
+                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], kernel_size=7,
                  drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
                  ):
         super().__init__()
@@ -249,7 +249,7 @@ class ConvNeXtReverse(tf.keras.Model):
         cur = 0
         for i in range(len(depths)):
             stage = tf.keras.Sequential(
-                [Block(dim=dims[i], drop_path=dp_rates[cur + j],
+                [Block(dim=dims[i], drop_path=dp_rates[cur + j], kernel_size=kernel_size,
                        layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
             )
             self.stages.append(stage)
@@ -281,7 +281,7 @@ class ConvNeXtReverseWithSkip(tf.keras.Model):
     """
 
     def __init__(self, 
-                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
+                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], kernel_size=7,
                  drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
                  ):
         super().__init__()
@@ -304,7 +304,7 @@ class ConvNeXtReverseWithSkip(tf.keras.Model):
         cur = 0
         for i in range(len(depths)):
             stage = tf.keras.Sequential(
-                [Block(dim=dims[i], drop_path=dp_rates[cur + j],
+                [Block(dim=dims[i], kernel_size=kernel_size, drop_path=dp_rates[cur + j],
                        layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
             )
             self.stages.append(stage)

cvae_low.py

@@ -0,0 +1,498 @@
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+import tensorflow as tf
+
+gpus = tf.config.experimental.list_physical_devices('GPU')
+if gpus:
+    try:
+        # Currently, memory growth needs to be the same across GPUs
+        for gpu in gpus:
+            tf.config.experimental.set_memory_growth(gpu, True)
+        logical_gpus = tf.config.experimental.list_logical_devices('GPU')
+        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
+    except RuntimeError as e:
+        # Memory growth must be set before GPUs have been initialized
+        print(e)
+
+import time
+import argparse
+import wandb
+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
+
+
+name='cvae_low_1_d128'
+def init_wandb():
+    run = wandb.init(project='vae', name=name,
+            notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. fixed balancing of reconst and kl loss for baseline. \
+new implementation for kl loss. no lr schedule. 25 jumpsuits, average reconst_loss before calculation, linearly increasing balancing, up to 1 from 0.1. \
+offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \
+retrained body encoder with deoriented bodies \
+use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.01\
+fixed gender issue and retrained body encoder\
+remove shape from unposed garment\
+training on 100 garments , absolute, not offsets\
+using filled versions')
+    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('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)
+    res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis)
+    return res
+
+
+def kl_loss(mean, logvar):
+    kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar)
+    kl_loss = tf.reduce_mean(kl_loss)
+    kl_loss *= -0.5
+    return kl_loss
+
+
+@tf.function
+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 = 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)
+        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]
+
+        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)
+
+        reconstr_loss = l2_loss
+        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)
+        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' + 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'])
+    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/reg_coeff': reg_coeff, 'train/l2_loss': l2_loss }
+        
+
+
+@tf.function
+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_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']))
+    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])
+    b_loss = 0 #1e-3 * bend_loss(mesh_pred, smpl_faces, laplacian_loss, batch['mesh_mask'])
+    # 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)
+    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])
+    # 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 = 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
+    ms_ssim = tf.reduce_mean(ms_ssim)
+    l1_loss = tf.reduce_mean(l1_loss)
+    fft_loss = tf.reduce_mean(fft_loss)
+    b_loss = tf.reduce_mean(b_loss)
+    reconstr_loss = tf.reduce_mean(reconstr_loss)
+    mean = tf.reduce_mean(mean)
+    logvar = tf.reduce_mean(logvar)
+
+    return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim,
+             'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2, 'train/bend_loss': b_loss, 'train/reconstr_loss': reconstr_loss,
+             'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar }
+
+
+def plot_everything(batch: dict, type: str):
+    uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+    uv_curvature_low = tf.one_hot(batch['uv_curvature' + 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 = pred * batch['uv_mask' + res_modifier][:,:,:,tf.newaxis]
+    
+    mesh_pred = uv2mesh(pred)
+    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
+    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])]
+    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))]
+
+    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) }
+    
+    mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0), num_bins=model.latent_dim)
+    logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0), num_bins=model.latent_dim)
+
+    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)
+    count = 0
+    for batch in dataset_train:
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+        uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+        nr_of_samples=2
+
+        pred = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask' + res_modifier][0], nr_of_samples=nr_of_samples)
+        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_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]
+
+        body = batch['mesh_body_posed'][0].numpy()
+        body_color = np.ones_like(body, dtype=np.int32)*100
+
+        F = np.load("custom_utils/smpl_faces.npy")
+        mask = batch['mesh_mask'].numpy().astype(bool)[0]
+        _, faces = _maskFaces(np.ones((14475, 3)), F, mask)
+
+
+        pred_zero = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask' + res_modifier][0], nr_of_samples=nr_of_samples, zero_latent=True)
+        mesh_pred_zero = uv2mesh(tf.transpose(pred_zero, [0,3,1,2]))
+        mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
+
+
+
+
+        meshes = []
+        meshes.append(go.Mesh3d(
+                    # 8 vertices of a cube
+                    x=mesh_orig_reposed[mask,0],
+                    y=mesh_orig_reposed[mask,2],
+                    z=mesh_orig_reposed[mask,1],
+                    # i, j and k give the vertices of triangles
+                    i = faces[:,0],
+                    j = faces[:,1],
+                    k = faces[:,2]
+                    ))
+
+        for i in range(nr_of_samples):
+            x = mesh_posed[i][mask, :].copy()
+            # x[:, 2] += 1 * (i + 1)
+            x[:, 2] += 1
+            meshes.append(go.Mesh3d(
+                        # 8 vertices of a cube
+                        x=x[:,0],
+                        y=x[:,2],
+                        z=x[:,1],
+                        # i, j and k give the vertices of triangles
+                        i = faces[:,0],
+                        j = faces[:,1],
+                        k = faces[:,2]
+                        ))
+
+        
+        for i in range(nr_of_samples):
+            x = mesh_posed_zero[i][mask, :].copy()
+            # x[:, 2] += 1 * (i + 1)
+            x[:, 2] += 2
+            meshes.append(go.Mesh3d(
+                        # 8 vertices of a cube
+                        x=x[:,0],
+                        y=x[:,2],
+                        z=x[:,1],
+                        # i, j and k give the vertices of triangles
+                        i = faces[:,0],
+                        j = faces[:,1],
+                        k = faces[:,2]
+                        ))
+
+        # for i in range(5):
+        #     x = body.copy()
+        #     x[:, 2] += 1 * i
+        #     meshes.append(go.Mesh3d(
+        #                 # 8 vertices of a cube
+        #                 x=x[:,0],
+        #                 y=x[:,2],
+        #                 z=x[:,1],
+        #                 # i, j and k give the vertices of triangles
+        #                 i = F[:,0],
+        #                 j = F[:,1],
+        #                 k = F[:,2],
+        #                 vertexcolor=body_color
+        #                 ))
+
+
+        fig = go.Figure(data=meshes)
+
+        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1))
+        # fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1))
+        fig.update_layout(
+            scene = dict(
+                xaxis = dict(nticks=4, range=[-1,1],),
+                            # yaxis = dict(nticks=4, range=[-1,nr_of_samples+1],),
+                            yaxis = dict(nticks=4, range=[-1,3],),
+                            zaxis = dict(nticks=4, range=[-1,1],),)
+                            )
+
+        # wandb.log({'plot_' + str(count): fig})
+        fig.show()
+
+        pred = pred.numpy()
+        frames = [Image.fromarray(((pred[i] + 1) * 128).astype(np.int8), mode='RGB') for i in range(pred.shape[0])]
+        frames[0].save("temp/garment_1.gif", format="GIF", append_images=frames[1:],
+               save_all=True, duration=500, loop=0)
+
+
+        if count > 10:
+            break
+
+        print(count)
+        count += 1
+
+        print()
+
+if __name__ == '__main__':
+    init_wandb()
+    allow_memory_growth()
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+
+    batch_size = 4
+    
+    smpl = SMPL()
+
+    resolution = 128
+    res_modifier = ''
+    if resolution != 512:
+        res_modifier = '_' + str(resolution)
+
+    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,),
+                          '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_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,),
+                        #  'mesh_static': (14475, 3,),
+                         'mesh_mask': (14475,),
+                         'uv_mask' + res_modifier: (resolution, resolution//2,),
+                        #  '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=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_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',
+                 freeze_encoders=True,
+                 use_full_res_encoders=True)
+
+    # 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
+    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)
+
+    step = 0
+    for batch in dataset_train:
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, 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)
+        valid_loss = {}
+
+        to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
+        if step % 2500 == 0:
+            train_figs = plot_everything(batch, 'train')
+            # valid_figs = plot_everything(data_valid, 'valid')
+            valid_figs = {}
+            to_log = { **to_log, **train_figs, **valid_figs }
+
+        if step % 10000 == 0:
+            manager.save(checkpoint_number=step)
+
+        if train_loss['train/euclid_error'] < 0.01:
+            if balancing_activated_step == 0:
+                balancing_activated_step = step
+            balancing_coeff = balancing_coeff + 1 / 15000
+
+        # gamma = 0.99 * gamma + 0.01 * train_loss['train/reconstr_loss']        
+
+        to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } }
+        
+        wandb.log(to_log, step = step)
+
+        print(step)
+
+        step += 1

cvae_pyramid.py

@@ -1,5 +1,5 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+os.environ["CUDA_VISIBLE_DEVICES"] = "2"
 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_2_d128'
+name='cvae_pyr_3_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,11 +49,13 @@ fixed gender issue and retrained body encoder\
 remove shape from unposed garment\
 training on 100 garments with offsets over the average of these\
 l2 loss on both 128 and full res, bottom predicts base, high res predicts offsets over this\
-low and high res both use same full res condition encoders, but low res also gets low res versions for main encoder input')
+low and high res both use same full res condition encoders, but low res also gets low res versions for main encoder input\
+not using filled in versions, but using custom MLP upscaling')
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae_pyramid.py')
     artifact.add_file('models/cvae_model.py')
     artifact.add_file('models/cvae_pyramid_model.py')
+    artifact.add_file('models/upscaler.py')
     artifact.add_file('dataset.py')
     artifact.add_file('losses.py')
     artifact.add_file('convnext/models/convnext_tf.py')
@@ -135,8 +137,8 @@ def train_step(batch):
         uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
         uv_curvature_low = tf.one_hot(batch['uv_curvature' + resolution_modifier], 7)
 
-        pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, 
-                                                                                 batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low])
+        pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'],
+                                                                                 batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low, batch['uv_mask' + resolution_modifier],])
 
         mesh_pred_full = uv2mesh(pred_full)
         mesh_pred_posed_full = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_full)
@@ -152,11 +154,12 @@ 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) * 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_full) ** 2), axis=-1) * batch['uv_mask'], axis=[1, 2])
+        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)
         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])
@@ -236,8 +239,8 @@ def valid_step(batch):
 def plot_everything(batch: dict, type: str):
     uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
     uv_curvature_low = tf.one_hot(batch['uv_curvature' + resolution_modifier], 7)
-    pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, 
-                                                                             batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low])
+    pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'],
+                                                                             batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low, batch['uv_mask' + resolution_modifier]])
 
     mesh_pred_full = uv2mesh(pred_full)
     mesh_pred_low = uv2mesh(pred_low)
@@ -409,7 +412,7 @@ if __name__ == '__main__':
     args = vars(parser.parse_args())
 
     batch_size = 4
-    low_resolution = 128
+    low_resolution = 32
     resolution_modifier = '_' + str(low_resolution)
     
     smpl = SMPL()
@@ -433,11 +436,12 @@ if __name__ == '__main__':
                         #  '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,),
                         }
 
-    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=True, resolutions=[512, low_resolution])
-    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=True, resolutions=[512, low_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, low_resolution])
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=[512, low_resolution])
     dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=batch_size, signature=signature_dynamic)
     static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
 
@@ -447,7 +451,8 @@ if __name__ == '__main__':
                     freeze_encoders=True, 
                     static_model='checkpoints/ae_static_4',
                     body_model='checpoints/ae_body_7',
-                    curvature_model='checpoints/ae_curvature_7')
+                    curvature_model='checpoints/ae_curvature_7',
+                    upscaler_model='checkpoints/up_2_3')
 
     # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)
     optimizer = tf.optimizers.Adam(learning_rate=0.0001)

models/cvae_model.py

@@ -358,11 +358,11 @@ class CVAE32(tf.keras.Model):
 
 
     def call(self, inputs, training=None, zero_latent=False):
-        uv_unposed, uv_template, uv_body_posed, uv_curv = inputs
+        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)
-        z_body = self.body_enc(uv_body_posed)
-        z_curv = self.curv_enc(uv_curv)
+        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)
 

models/cvae_pyramid_model.py

@@ -8,6 +8,7 @@ import numpy as np
 
 from custom_utils.custom_utils import load_model
 from models.cvae_model import CVAE, CVAE128, CVAE32
+from models.upscaler import UpscalerBasic as Upscaler
 
 
 class Pyramid(tf.keras.Model):
@@ -18,6 +19,7 @@ class Pyramid(tf.keras.Model):
                  static_model='',
                  body_model='',
                  curvature_model='',
+                 upscaler_model='',
                  freeze_encoders=False,
                  use_full_res_encoders=True):
         super(Pyramid, self).__init__()
@@ -43,6 +45,18 @@ class Pyramid(tf.keras.Model):
                                curvature_model=curvature_model, 
                                freeze_encoders=freeze_encoders)
 
+        self.upscale_enhancer = None
+        if upscaler_model != '':
+            self.upscale_enhancer = Upscaler(512, low_res)
+            load_model(self.upscale_enhancer, upscaler_model)
+            self.upscale_enhancer = tf.keras.Sequential([
+                self.upscale_enhancer.flatten,
+                self.upscale_enhancer.seq,
+                self.upscale_enhancer.reshape
+            ])
+
+            self.upscale_enhancer.trainable = False
+        
         self.use_full_res_encoders = use_full_res_encoders
 
         if avg_garment_path != '':
@@ -51,18 +65,27 @@ class Pyramid(tf.keras.Model):
             self.avg_garment = None
 
     def call(self, inputs, training=None):
-        uv_unposed, uv_static, uv_body_posed, uv_curv, uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low = inputs
+        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
 
         if self.use_full_res_encoders:
             pred_low, mean_low, logvar_low, _ = self.model_low([uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low, uv_static, uv_body_posed, uv_curv], training=training)
         else:
             pred_low, mean_low, logvar_low, _ = self.model_low([uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low], training=training)
 
-        x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+        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
+        
+        x = tf.image.resize(x, (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)
 
         x = x + pred_full
+        x = x * uv_mask[:,:,:,tf.newaxis]
         
         return x, mean_full, logvar_full, pred_low, mean_low, logvar_low,