From fafa28cc69f32246d3cdefe52eaa74dc0daad436 Mon Sep 17 00:00:00 2001
From: Hunor Laczko <luni741@gmail.com>
Date: Tue, 20 Sep 2022 22:53:57 +0000
Subject: [PATCH] custom upscaling

---
 models/upscaler.py | 107 +++++++++++++++++++++++
 upscale_test.py    | 214 +++++++++++++++++++++++++++++++++++++++++++++
 verify_upscaler.py | 102 +++++++++++++++++++++
 3 files changed, 423 insertions(+)
 create mode 100644 models/upscaler.py
 create mode 100644 upscale_test.py
 create mode 100644 verify_upscaler.py

diff --git a/models/upscaler.py b/models/upscaler.py
new file mode 100644
index 0000000..e8a457a
--- /dev/null
+++ b/models/upscaler.py
@@ -0,0 +1,107 @@
+import tensorflow as tf
+from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseWithSkip, ConvNeXtWithSkip
+
+
+class Bottleneck(tf.keras.Model):
+    
+    def __init__(self, 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)]
+            )
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.bottleneck(x, training=training)
+        
+        return x
+
+
+
+class Upscaler(tf.keras.Model):
+    
+    def __init__(self, bottleneck_size):
+        super(Upscaler, self).__init__()
+
+        depths = [27, 27]
+        dims = [512, 512]
+        reverse_dims = [16] + dims[:-1]
+        self.encoder = ConvNeXtWithSkip(depths=depths, dims=dims, kernel_size=3)
+
+        self.bottleneck = Bottleneck(dims[-1], kernel_size=3)
+
+        self.decoder = ConvNeXtReverseWithSkip(depths=depths, dims=reverse_dims, kernel_size=3)
+
+    def call(self, inputs, training=None):
+        img, mask = inputs
+
+        img = tf.transpose(img, [0,2,3,1])
+
+        x, skip_features = self.encoder(img, training=training)
+
+        x = self.bottleneck(x, training=training)
+
+        x = self.decoder([x, skip_features], training=training)
+
+        x = x * (1 - mask[:,:,:,tf.newaxis])
+        x = x + img
+        x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
+
+        return x
+
+
+class UpscalerBasic(tf.keras.Model):
+    
+    def __init__(self, bottleneck_size, size=32):
+        super(UpscalerBasic, self).__init__()
+
+        self.flatten = tf.keras.layers.Flatten()
+
+        self.seq = tf.keras.Sequential([
+                    tf.keras.layers.Dense(1024, activation='relu'),
+                    tf.keras.layers.BatchNormalization(),
+                    tf.keras.layers.Dense(1024, activation='relu'),
+                    tf.keras.layers.BatchNormalization(),
+                    tf.keras.layers.Dense(1024, activation='relu'),
+                    tf.keras.layers.BatchNormalization(),
+                    tf.keras.layers.Dense(size * (size//2) * 3),
+        ])
+
+        self.reshape = tf.keras.layers.Reshape((size, (size//2), 3))
+
+    def call(self, inputs, training=None):
+        img, mask = inputs
+
+        if img.shape[1] == 3:
+            img = tf.transpose(img, [0,2,3,1])
+
+        x = self.flatten(img)
+        x = self.seq(x, training=training)
+        x = self.reshape(x)
+
+        x = x * (1 - mask[:,:,:,tf.newaxis])
+        x = x + img
+        x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
+
+        return x
+
+# class Upscaler(tf.keras.Model):
+    
+#     def __init__(self, bottleneck_size):
+#         super(Upscaler, self).__init__()
+
+#         self.decoder = ConvNeXtReverse(depths=list(reversed([9, 9, 9])), dims=list(reversed([bottleneck_size, 128, 32, 16])))
+
+#     def call(self, inputs, training=None):
+#         x, mask = inputs
+
+#         x = tf.transpose(x, [0,2,3,1])
+#         x = self.decoder(x)
+
+#         x = x * mask[:,:,:,tf.newaxis]
+
+#         return x
+
diff --git a/upscale_test.py b/upscale_test.py
new file mode 100644
index 0000000..67a2ae0
--- /dev/null
+++ b/upscale_test.py
@@ -0,0 +1,214 @@
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "2"
+
+import time
+import argparse
+import wandb
+import random
+import cv2
+import numpy as np
+import tensorflow as tf
+from custom_utils.custom_utils import UV2Mesh, 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.upscaler import Upscaler
+from models.upscaler import UpscalerBasic as Upscaler
+
+name='up_2_3_d128'
+def init_wandb():
+    run = wandb.init(project='vae', name=name,
+            notes='convnext upscaling\
+changed ms-ssim scale weights\
+inpainting style\
+MLP based\
+smaller batch\
+use mask for loss\
+    ')
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('upscale_test.py')
+    artifact.add_file('models/upscaler.py')
+    artifact.add_file('losses.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):
+    mesh_pred_unposed[:,2] += 1
+    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],
+        )
+    ])
+
+    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_unposed' + resolution_modifier], batch['uv_mask' + resolution_modifier]])
+
+        mesh_pred = uv2mesh(pred)
+        mesh_orig = batch['mesh_unposed']
+
+        l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred * batch['uv_mask'][:,:,:,tf.newaxis]) )), axis=1))
+        ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred * batch['uv_mask'][:,:,:,tf.newaxis] + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
+        # ms_ssim = tf.reduce_mean(tf.image.ssim(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))
+        loss = 1e3 * (1 - ms_ssim)
+
+    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'])
+    gradients = tape.gradient(loss, model.trainable_variables)
+    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
+    return { 'train/loss': loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim }
+        
+
+
+# UPDATE!
+@tf.function
+def valid_step(batch):
+    pred = model(batch['uv_unposed' + resolution_modifier], batch['uv_mask'])
+
+    mesh_pred = uv2mesh(pred)
+    mesh_orig = batch['mesh_unposed']
+
+    l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_unposed'] - pred) )), axis=1))
+    ms_ssim = tf.reduce_mean(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.reduce_mean(tf.image.ssim(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))
+    loss = 1e3 * (1 - ms_ssim)
+
+    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'])
+    return { 'valid/loss': loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim }
+   
+
+
+def plot_everything(batch: dict, type: str):
+
+    pred = model([batch['uv_unposed' + resolution_modifier], batch['uv_mask' + resolution_modifier]])
+    
+    mesh_pred = uv2mesh(pred)
+    mesh_orig = batch['mesh_unposed']
+    
+    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])]
+    gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
+    
+    figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[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) }
+    
+    orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
+    img = tf.concat((orig_img, pred), axis=2)
+    imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
+
+    return { **figs, **{ type + '/images': imgs} }
+
+
+
+def train():
+    valid_iterator = iter(dataset_valid)
+
+    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)
+
+        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.save(checkpoint_number=step)
+
+
+
+        print(step)
+
+        step += 1
+
+if __name__ == '__main__':
+    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)
+
+    low_res = 128
+    uv_mask_body = tf.ones((low_res, low_res//2))
+    resolution_modifier = '_' + str(low_res)
+
+    factory = DatasetFactory()
+    signature_dynamic = { 'id': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_unposed' + resolution_modifier: (3, low_res, low_res//2,),
+                          'uv_unposed': (3, 512, 256,),
+                          'mesh_unposed': (14475, 3,),
+                        }
+    signature_static = { 
+                         'mesh_mask': (14475,),
+                         'uv_mask': (512, 256,),
+                         'uv_mask' + resolution_modifier: (low_res, low_res//2,),
+                        }
+
+    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=8, signature=signature_dynamic, use_filled=False, resolutions=[512, low_res])
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=[512, low_res])
+    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)
+
+    model = Upscaler(bottleneck_size=512, size=low_res)
+
+    optimizer = tf.optimizers.Adam(learning_rate=0.0001)
+    folder = 'models/ae_logs/'
+
+    ckpt = tf.train.Checkpoint(model=model)
+    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3)
+    ckpt.restore(manager.latest_checkpoint)
+
+    train()
\ No newline at end of file
diff --git a/verify_upscaler.py b/verify_upscaler.py
new file mode 100644
index 0000000..5167c2c
--- /dev/null
+++ b/verify_upscaler.py
@@ -0,0 +1,102 @@
+# check if upscaler is better then filling
+import os
+
+from models.upscaler import UpscalerBasic as Upscaler
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+
+import tqdm
+import numpy as np
+import plotly.graph_objects as go
+
+from custom_utils.custom_utils import UV2Mesh, _maskFaces, load_model
+# from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces
+
+
+if __name__ == '__main__':
+    folder = '/home/workspace/local-data/train_100'
+    mask = np.load('custom_utils/mask_256x512.npy')
+    F = np.load("custom_utils/smpl_faces.npy")
+    uv2mesh = UV2Mesh(res=256)
+    res = 32
+    res_modifier = '_' + str(res)
+
+    upscale_enhancer = Upscaler(512, res)
+    load_model(upscale_enhancer, 'checkpoints/up_2_3')
+
+    frames = os.listdir(os.path.join(folder, 'pose_shape'))
+    visited = set()
+    meshes = []
+    count = 0
+
+    for frame in tqdm.tqdm(frames):
+
+        id, frame_nr = frame.split('_')
+        frame_nr = frame_nr.split('.')[0]
+        if id in visited:
+            continue
+        visited.add(id)
+        print(id)
+        uv_unposed_no_shape_low = np.load(os.path.join(folder, 'uv_unposed_no_shape' + res_modifier, id + '_' + frame_nr + '.npz'));  uv_unposed_no_shape_low = uv_unposed_no_shape_low['uv_unposed_no_shape']
+        uv_unposed_no_shape_filled_low = np.load(os.path.join(folder, 'uv_unposed_no_shape_filled' + res_modifier, id + '_' + frame_nr + '.npz'));  uv_unposed_no_shape_filled_low = uv_unposed_no_shape_filled_low['uv_unposed_no_shape_filled']
+        mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy'))
+        uv_mask = np.load(os.path.join(folder, 'uv_mask' + res_modifier, id + '.npz')); uv_mask = uv_mask['uv_mask']
+        mesh = uv2mesh(uv_unposed_no_shape_low[np.newaxis,:,:,:])[0].numpy()
+        mesh_filled = uv2mesh(uv_unposed_no_shape_filled_low[np.newaxis,:,:,:])[0].numpy()
+        mesh_orig = np.load(os.path.join(folder, 'mesh_unposed_no_shape', id + '_' + frame_nr + '.npz')); mesh_orig = mesh_orig['mesh_unposed_no_shape']
+
+        uv_enhanced = upscale_enhancer([uv_unposed_no_shape_low[np.newaxis,:,:,:], uv_mask[np.newaxis,:,:].astype(np.float32)], training=False)
+        mesh_enhanced = uv2mesh(uv_enhanced)[0].numpy()
+
+        metric = np.sum(np.sqrt(np.sum(mesh_mask[:,np.newaxis] * (mesh_orig - mesh) ** 2, axis=-1))) / np.sum(mesh_mask)
+        metric_filled = np.sum(np.sqrt(np.sum(mesh_mask[:,np.newaxis] * (mesh_orig - mesh_filled) ** 2, axis=-1))) / np.sum(mesh_mask)
+        metric_enhanced = np.sum(np.sqrt(np.sum(mesh_mask[:,np.newaxis] * (mesh_orig - mesh_enhanced) ** 2, axis=-1))) / np.sum(mesh_mask)
+        print(metric)
+        print(metric_filled)
+        print(metric_enhanced)
+        print()
+
+        mesh, faces = _maskFaces(mesh, F, mesh_mask)
+        mesh_filled, _ = _maskFaces(mesh_filled, F, mesh_mask)
+        mesh_enhanced, _ = _maskFaces(mesh_enhanced, F, mesh_mask)
+        x = mesh
+        y = mesh_filled
+        z = mesh_enhanced
+        y[:, 2] += 0.5
+        z[:, 2] += 1
+        meshes = [go.Mesh3d(
+                    x=x[:,0],
+                    y=x[:,2],
+                    z=x[:,1],
+                    i = faces[:,0],
+                    j = faces[:,1],
+                    k = faces[:,2]
+                    ),
+                    go.Mesh3d(
+                    x=y[:,0],
+                    y=y[:,2],
+                    z=y[:,1],
+                    i = faces[:,0],
+                    j = faces[:,1],
+                    k = faces[:,2]
+                    ),
+                    go.Mesh3d(
+                    x=z[:,0],
+                    y=z[:,2],
+                    z=z[:,1],
+                    i = faces[:,0],
+                    j = faces[:,1],
+                    k = faces[:,2]
+                    )]
+
+        fig = go.Figure(data=meshes)
+        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],),)
+                            )
+
+        fig.show()
+
+        print()
\ No newline at end of file