diff --git a/custom_data/garment_uv_mask_full.npy b/custom_data/garment_uv_mask_full.npy new file mode 100644 index 0000000..c79c38e Binary files /dev/null and b/custom_data/garment_uv_mask_full.npy differ diff --git a/data/dataset.py b/data/dataset.py index 188f1b2..e64dde9 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -30,14 +30,15 @@ def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file= if is_cape: frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0] + frames = [Frame(frame.split('.')[0], frame.split('.')[1]) for frame in frames] else: frames = [frame.split('.')[0] for frame in frames] # .split('_')[0] + frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames] - if testing: - frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') )) - frames = [frame.split('.')[0] + '_4' for frame in frames] + # if testing: + # frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') )) + # frames = [frame.split('.')[0] + '_4' for frame in frames] - frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames] if frames_file is not None: frames = [] @@ -65,12 +66,12 @@ class StaticCache(): if is_cape: self.uv_masks = dict() - uv_mask = np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32) - self.uv_masks[512] = uv_mask.astype(np.float32) + uv_mask = torch.tensor(np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32), device=device) + self.uv_masks[512] = uv_mask.to(torch.float32).cpu().numpy() for res in [32, 64, 128, 256]: - down = Resize(uv_mask[:,:,None], size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST) - down = np.squeeze(down.numpy().astype(np.float32)) - self.uv_masks[res] = down + down = Resize(size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST)(uv_mask[None,:,:]) + down = torch.squeeze(down).to(torch.float32) + self.uv_masks[res] = down.cpu().numpy() def read(self, id): @@ -127,7 +128,7 @@ class Dataset(Dataset): self.frames = read_frames(folder, frames_file=frames_file, is_cape=is_cape, shuffle=True) self.body_faces = np.load('smpl/smpl_faces.npy') - self.static_cache = StaticCache(folder, signature, is_cape=is_cape) + self.static_cache = StaticCache(folder, signature, is_cape=is_cape, resolutions=resolutions, device=device) self.cache = dict() if use_cache else None diff --git a/models/upscaler.py b/models/upscaler.py index c03e663..b857ea3 100644 --- a/models/upscaler.py +++ b/models/upscaler.py @@ -1,172 +1,35 @@ -import tensorflow as tf -from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseWithSkip, ConvNeXtWithSkip +import torch +import torch.nn as nn - -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): - x = inputs - - x = self.bottleneck(x, training=training) - - return x - - - -class Upscaler(tf.keras.Model): - - def __init__(self, bottleneck_size): +class Upscaler(nn.Module): + def __init__(self, low_res=32, kernel_size=3): 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): - 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, low_res=32, high_res=64): - super(UpscalerBasic, self).__init__() - - self.low_res = low_res - self.high_res = high_res - 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(low_res * (low_res//2) * 3), - ]) - - self.reshape = tf.keras.layers.Reshape((low_res, (low_res//2), 3)) - - def call(self, inputs, training): - img, mask = inputs - - 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=(self.high_res, self.high_res//2), method=tf.image.ResizeMethod.LANCZOS5) - - return x - - - - -class UpscalerBasicNew(tf.keras.Model): - - def __init__(self, low_res=32): - super(UpscalerBasicNew, self).__init__() - - self.low_res = low_res - self.flatten = tf.keras.layers.Flatten(name='upscale_enhancer_flattent') - - self.seq = tf.keras.Sequential([ - tf.keras.layers.Dense(4096, activation='relu', name='upscale_enhancer_dense_1'), - tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_1'), - tf.keras.layers.Dense(2048, activation='relu', name='upscale_enhancer_dense_2'), - tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_2'), - tf.keras.layers.Dense(2048, activation='relu', name='upscale_enhancer_dense_3'), - tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_3'), - tf.keras.layers.Dense(low_res * (low_res//2) * 3, activation='tanh', name='upscale_enhancer_dense_final'), - ]) - - self.reshape = tf.keras.layers.Reshape((low_res, (low_res//2), 3), name='upscale_enhancer_reshape') - - def call(self, inputs, training): - img, mask = inputs - - x = self.flatten(img) - x = self.seq(x, training=training) - x = self.reshape(x) - - x = x * (1 - mask[:,:,:,tf.newaxis]) - x = x + img - - return x - - -class UpscalerBasicConv(tf.keras.Model): - - def __init__(self, low_res=32, kernel_size=3): - super(UpscalerBasicConv, self).__init__() - self.low_res = low_res - self.seq = tf.keras.Sequential([ - tf.keras.layers.Conv2D(256, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_1'), - tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_1'), - tf.keras.layers.Conv2D(128, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_2'), - tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_2'), - tf.keras.layers.Conv2D(64, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_3'), - tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_3'), - tf.keras.layers.Conv2D(32, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_4'), - tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_4'), - tf.keras.layers.Dense(3, activation='tanh', name='upscale_enhancer_dense_final'), - ]) + self.seq = nn.Sequential( + nn.Conv2d(3, 256, kernel_size=kernel_size, padding='same'), + nn.ReLU(), + nn.BatchNorm2d(256), + nn.Conv2d(256, 128, kernel_size=kernel_size, padding='same'), + nn.ReLU(), + nn.BatchNorm2d(128), + nn.Conv2d(128, 64, kernel_size=kernel_size, padding='same'), + nn.ReLU(), + nn.BatchNorm2d(64), + nn.Conv2d(64, 32, kernel_size=kernel_size, padding='same'), + nn.ReLU(), + nn.BatchNorm2d(32), + nn.Conv2d(32, 3, kernel_size=kernel_size, padding='same'), + nn.Tanh() + ) - - def call(self, inputs, training): + def forward(self, inputs) -> torch.Tensor: img, mask = inputs - x = self.seq(img, training=training) + x = self.seq(img) - x = x * (1 - mask[:,:,:,tf.newaxis]) + x = x * (1 - mask.unsqueeze(1)) x = x + img 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): -# 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/scripts/upscaler_train.py b/scripts/upscaler_train.py new file mode 100644 index 0000000..7abf6f5 --- /dev/null +++ b/scripts/upscaler_train.py @@ -0,0 +1,284 @@ +import os +import sys +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') + +import importlib +import argparse +parser = argparse.ArgumentParser(description='') +parser.add_argument('--name', type=str) +parser.add_argument('--config', type=str) +parser.add_argument('--debug', action='store_true') +args = vars(parser.parse_args()) + +os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" +os.environ["CUDA_VISIBLE_DEVICES"] = '7' + + +import time +import argparse +import wandb +import random +import numpy as np +import torch +# from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure as MSSSIM +from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, get_summary_str +import plotly.graph_objects as go +import plotly.express as ex +import torch.optim as optim +import torch.nn.functional as F + +from data.dataset import Dataset +# from torch.utils.data import DataLoader +from custom_utils.custom_utils import plot_3d +from models.upscaler import Upscaler +from utils.checkpoints import CheckpointSaver, load_model +from custom_utils.custom_utils import Accumulator +from utils.utils import InfiniteDataLoader as DataLoader + +from pytorch_msssim import MS_SSIM + + +low_res = 32 +is_skirt = False +is_cape = True +is_male = True +skirt_modifier = '_skirt' if is_skirt else '' +cape_modifier = '_cape' if is_cape else '' +gender_modifier = '' if not is_cape else ('_m' if is_male else '_f') +smpl_res = 3638 if is_cape else 14475 +use_cache = False +name=f'iccv_{low_res}' + skirt_modifier + cape_modifier + gender_modifier + +if not is_skirt: + train_dir = '/data-local/127/train' + valid_dir = '/data-local/127/valid' +else: + train_dir = '/data-local/128/train_skirt' + valid_dir = '/data-local/127/valid_skirt' + +if is_cape: + if is_male: + train_dir = '/data-local/127/cape_processed/train/male' + valid_dir = '/data-local/127/cape_processed/test/male' + else: + train_dir = '/data-local/127/cape_processed/train/female' + valid_dir = '/data-local/127/cape_processed/test/female' + + + +def init_wandb(): + run = wandb.init(project='upscale', name=name, notes='', dir='/data-local/127/temp', tags=['upscaler', 'iccv']) + artifact = wandb.Artifact('code', type='code') + artifact.add_file('scripts/upscaler_train.py') + artifact.add_file('losses/losses.py') + artifact.add_file('data/dataset.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 = UV2Mesh(res=256, is_skirt=is_skirt, is_smpl_orig=is_cape) + + +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 + + +def train_step(batch): + model.train() + with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp): + pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]]) + pred = F.interpolate(pred, size=(512, 256), mode='bilinear') + + l1_loss = torch.mean(torch.sum(torch.abs(batch['uv_unposed'] - pred * batch['uv_mask'].unsqueeze(1)), dim=-1)) + if low_res == 32: + ms_ssim = torch.mean(_MSSSIM_small(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1)) + else: + ms_ssim = torch.mean(_MSSSIM(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1)) + + loss = 1e3 * (1 - ms_ssim) + + scaler.scale(loss).backward() + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad() + + model.eval() + with torch.inference_mode(): + mesh_pred = uv2mesh(pred) + mesh_orig = batch['mesh_unposed'] + metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'].unsqueeze(2) * (mesh_orig - mesh_pred) ** 2, dim=-1))) / torch.sum(batch['mesh_mask']) + + return { 'train/loss': loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, } + + +def valid_step(batch): + model.eval() + + with torch.inference_mode(): + pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]]) + pred = F.interpolate(pred, size=(512, 256), mode='bilinear') + + l1_loss = torch.mean(torch.sum(torch.abs(batch['uv_unposed'] - pred * batch['uv_mask'].unsqueeze(1)), dim=-1)) + if low_res == 32: + ms_ssim = torch.mean(_MSSSIM_small(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1)) + else: + ms_ssim = torch.mean(_MSSSIM(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1)) + + loss = 1e3 * (1 - ms_ssim) + mesh_pred = uv2mesh(pred) + mesh_orig = batch['mesh_unposed'] + metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'].unsqueeze(2) * (mesh_orig - mesh_pred) ** 2, dim=-1))) / torch.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): + model.eval() + with torch.inference_mode(): + pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]]) + pred = F.interpolate(pred, size=(512, 256), mode='bilinear') + + mesh_pred = uv2mesh(pred) + mesh_orig = batch['mesh_unposed'] + + _F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if is_cape else "") + ".npy") + mask = batch['mesh_mask'].cpu().numpy().astype(bool) + garment_unposed, faces = maskFacesBatch(mesh_pred.cpu().numpy(), _F, mask) + # gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])] + gt = [mesh_orig[i].cpu().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 = batch['uv_unposed'] + img = torch.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(): + global step + valid_iterator = iter(valid_dataloader) + + for batch in train_dataloader: + train_loss = train_step(batch) + train_accumulator.update(train_loss) + + if step % 25 == 0: + data_valid = next(valid_iterator) + valid_loss = valid_step(data_valid) + valid_accumulator.update(valid_loss) + + if step % 1000 == 0 or (step < 1000 and step % 1000 == 0): + to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } }) + if step % 5000 == 0 or (step < 5000 and step % 1000 == 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)) } } + + if not args['debug']: + wandb.log(to_log, step=step) + train_accumulator.reset() + valid_accumulator.reset() + print(step) + + if step % 10000 == 0: + ckpt_saver(model, step, optimizer, scaler) + + step += 1 + + +if __name__ == '__main__': + if not args['debug']: + init_wandb() + + use_amp = True + uv_mask_body = torch.ones((low_res, low_res//2), device='cuda') + low_resolution_modifier = '_' + str(low_res) + + signature = { 'id', 'frame_nr', + 'pose_shape', + 'uv_unposed' + low_resolution_modifier, + 'uv_unposed', + 'mesh_unposed', + 'mesh_posed', + 'mesh_mask', + 'uv_mask', + 'uv_mask' + low_resolution_modifier, + } + + # dataset = Dataset('/data-local/127/train', config['signature'], resolutions=[512], device=device, frames_file='frames_jumpsuit.txt') + dataset_train = Dataset(train_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=None, resolutions=[512, low_res], is_cape=is_cape) + dataset_valid = Dataset(valid_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=None, resolutions=[512, low_res], is_cape=is_cape) + train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True) + valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True) + train_accumulator = Accumulator() + valid_accumulator = Accumulator() + + if low_res == 32 or low_res == 64: + kernel_size = 3 + elif low_res == 128 or low_res == 256: + kernel_size = 5 + else: + kernel_size = 7 + + model : Upscaler = Upscaler(kernel_size=kernel_size).to('cuda') + optimizer = optim.Adam(model.parameters(), lr=0.0001) + scaler = torch.cuda.amp.GradScaler(enabled=use_amp) + + # _MSSSIM = MSSSIM(data_range=2).to(config['device']) + _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3, weights=[0.5, 0.2, 0.2, 0.1]) + _MSSSIM_small = MS_SSIM(data_range=2, size_average=True, channel=3, weights=[1.0]) + + step = 0 + ckpt_saver = CheckpointSaver('checkpoints/iccv/torch/' + name, top_n=5) + + print(f'started training: {name}') + + if not args['debug']: + # model = torch.compile(model) + wandb.watch(model, log='all', log_freq=1000, log_graph=True) + + train() + + + diff --git a/utils/checkpoints.py b/utils/checkpoints.py index c098f44..07ebcae 100644 --- a/utils/checkpoints.py +++ b/utils/checkpoints.py @@ -16,9 +16,9 @@ class CheckpointSaver: self.top_n = top_n self.top_model_paths = [] - def __call__(self, model, iteration, optimizer): + def __call__(self, model, iteration, optimizer, scaler): model_path = os.path.join(self.dirpath, f'iteration_{iteration}.pt') - torch.save({ 'model': model.state_dict(), 'optimizer': optimizer }, model_path) + torch.save({ 'model': model.state_dict(), 'optimizer': optimizer, 'scaler': scaler }, model_path) self.top_model_paths.append({'path': model_path}) if len(self.top_model_paths)>self.top_n: self.cleanup() @@ -31,12 +31,13 @@ class CheckpointSaver: self.top_model_paths = self.top_model_paths[len(self.top_model_paths)-self.top_n:] -def load_model(model, checkpoint, optimizer=None, iteration=''): +def load_model(model, checkpoint, optimizer=None, scaler=None, iteration=''): # list of checkpoint files in the folder checkpoints checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith(f'{iteration}.pt')] checkpoint = torch.load(checkpoints[-1]) iteration = checkpoints[-1].split('_')[-1].split('.')[0] model.load_state_dict(checkpoint['model']) if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer']) + if scaler is not None: scaler.load_state_dict(checkpoint['scaler']) return int(iteration) if iteration != '' else 0