cvpr/auto_encoder_curv.py

import os
import sys
from configs.config_ae_curv_1_orig_curv_10_10p_r2 import config
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']

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 cmapy as cm


from custom_utils.custom_utils import UV2Mesh, UV2MeshLabel, get_summary_html, maskFacesBatch
import plotly.graph_objects as go
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 models.cond_encoder import AutoEncoder


testing = False
def init_wandb():
    if not testing:
        run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp/', tags=config['tags'])
        artifact = wandb.Artifact('code', type='code')
        artifact.add_file('auto_encoder_curv.py')
        artifact.add_file('losses.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')
        config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
        for config_file_name in config_file_names:
            file_name = '/'.join(config_file_name.split('.')) + '.py'
            if os.path.exists(file_name):
                artifact.add_file(file_name)
        run.log_artifact(artifact)

uv2mesh = UV2MeshLabel(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])


g_mean_iou = tf.keras.metrics.MeanIoU(num_classes=7)
def get_mean_iou(y_true, y_pred):
    g_mean_iou.reset_state()
    g_mean_iou.update_state(y_true, y_pred)
    return g_mean_iou.result()

g_mean_iou_mesh = tf.keras.metrics.MeanIoU(num_classes=7)
def get_mean_iou_mesh(y_true, y_pred):
    g_mean_iou_mesh.reset_state()
    g_mean_iou_mesh.update_state(y_true, y_pred)
    return g_mean_iou_mesh.result()

# https://github.com/NifTK/NiftyNet/blob/dev/niftynet/layer/loss_segmentation.py
def dice_dense(prediction, ground_truth, weight_map=None):
    """
    Computing mean-class Dice similarity.
    :param prediction: last dimension should have ``num_classes``
    :param ground_truth: segmentation ground truth (encoded as a binary matrix)
        last dimension should be ``num_classes``
    :param weight_map:
    :return: ``1.0 - mean(Dice similarity per class)``
    """

    if weight_map is not None:
        raise NotImplementedError
    prediction = tf.cast(prediction, dtype=tf.float32)
    ground_truth = tf.cast(ground_truth, dtype=tf.float32)
    ground_truth = tf.reshape(ground_truth, prediction.shape)
    # computing Dice over the spatial dimensions
    reduce_axes = list(range(len(prediction.shape) - 1))
    dice_numerator = 2.0 * tf.reduce_sum(
        prediction * ground_truth, axis=reduce_axes)
    dice_denominator = \
        tf.reduce_sum(tf.square(prediction), axis=reduce_axes) + \
        tf.reduce_sum(tf.square(ground_truth), axis=reduce_axes)

    epsilon = 0.00001

    dice_score = (dice_numerator + epsilon) / (dice_denominator + epsilon)
    return tf.reduce_mean(dice_score)


@tf.function
def train_step(batch):	
    with tf.GradientTape() as tape:	
        pred = model([batch['uv_curvature'], batch['uv_mask']], training=True)	
        ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred)	
        ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1))	
        loss = ce_uv_loss	
    pred_mesh = tf.math.argmax(uv2mesh(pred), axis=-1)	
    mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1))	
    mean_iou_mesh = get_mean_iou(tf.cast(batch['mesh_curvature'], tf.float32) * batch['mesh_mask'], tf.cast(pred_mesh, tf.float32) * batch['mesh_mask'])	
    dice = dice_dense(pred, batch['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/mean_iou_mesh': mean_iou_mesh, 'train/dice': dice }	
        	
@tf.function	
def valid_step(batch):	
    pred = model([batch['uv_curvature'], batch['uv_mask']], training=False)	
    ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred)	
    ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1))	
    loss = ce_uv_loss	
    pred_mesh = tf.math.argmax(uv2mesh(pred), axis=-1)	
    mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1))	
    mean_iou_mesh = get_mean_iou(tf.cast(batch['mesh_curvature'], tf.float32) * batch['mesh_mask'], tf.cast(pred_mesh, tf.float32) * batch['mesh_mask'])	
    dice = dice_dense(pred, batch['uv_curvature'])	
    return { 'valid/loss': loss, 'valid/ce_uv_loss': ce_uv_loss, 'valid/ssim': ms_ssim, 'valid/mean_iou': mean_iou, 'valid/mean_iou_mesh': mean_iou_mesh, 'valid/dice': dice }

# def train_step(batch):
#     with tf.GradientTape() as tape:
#         pred = model([batch['uv_curvature'], batch['uv_mask']], training=True)

#         ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred)
#         ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1))

#         mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,tf.newaxis]
#         mesh_orig = tf.cast(uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.int32)), tf.float32) * batch['mesh_mask']
#         ce_mesh_loss = scce(mesh_orig, mesh_pred)

#         loss = ce_uv_loss + ce_mesh_loss

#     mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1))
#     dice = dice_dense(pred, batch['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, 'train/ce_mesh_loss': ce_mesh_loss }
        


# @tf.function
# def valid_step(batch):
#     pred = model([batch['uv_curvature'], batch['uv_mask']], training=False)

#     ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred)
#     ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1))

#     mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,tf.newaxis]
#     mesh_orig = tf.cast(uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.int32)), tf.float32) * batch['mesh_mask']
#     ce_mesh_loss = scce(mesh_orig, mesh_pred)
#     loss = ce_uv_loss + ce_mesh_loss

#     mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1))
#     dice = dice_dense(pred, batch['uv_curvature'])

#     return { 'valid/loss': loss, 'valid/ce_uv_loss': ce_uv_loss, 'valid/ssim': ms_ssim, 'valid/mean_iou': mean_iou,
#              'valid/dice': dice, 'valid/ce_mesh_loss': ce_mesh_loss }
        

def plot_everything(batch: dict, type: str):
    pred = model([batch['uv_curvature'], batch['uv_mask']], training=False)
    
    orig_img = tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.dtypes.int32)
    pred_img = tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)
    img = tf.concat((orig_img, pred_img), axis=2).numpy()

    colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, 6+1)])
    color_img = np.zeros((img.shape[0], img.shape[1], img.shape[2], 3), np.uint8)
    for i in range(7):
        color_img[img == i, :] = colors[i]
    color_img = [wandb.Image(color_img[i]) for i in range(4)]

    imgs = { type + '/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(img.shape[0])], type + '/color_images': color_img }

    return 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, 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, config['signature_dynamic']))
        valid_loss = valid_step(data_valid)

        to_log = { **train_loss, **valid_loss }
        if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
            train_figs = plot_everything(batch, 'train')
            valid_figs = plot_everything(data_valid, 'valid')
            to_log = { **to_log, **train_figs, **valid_figs }
        if step == 0:
            to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
 
        wandb.log(to_log, step=step)  
        print(step)

        if step % 10000 == 0:
            manager_enc.save(checkpoint_number=step)
            manager_dec.save(checkpoint_number=step)

        step += 1


def test():
    #TODO
    # ckpt.restore(manager.latest_checkpoint)
    for batch in dataset_train:
        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']])
        orig_img = batch['uv_curvature']
        pred_img = tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)
        img = tf.concat((orig_img, pred_img), axis=2).numpy()

        ex.imshow(img[0] / 6.0).show()
        print()



if __name__ == '__main__':
    init_wandb()
    # allow_memory_growth()
    parser = argparse.ArgumentParser(description='')
    parser.add_argument('--name', type=str)
    args = vars(parser.parse_args())

    scce = tf.keras.losses.SparseCategoricalCrossentropy()

    resolution_modifier = ''
    if config['resolution'] != 512:
        resolution_modifier = '_' + str(config['resolution'])

    factory = DatasetFactory()

    dataset_train = factory.get_dataset(folder=config['train_dir'],
                                        batch_size=config['batch_size'], 
                                        signature=config['signature_dynamic'],
                                        resolutions=[config['resolution']], 
                                        in_memory=config['dataset_in_memory'],
                                        use_no_orient=config['remove_global_orientation'],
                                        frames_file=config['frames_file'] if 'frames_file' in config.keys() else None, 
                                        is_cape=config['is_cape'],
                                        curv_neighbourhood=config['curv_neighbourhood'])
    static_dataset_train = StaticDataset(folder=config['train_dir'],
                                         signature=config['signature_static'], 
                                         resolutions=[config['resolution']], 
                                         in_memory=True, 
                                        is_cape=config['is_cape'])

    dataset_valid = factory.get_dataset(folder=config['valid_dir'],
                                        batch_size=config['batch_size'], 
                                        signature=config['signature_dynamic'], 
                                        resolutions=[config['resolution']], 
                                        in_memory=config['dataset_in_memory'], 
                                        use_no_orient=config['remove_global_orientation'], 
                                        frames_file=config['frames_file_valid'] if 'frames_file_valid' in config.keys() else None, 
                                        is_cape=config['is_cape'],
                                        curv_neighbourhood=config['curv_neighbourhood'])
    static_dataset_valid = StaticDataset(folder=config['valid_dir'],
                                         signature=config['signature_static'], 
                                         resolutions=[config['resolution']], 
                                         in_memory=True, 
                                        is_cape=config['is_cape'])

    model = AutoEncoder(config['cond_ae'])
    optimizer = tf.optimizers.Adam(learning_rate=config['lr'])

    ckpt_enc = tf.train.Checkpoint(model=model.encoder)
    manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/post/' + config['name'] + '_enc', max_to_keep=20)
    ckpt_dec = tf.train.Checkpoint(model=model.decoder)
    manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/post/' + config['name'] + '_dec', max_to_keep=20)

    if testing:
        test()
    else:
        train()