update

scripts/cvae_norm.py

@@ -12,7 +12,7 @@ args = vars(parser.parse_args())
 
 
 if args['debug']:
-    from configs.iccvw.cvae_norm_1 import config
+    from configs.iccvw.cvae_norm_5 import config
 else:
     config_module = importlib.import_module(args['config'])
     config = config_module.config
@@ -28,6 +28,7 @@ import wandb
 import tqdm
 import random
 import numpy as np
+import kaleido
 import plotly.graph_objects as go
 import plotly.express as ex
 from PIL import Image
@@ -93,7 +94,7 @@ def train_step(batch):
     kl = kl_loss(mean, logvar)
 
     reconstr_loss = l2_loss + l2_mesh_loss + (1 - ms_ssim)
-    loss = reconstr_loss + g_balancing_coeff * kl
+    loss: torch.Tensor = reconstr_loss + g_balancing_coeff * kl
 
     optimizer.zero_grad()
     loss.backward()
@@ -164,125 +165,6 @@ def plot_everything(batch: dict, type: str):
     return { **{ 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, config['signature_dynamic']))
-        uv_normals = tf.one_hot(batch['uv_normals'], 7)
-        nr_of_samples=2
-
-        pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][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(pred)
-        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((config['smpl_res'], 3)), F, mask)
-
-        pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
-        mesh_pred_zero = uv2mesh(pred_zero)
-        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()
-
-
 
 def train():
     global g_balancing_coeff
@@ -314,7 +196,7 @@ def train():
             print(step)
 
         if train_loss['train/euclid_error'] < config['balancing_activation_limit']:
-            g_balancing_coeff += 1 / 10000 
+            g_balancing_coeff += 1 / 2000 
 
         if step % 20000 == 0:
             ckpt_saver(model, step, optimizer, None)
@@ -322,6 +204,49 @@ def train():
         step += 1
 
 
+
+def sample():
+    model.eval()
+    load_model(model, f'checkpoints/iccvw/{config["name"]}')
+    count = 0
+
+    for batch in train_dataloader:
+        uv_static = batch['uv_static']
+        uv_body_posed = batch['uv_body_posed']
+        uv_mask = batch['uv_mask']
+
+        batch_size = uv_static.size(0)
+
+        for i in range(batch_size):
+            uv_static_tiled = uv_static[i][None,:,:,:].repeat(4, 1, 1, 1)
+            uv_body_posed_tiled = uv_body_posed[i][None,:,:,:].repeat(4, 1, 1, 1)
+            uv_mask_tiled = uv_mask[i][None,:,:].repeat(4, 1, 1)
+        
+            with torch.inference_mode():
+                pred = model.sample(uv_static_tiled, uv_body_posed_tiled, uv_mask_tiled)
+            
+            pred = pred * 1.4 
+
+            pred = pred.cpu().numpy().transpose(0, 2, 3, 1)
+            pred = (pred + 1.4) / 2.8
+            pred = pred * uv_mask_tiled[:,:,:,None].cpu().numpy()
+            gt = (batch['uv_normals'][i].cpu().numpy().transpose(1, 2, 0) + 1.4) / 2.8 * uv_mask_tiled[0,:,:,None].cpu().numpy()
+            img = np.concatenate(pred, axis=1)
+            img = np.concatenate([gt, img], axis=1)
+            # flip image along its shorter edge
+            img = img[::-1,:,:]
+
+            fig = ex.imshow(img)
+            fig.update_layout(margin=dict(l=0, r=0, b=0, t=0, pad=0), width=img.shape[1], height=img.shape[0])
+            fig.update_xaxes(visible=False)
+            fig.update_yaxes(visible=False)
+            fig.write_image(f'temp/normals/{batch["id"][i]}_{batch["frame_nr"][i].cpu().numpy()}_{batch["garment"][i]}.png')
+            count += 1
+
+        if count > 100:
+            break
+
+
 if __name__ == '__main__':
     if not args['debug']:
         init_wandb()
@@ -338,6 +263,10 @@ if __name__ == '__main__':
 
     model = CVAENorm(config)
 
+    sample()
+    exit()
+
+
     optimizer = optim.Adam(model.parameters(), lr=config['lr'])
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
     g_balancing_coeff = config['starting_balancing_coeff']

scripts/cvae_pyramid.py

@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.pyr_64_128_256_512_cape_f import config
+    from configs.iccvw.pyr_64_cape_f_skin import config
 else:
     config_module = importlib.import_module(args['config'])
     config = config_module.config

scripts/cvae_pyramid_e2e.py

@@ -0,0 +1,570 @@
+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())
+
+if args['debug']:
+    from configs.iccvw.pyr_64_128_256_512_cape_m_e2e import config
+else:
+    config_module = importlib.import_module(args['config'])
+    config = config_module.config
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+import time
+import argparse
+import wandb
+import tqdm
+import random
+import numpy as np
+import plotly.graph_objects as go
+import plotly.express as ex
+from PIL import Image
+
+if config['is_skirt']:
+    from losses.losses import edge_loss_skirt as edge_loss
+else:
+    from losses.losses import edge_loss
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+from torchvision.transforms import Resize
+from torchvision.transforms.functional import InterpolationMode
+from data.dataset import Dataset
+from utils.checkpoints import CheckpointSaver, load_model
+from utils.utils import InfiniteDataLoader as DataLoader
+from utils.utils import text_to_html
+from utils.postprocess import fix_collision
+from losses.losses import collision_loss, normal_loss, kl_loss
+from pytorch_msssim import MS_SSIM
+from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary_html, maskFacesBatch, get_summary_str, plot_3d_collision, group_inputs_for_levels
+from smpl.smpl_torch import SMPL
+from plotly.subplots import make_subplots
+from models.pyramid_model import Pyramid
+
+
+def init_wandb():
+    wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
+    run = wandb.init(project='iccvw', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('scripts/cvae_pyramid_e2e.py')
+    artifact.add_file('losses/losses.py')
+    artifact.add_file('models/base.py')
+    artifact.add_file('models/cond_encoder.py')
+    artifact.add_file('models/convnext.py')
+    artifact.add_file('models/pyramid_model.py')
+    artifact.add_file('models/pyramid_level.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=config['is_skirt'], is_smpl_orig=config['is_cape'])
+
+def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, collision_faces, collision_color):
+    fig = make_subplots(rows=1, cols=1+len(mesh_pred_unposed_levels)+1+1, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(1+len(mesh_pred_unposed_levels)+1+1)] ])
+    
+    fig.add_trace(go.Mesh3d(
+            x=mesh_gt[:,0],
+            y=mesh_gt[:,2],
+            z=mesh_gt[:,1],
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2],
+            text="Ground Truth"
+        ),
+        row=1, col=1)
+
+    for i, mesh in enumerate(mesh_pred_unposed_levels):
+        fig.add_trace(go.Mesh3d(
+            x=mesh[:,0],
+            y=mesh[:,2],
+            z=mesh[:,1],
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2],
+            text='Unposed level_' + str(i)
+        ),
+        row=1, col=i+2)
+
+    fig.add_trace(go.Mesh3d(
+            x=mesh_posed[:,0],
+            y=mesh_posed[:,2],
+            z=mesh_posed[:,1],
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2],
+            text='Posed GT'
+        ),
+        row=1, col=1+len(mesh_pred_unposed_levels)+1)
+
+    fig.add_trace(go.Mesh3d(
+            x=collision[:,0],
+            y=collision[:,2],
+            z=collision[:,1],
+            i = collision_faces[:,0],
+            j = collision_faces[:,1],
+            k = collision_faces[:,2],
+            vertexcolor=collision_color,
+            text='Posed output on body'
+        ),
+        row=1, col=1+len(mesh_pred_unposed_levels)+1+1)
+
+    if len(mesh_pred_unposed_levels) == 1:
+        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1), 
+                        scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1), 
+                        scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1), 
+                        scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1), 
+                        )
+        helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
+                    yaxis = dict(nticks=4, range=[-1,1],),
+                    zaxis = dict(nticks=4, range=[-1,1],),)
+        fig.update_layout(scene = helper,
+                            scene2 = helper,
+                            scene3 = helper,
+                            scene4 = helper,
+                            )
+    if len(mesh_pred_unposed_levels) == 2:
+        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1), 
+                        scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1), 
+                        scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1), 
+                        scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1), 
+                        scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1), 
+                        )
+        helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
+                    yaxis = dict(nticks=4, range=[-1,1],),
+                    zaxis = dict(nticks=4, range=[-1,1],),)
+        fig.update_layout(scene = helper,
+                            scene2 = helper,
+                            scene3 = helper,
+                            scene4 = helper,
+                            scene5 = helper,
+                            )
+    if len(mesh_pred_unposed_levels) == 3:
+        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
+                        scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
+                        scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
+                        scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
+                        scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1),
+                        scene6_aspectmode='manual', scene6_aspectratio=dict(x=1, y=1, z=1,) 
+                        )
+        helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
+                    yaxis = dict(nticks=4, range=[-1,1],),
+                    zaxis = dict(nticks=4, range=[-1,1],),)
+        fig.update_layout(scene = helper,
+                            scene2 = helper,
+                            scene3 = helper,
+                            scene4 = helper,
+                            scene5 = helper,
+                            scene6 = helper,
+                            )
+    if len(mesh_pred_unposed_levels) == 4:
+        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
+                        scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
+                        scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
+                        scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
+                        scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1),
+                        scene6_aspectmode='manual', scene6_aspectratio=dict(x=1, y=1, z=1,), 
+                        scene7_aspectmode='manual', scene7_aspectratio=dict(x=1, y=1, z=1,), 
+                        )
+        helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
+                    yaxis = dict(nticks=4, range=[-1,1],),
+                    zaxis = dict(nticks=4, range=[-1,1],),)
+        fig.update_layout(scene = helper,
+                            scene2 = helper,
+                            scene3 = helper,
+                            scene4 = helper,
+                            scene5 = helper,
+                            scene6 = helper,
+                            scene7 = helper,
+                            )
+    if len(mesh_pred_unposed_levels) == 5:
+        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
+                        scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
+                        scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
+                        scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
+                        scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1),
+                        scene6_aspectmode='manual', scene6_aspectratio=dict(x=1, y=1, z=1,), 
+                        scene7_aspectmode='manual', scene7_aspectratio=dict(x=1, y=1, z=1,), 
+                        scene8_aspectmode='manual', scene8_aspectratio=dict(x=1, y=1, z=1,), 
+                        )
+        helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
+                    yaxis = dict(nticks=4, range=[-1,1],),
+                    zaxis = dict(nticks=4, range=[-1,1],),)
+        fig.update_layout(scene = helper,
+                            scene2 = helper,
+                            scene3 = helper,
+                            scene4 = helper,
+                            scene5 = helper,
+                            scene6 = helper,
+                            scene7 = helper,
+                            scene8 = helper,
+                            )
+    return fig
+
+
+def train_step(batch):
+    if config['is_cape']:
+        batch['mesh_fix'] = {'mesh_mask': batch['mesh_mask']}
+
+    model.train()
+    # with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp):
+    levels_input, conds = group_inputs_for_levels(batch, config)
+
+    out_levels = model([levels_input, conds])
+
+    if config['use_skinning']:
+        mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
+    else:
+        mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
+
+    output = dict()
+    reconstr_loss_total = 0
+    kl_total = 0
+    posed_or_not = 'unposed' if config['use_skinning'] else 'posed'
+
+    for i, res in enumerate(config['pyramid']['resolutions']):
+        # i = -1
+        res = config['pyramid']['resolutions'][i]
+        res_modifier = '_' + str(res) if res != 512 else ''
+        pred, mean, logvar = out_levels[i]
+        if res != 512:
+            pred = resizer(pred) * batch['uv_mask'][:,None,:,:]
+
+        l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2])) # levels_input[i][0] = uv_orig
+        kl = kl_loss(mean, logvar)
+
+        mesh_pred = uv2mesh(pred)
+
+        if config['use_skinning']:
+            shape = batch['pose_shape'][:, 72:]
+            gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
+            shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+
+            mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_fix']['mesh_mask'][:,:,None]
+            mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+        else:
+            mesh_pred_posed = mesh_pred
+
+        e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_fix']['mesh_mask']))
+        c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_fix']['mesh_mask'])
+        c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
+        n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_fix']['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+        l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1))
+        
+        reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss)
+        if not config['is_cape']:
+            reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + c_loss) * config['reg_loss_coeff'] 
+            if config['use_norm_loss'] and step > 2000:
+                reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * n_loss * config['reg_loss_coeff']
+        
+        kl_total += kl
+    loss: torch.Tensor = reconstr_loss_total + g_balancing_coeff * kl_total
+
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+    model.eval()
+    with torch.inference_mode():
+        mean = torch.mean(mean)
+        logvar = torch.mean(logvar)
+
+        ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1))
+        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_fix']['mesh_mask'])
+
+    output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,'train/euclid_error' + res_modifier: metric, 
+                                'train/mean' + res_modifier: mean, 'train/logvar' + res_modifier: logvar, 'train/kl_loss' + res_modifier: kl,
+                                'train/var' + res_modifier: np.e ** logvar, 'train/l2_loss' + res_modifier: l2_loss, 'train/normal_loss': n_loss,
+                                'train/euclid_error': metric, 'train/l2_mesh_loss': l2_mesh_loss, 'train/ms_ssim' + res_modifier: ms_ssim,
+                                'train/collision_ratio': c_ratio } } # 
+
+
+    return { **{ 'train/loss': loss, 'train/reconstr_loss_total': reconstr_loss_total, 'train/kl_total': kl_total }, **output }
+
+
+def valid_step(batch):
+    if config['is_cape']:
+        batch['mesh_fix'] = {'mesh_mask': batch['mesh_mask']}
+
+    model.eval()
+    with torch.inference_mode():
+        levels_input, conds = group_inputs_for_levels(batch, config)
+
+        out_levels = model([levels_input, conds])
+
+        if config['use_skinning']:
+            mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
+        else:
+            mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
+
+        output = dict()
+        reconstr_loss_total = 0
+        kl_total = 0
+
+        for i, res in enumerate(config['pyramid']['resolutions']):
+            # i = -1
+            res = config['pyramid']['resolutions'][i]
+            res_modifier = '_' + str(res) if res != 512 else ''
+            pred, mean, logvar = out_levels[i]
+            if res != 512:
+                pred = resizer(pred) * batch['uv_mask'][:,None,:,:]
+
+            mesh_pred = uv2mesh(pred)
+
+            if config['use_skinning']:
+                shape = batch['pose_shape'][:, 72:]
+                gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
+                shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+
+                mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_fix']['mesh_mask'][:,:,None]
+                mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+            else:
+                mesh_pred_posed = mesh_pred
+
+            e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_fix']['mesh_mask']))
+            c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_fix']['mesh_mask'])
+            c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
+            l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1))
+            n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_fix']['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+            l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2]))
+            
+            kl = kl_loss(mean, logvar)
+
+            if config['use_extra_losses']:
+                reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss)
+                if not config['is_cape']:
+                    reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + c_loss) * config['reg_loss_coeff'] 
+                    if config['use_norm_loss'] and step > 2000:
+                        reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * n_loss * config['reg_loss_coeff']
+            else:
+                reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
+
+            kl_total += kl
+        loss = reconstr_loss_total + g_balancing_coeff * kl_total
+
+        mean = torch.mean(mean)
+        logvar = torch.mean(logvar)
+
+        ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1))
+        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_fix']['mesh_mask'])
+
+        output = { **output, **{ 'valid/edge_loss' + res_modifier: e_loss, 'valid/collision_loss' + res_modifier: c_loss,'valid/euclid_error' + res_modifier: metric, 
+                                    'valid/mean' + res_modifier: mean, 'valid/logvar' + res_modifier: logvar, 'valid/kl_loss' + res_modifier: kl,
+                                    'valid/var' + res_modifier: np.e ** logvar, 'valid/l2_loss' + res_modifier: l2_loss, 'valid/normal_loss': n_loss,
+                                    'valid/euclid_error': metric, 'valid/l2_mesh_loss': l2_mesh_loss, 'valid/ms_ssim' + res_modifier: ms_ssim, } } # 
+
+
+
+    return { **{ 'valid/loss': loss, 'valid/reconstr_loss_total': reconstr_loss_total, 'valid/kl_total': kl_total }, **output }
+
+
+def plot_everything(batch: dict, type: str):
+    if config['is_cape']:
+        batch['mesh_fix'] = {'mesh_mask': batch['mesh_mask']}
+
+    if config['use_skinning']:
+        mesh_orig, uv_orig = batch['mesh_unposed'].cpu().numpy(), batch['uv_unposed']
+    else:
+        mesh_orig, uv_orig = batch['mesh_posed'].cpu().numpy(), batch['uv_posed']
+    
+    model.eval()
+    with torch.inference_mode():
+        levels_input = group_inputs_for_levels(batch, config)
+        output_levels = model(levels_input)
+
+    mesh_pred_levels = [uv2mesh(output_levels[i][0]).cpu() for i in range(len(output_levels))]
+    mesh_orig_posed = batch['mesh_posed'].cpu()
+    body = batch['mesh_body_posed'].cpu().numpy()
+    
+    # mesh_pred = torch.cat((mesh_pred[:, :, 0:1], -mesh_pred[:, :, 2:3], mesh_pred[:, :, 1:2]), dim=2)
+    # mesh_pred_levels = [torch.cat((mesh_pred_levels[i][:, :, 0:1], -mesh_pred_levels[i][:, :, 2:3], mesh_pred_levels[i][:, :, 1:2]), dim=2).numpy() for i in range(len(mesh_pred_levels))]
+    # mesh_orig_posed = torch.cat((mesh_orig_posed[:, :, 0:1], -mesh_orig_posed[:, :, 2:3], mesh_orig_posed[:, :, 1:2]), dim=2).numpy()
+
+    orig_img = (torch.permute(uv_orig, [0,2,3,1]) + 1) / 2
+    uv_static_img = (torch.permute(batch['uv_static'], [0,2,3,1]) + 1) / 2
+    uv_body_posed_img = (torch.permute(batch['uv_body_posed'], [0,2,3,1]) + 1) / 2
+    uv_normals = (torch.permute(batch['uv_normals'], [0,2,3,1]) + 1.4) / 2.8
+    pred_img = torch.concatenate([torch.permute(model.resizers[512](output_levels[i][0]), [0,2,3,1]) for i in range(len(output_levels))], dim=2)
+    # pred_img = (pred_img - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis]) / (pred_img.max(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis] - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis])
+    pred_img = (pred_img - pred_img.flatten().reshape(pred_img.shape[0], -1).min(1)[0][:,None,None,None]) / (pred_img.flatten().reshape(pred_img.shape[0], -1).max(1)[0][:,None,None,None] - pred_img.flatten().reshape(pred_img.shape[0], -1).min(1)[0][:,None,None,None])
+
+    img = torch.concat((uv_static_img, uv_body_posed_img, uv_normals, orig_img, pred_img), dim=2)
+    imgs = [wandb.Image(img[i].cpu().numpy()) for i in range(4)]
+
+    _F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
+    mask = batch['mesh_fix']['mesh_mask'].cpu().numpy().astype(bool)
+    if 'mesh_fix' in batch and not config['is_cape']:
+        faces = [np.load( (dataset_train if type == 'train' else dataset_valid).folder + "/mesh_fixed/" + batch['id'][i] + '_' + batch['garment'][i] + ".npz" )['faces'] for i in range(len(batch['id']))]
+    else:
+        _, faces = maskFacesBatch(np.ones_like(mesh_orig), _F, mask)
+    
+    if config['use_skinning']:
+        shape = batch['pose_shape'][:, 72:]
+        gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
+        shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])]).cpu().numpy()
+
+        gt = [mesh_orig[i][mask[i], :] + shape_offsets[i, mask[i], :] for i in range(mask.shape[0])]
+
+        mesh_unposed_shaped_levels = [[] for _ in range(mask.shape[0])]
+        for i in range(len(output_levels)):
+            mesh_unposed_shaped = [mesh_pred_levels[i][j, mask[j], :] + shape_offsets[j, mask[j]] for j in range(mask.shape[0])]
+            for k in range(mask.shape[0]):
+                mesh_unposed_shaped_levels[k].append(mesh_unposed_shaped[k])
+
+        # full res version
+        mesh_unposed_shaped_full = np.zeros((mask.shape[0], config['smpl_res'], 3), dtype=np.float32)
+        for i in range(mask.shape[0]):
+            mesh_unposed_shaped_full[i, mask[i], :] = mesh_unposed_shaped[i]
+        mesh_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], torch.tensor(mesh_unposed_shaped_full, device=config['device']))[0].cpu().numpy()
+    
+        # rotate gt 180 degrees along z axis
+        # gt = [np.matmul(gt[i], np.array([[np.cos(np.pi), -np.sin(np.pi), 0], [np.sin(np.pi), np.cos(np.pi), 0], [0, 0, 1]])) for i in range(mask.shape[0])]
+
+        # gt_full = np.zeros((mask.shape[0], config['smpl_res'], 3), dtype=np.float32)
+        # for i in range(mask.shape[0]):
+        #     gt_full[i, mask[i], :] = gt[i]
+        # mesh_orig_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], torch.tensor(np.stack(gt_full, axis=0), device=config['device'], dtype=torch.float32))[0].cpu().numpy()
+    else:
+        mesh_posed = mesh_pred_levels[-1]
+        gt = [mesh_orig[i][mask[i], :] for i in range(mesh_posed.shape[0])]
+        mesh_unposed_shaped_levels = [[mesh_pred_levels[i][j, mask[j], :] for i in range(len(output_levels))] for j in range(mask.shape[0])]
+
+
+    gt_posed = [mesh_orig_posed[i, mask[i], :] for i in range(mask.shape[0])]
+    mesh_posed_fixed = fix_collision(torch.Tensor(mesh_posed).to('cuda'), torch.Tensor(body).to('cuda'), torch.Tensor(mask).to('cuda')).cpu().numpy()
+
+    mesh_posed = [mesh_posed[j, mask[j], :] for j in range(mask.shape[0])]
+    
+    if config['is_cape']:
+        collision = mesh_posed
+        collision_faces = faces
+        collision_color = [np.ones_like(mesh_posed[i], dtype=np.int32)*230 for i in range(len(mesh_posed))]
+    else:
+        collision = [np.concatenate((body[i], mesh_posed[i]), axis=0) for i in range(len(mesh_posed))]
+        collision_faces = [np.concatenate((_F, faces[i] + config['smpl_res']), axis=0) for i in range(len(mesh_posed))]
+        collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(mesh_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(mesh_posed))]
+
+    figs = { type + "/garment_" + str(i): plot_3d(gt[i], mesh_unposed_shaped_levels[i], faces[i], gt_posed[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
+    
+    if config['is_cape']:
+        figs_coll = dict()
+    else:
+        garment_posed_fixed = [mesh_posed_fixed[i, mask[i], :] for i in range(body.shape[0])]
+        collision_fixed = [np.concatenate((body[i], garment_posed_fixed[i]), axis=0) for i in range(len(mesh_posed))]
+        gt_posed = [batch['mesh_posed'].cpu().numpy()[i, mask[i], :] for i in range(mask.shape[0])]
+        collision_gt = [np.concatenate((body[i], gt_posed[i]), axis=0) for i in range(len(gt_posed))]
+        figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i]) for i in range(4) }
+
+    batch_id = batch['id']
+    batch_frame_nr = batch['frame_nr']
+    batch_garment = batch['garment']
+    batch_info = [f'batch_info: id: {batch_id[i]}, frame_nr: {batch_frame_nr[i]}, garment: {batch_garment[i]}' for i in range(len(batch_id))]
+    batch_info = text_to_html(batch_info)
+    batch_info = wandb.Html(batch_info)
+
+    return { **figs, **figs_coll, **{ type + '/images': imgs, }, **{ f'{type}/batch_info': batch_info }, }
+
+
+def speed_step(levels_input, conds):
+
+    return model.sample([levels_input, conds], training=False)
+
+
+
+def train():
+    global g_balancing_coeff
+    global step
+    valid_iterator = iter(valid_dataloader)
+
+    for batch in train_dataloader:
+        train_loss = train_step(batch)
+        train_accumulator.update(train_loss)
+
+        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 % 100 == 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 train_loss['train/euclid_error'] < config['balancing_activation_limit']:
+        #     g_balancing_coeff += 1 / 50000 
+
+        if step % 5000 == 0 and not args['debug']:
+            ckpt_saver(model, step, optimizer, None)
+
+        step += 1
+
+
+if __name__ == '__main__':
+    if not args['debug']:
+        init_wandb()
+
+    use_amp = True
+    smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+    dataset_resolutions = config['pyramid']['resolutions'].copy()
+    if 512 not in dataset_resolutions:
+        dataset_resolutions.append(512)
+
+    dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
+    dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config['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()
+
+    model: Pyramid = Pyramid(config)
+
+    # optimizer = optim.Adam(model.parameters(), lr=config['lr'])
+    optimizer = optim.Adam(model.parameters(), lr=config['lr'])
+    _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
+    ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=50)
+    scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
+    resizer = Resize(size=(512, 256), interpolation=InterpolationMode.BILINEAR).to(config['device'])
+
+    g_balancing_coeff = config['starting_balancing_coeff']
+    config['pyramid']['loss_weights'] = [1,2,3,4,5]
+
+    step = 0
+    if config['pretrained_model'] != '':
+        if 'ignore_saved_optimizer' in config.keys()  and config['ignore_saved_optimizer']:
+            step = load_model(model, config['pretrained_model'])
+        else:
+            step = load_model(model, config['pretrained_model'], optimizer=optimizer, iteration=config['pretrained_model_iteration'])
+
+    for param in model.static_enc.parameters(): param.requires_grad = False
+    for param in model.body_enc.parameters(): param.requires_grad = False
+    for param in model.norm_enc.parameters(): param.requires_grad = False
+    for i in range(len(model.levels)-1):
+        for param in model.levels[i].upscale_enhancer.parameters(): param.requires_grad = False
+
+    train()

scripts/cvae_pyramid_fine.py

@@ -40,6 +40,8 @@ import torch
 import torch.nn as nn
 import torch.optim as optim
 import torch.nn.functional as F
+from torchvision.transforms import Resize
+from torchvision.transforms.functional import InterpolationMode
 from data.dataset import Dataset
 from utils.checkpoints import CheckpointSaver, load_model
 from utils.utils import InfiniteDataLoader as DataLoader
@@ -57,7 +59,7 @@ def init_wandb():
     wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
     run = wandb.init(project='iccvw', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
-    artifact.add_file('scripts/cvae_pyramid.py')
+    artifact.add_file('scripts/cvae_pyramid_fine.py')
     artifact.add_file('losses/losses.py')
     artifact.add_file('models/base.py')
     artifact.add_file('models/cond_encoder.py')
@@ -241,8 +243,10 @@ def train_step(batch):
         res = config['pyramid']['resolutions'][i]
         res_modifier = '_' + str(res) if res != 512 else ''
         pred, mean, logvar = out_levels[i]
+        if res != 512:
+            pred = resizer(pred) * batch['uv_mask'][:,None,:,:]
 
-        l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((levels_input[i][0] - pred) ** 2), dim=-1), dim=[1, 2])) # levels_input[i][0] = uv_orig
+        l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2])) # levels_input[i][0] = uv_orig
         kl = kl_loss(mean, logvar)
 
         mesh_pred = uv2mesh(pred)
@@ -281,7 +285,7 @@ def train_step(batch):
         mean = torch.mean(mean)
         logvar = torch.mean(logvar)
 
-        ms_ssim = torch.mean(_MSSSIM(pred + 1, levels_input[i][0] + 1))
+        ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1))
         metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_fix']['mesh_mask'])
 
     output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,'train/euclid_error' + res_modifier: metric, 
@@ -318,6 +322,8 @@ def valid_step(batch):
             res = config['pyramid']['resolutions'][i]
             res_modifier = '_' + str(res) if res != 512 else ''
             pred, mean, logvar = out_levels[i]
+            if res != 512:
+                pred = resizer(pred) * batch['uv_mask'][:,None,:,:]
 
             mesh_pred = uv2mesh(pred)
 
@@ -336,7 +342,7 @@ def valid_step(batch):
             c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
             l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1))
             n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_fix']['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
-            l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((levels_input[i][0] - pred) ** 2), dim=-1), dim=[1, 2]))
+            l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2]))
             
             kl = kl_loss(mean, logvar)
 
@@ -355,7 +361,7 @@ def valid_step(batch):
         mean = torch.mean(mean)
         logvar = torch.mean(logvar)
 
-        ms_ssim = torch.mean(_MSSSIM(pred + 1, levels_input[i][0] + 1))
+        ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1))
         metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_fix']['mesh_mask'])
 
         output = { **output, **{ 'valid/edge_loss' + res_modifier: e_loss, 'valid/collision_loss' + res_modifier: c_loss,'valid/euclid_error' + res_modifier: metric, 
@@ -538,15 +544,18 @@ if __name__ == '__main__':
 
     model: Pyramid = Pyramid(config)
 
-    optimizer = optim.Adam(model.parameters(), lr=config['lr'] * 0.05)
+    # optimizer = optim.Adam(model.parameters(), lr=config['lr'])
+    optimizer = optim.Adam(model.parameters(), lr=config['lr'] * 0.1)
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
-    ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_fine_2', top_n=50)
+    ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=50)
     scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
+    resizer = Resize(size=(512, 256), interpolation=InterpolationMode.BILINEAR).to(config['device'])
 
     g_balancing_coeff = config['starting_balancing_coeff']
-    config['pyramid']['loss_weights'] = [1,3,5,10,25]
+    config['pyramid']['loss_weights'] = [1,2,3,4,5]
 
-    step = load_model(model, f'checkpoints/iccvw/{config["name"]}', iteration='25000')
+    step = 0
+    load_model(model, config['pretrained_model'], iteration=config['pretrained_model_iteration'])
 
     for param in model.static_enc.parameters(): param.requires_grad = False
     for param in model.body_enc.parameters(): param.requires_grad = False

scripts/generate_results.py

@@ -13,10 +13,10 @@ parser.add_argument('-g', '--gpu', type=str)
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.pyr_64_128_cape_f_2 import config
+    from configs.iccvw.pyr_64_128_cape_f_2 import config_1
 else:
     config_module = importlib.import_module(args['config'])
-    config = config_module.config
+    config_1 = config_module.config
 
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = args['gpu']
@@ -53,26 +53,26 @@ from models.pyramid_model import Pyramid
 from models.cvae_model import CVAE
 
 
-uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+uv2mesh = UV2Mesh(res=256, is_skirt=config_1['is_skirt'], is_smpl_orig=config_1['is_cape'])
 
 
 def step_pyramid(batch):
     model.eval()
     with torch.inference_mode():
-        levels_input, conds = group_inputs_for_levels(batch, config)
+        levels_input, conds = group_inputs_for_levels(batch, config_1)
 
         out_levels = model([levels_input, conds])
         pred, _, _ = out_levels[-1]
 
         mesh_pred = uv2mesh(pred)
 
-        if config['use_skinning']:
+        if config_1['use_skinning']:
             shape = batch['pose_shape'][:, 72:]
             gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
             shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
 
             mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
-            if not config['is_cape']:
+            if not config_1['is_cape']:
                 mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
             else:
                 mesh_pred_posed = mesh_pred_shaped
@@ -83,7 +83,7 @@ def step_pyramid(batch):
 
 
 def step(batch):
-    if config['use_skinning']:
+    if config_1['use_skinning']:
         uv_orig = batch['uv_unposed']
     else:
         uv_orig = batch['uv_posed']
@@ -92,7 +92,7 @@ def step(batch):
     pred, _, _ = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']])
     mesh_pred = uv2mesh(pred)
 
-    if config['use_skinning']:
+    if config_1['use_skinning']:
         shape = batch['pose_shape'][:, 72:]
         gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
         shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
@@ -141,33 +141,33 @@ def generate():
 
 if __name__ == '__main__':
 
-    smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
+    smpl = SMPL(orig_smpl=config_1['is_cape'], path='smpl', device=config_1['device'])
-    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+    skirt_modifier = '_skirt' if config_1['is_skirt'] else ''
     dataset_resolutions = []
-    if 'pyramid' in config:
+    if 'pyramid' in config_1:
-        dataset_resolutions = config['pyramid']['resolutions'].copy()
+        dataset_resolutions = config_1['pyramid']['resolutions'].copy()
     if 512 not in dataset_resolutions:
         dataset_resolutions.append(512)
 
     # remove all values from config['signature] that start with mesh_
-    config['signature'] = set([x for x in config['signature'] if not x.startswith('mesh_')])
+    config_1['signature'] = set([x for x in config_1['signature'] if not x.startswith('mesh_')])
 
-    dataset_test = Dataset(config['test_dir'], config['signature'], device=config['device'], frames_file=config['test_frames'], use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
+    dataset_test = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file=config_1['test_frames'], use_cache=config_1['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
     # dataset_test = Dataset(config['test_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
     test_dataloader = DataLoader(dataset_test, batch_size=4, drop_last=True, shuffle=True, )
 
-    is_pyramid = config['name'].startswith('pyr')
+    is_pyramid = config_1['name'].startswith('pyr')
 
     if is_pyramid:
-        model: Pyramid = Pyramid(config)
+        model: Pyramid = Pyramid(config_1)
     else:
-        model: CVAE = CVAE(config)
+        model: CVAE = CVAE(config_1)
 
-    it = load_model(model, 'checkpoints/iccvw/' + config['name'], iteration=args['iteration'])
+    it = load_model(model, 'checkpoints/iccvw/' + config_1['name'], iteration=args['iteration'])
     model.eval()
     # it = load_model(model, 'checkpoints/iccvw/' + config['name'])
 
-    name = config['name'] + '_' + str(it)
+    name = config_1['name'] + '_' + str(it)
 
     folder = '/data-local/127_2/results'
     out_folder = f'{folder}/{name}'

scripts/generate_results_all.py

@@ -0,0 +1,142 @@
+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('--iteration', type=str, default='')
+parser.add_argument('--debug', action='store_true')
+parser.add_argument('-g', '--gpu', type=str, default='1')
+args = vars(parser.parse_args())
+
+if args['debug']:
+    from configs.iccvw.pyr_64_128_256_512_e2e import config as config_1
+else:
+    config_module = importlib.import_module(args['config'])
+    config_1 = config_module.config
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = args['gpu']
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+import time
+import argparse
+import wandb
+import tqdm
+import random
+import numpy as np
+import plotly.graph_objects as go
+import plotly.express as ex
+from PIL import Image
+
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+from tqdm import tqdm
+from data.dataset import Dataset
+from utils.checkpoints import CheckpointSaver, load_model
+from torch.utils.data import DataLoader
+from utils.utils import text_to_html
+from utils.postprocess import fix_collision
+from losses.losses import collision_loss, normal_loss, kl_loss
+from pytorch_msssim import MS_SSIM
+from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary_html, maskFacesBatch, get_summary_str, plot_3d_collision, group_inputs_for_levels
+from smpl.smpl_torch import SMPL
+from plotly.subplots import make_subplots
+from models.pyramid_model import Pyramid
+from models.cvae_model import CVAE
+
+
+uv2mesh = UV2Mesh(res=256, is_skirt=config_1['is_skirt'], is_smpl_orig=config_1['is_cape'])
+
+
+def step_pyramid(batch):
+    model.eval()
+    with torch.inference_mode():
+        levels_input, conds = group_inputs_for_levels(batch, config_1)
+
+        out = []
+        out_levels = model([levels_input, conds])
+        for i in range(len(out_levels)):
+            pred, _, _ = out_levels[i]
+
+            mesh_pred = uv2mesh(pred)
+
+            if config_1['use_skinning']:
+                shape = batch['pose_shape'][:, 72:]
+                gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
+                shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+
+                mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
+                if not config_1['is_cape']:
+                    mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+                else:
+                    mesh_pred_posed = mesh_pred_shaped
+            else:
+                mesh_pred_posed = mesh_pred
+
+            out.append(mesh_pred_posed.cpu().numpy())
+
+    return out
+
+
+
+def generate():
+    for batch in tqdm(test_dataloader):
+        with torch.inference_mode():
+            mesh_pred = step_pyramid(batch)
+
+        ids = batch['id']
+        frame_nrs = batch['frame_nr']
+        garments = batch['garment']
+
+
+        for i in range(len(ids)):
+            np.savez_compressed(f'{out_folder}/{ids[i]}_{frame_nrs[i]}_{garments[i]}', { f'lvl{k}': mesh_pred[k][i] for k in range(len(config_1['pyramid']['resolutions'])) } )
+
+#np.load(f'{out_folder}/00462_67_Jumpsuit.npz', allow_pickle=True)['arr_0'].item()['lvl3']
+
+
+
+if __name__ == '__main__':
+
+    smpl = SMPL(orig_smpl=config_1['is_cape'], path='smpl', device=config_1['device'])
+    skirt_modifier = '_skirt' if config_1['is_skirt'] else ''
+    dataset_resolutions = []
+    if 'pyramid' in config_1:
+        dataset_resolutions = config_1['pyramid']['resolutions'].copy()
+    if 512 not in dataset_resolutions:
+        dataset_resolutions.append(512)
+
+    # remove all values from config['signature] that start with mesh_
+    config_1['signature'] = set([x for x in config_1['signature'] if not x.startswith('mesh_')])
+
+    dataset_test = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file=config_1['test_frames'], use_cache=config_1['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
+    # dataset_test = Dataset(config['test_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
+    test_dataloader = DataLoader(dataset_test, batch_size=4, drop_last=True, shuffle=True, )
+
+    is_pyramid = config_1['name'].startswith('pyr')
+
+    if is_pyramid:
+        model: Pyramid = Pyramid(config_1)
+    else:
+        model: CVAE = CVAE(config_1)
+
+    it = load_model(model, 'checkpoints/iccvw/' + config_1['name'], iteration=args['iteration'])
+    model.eval()
+    # it = load_model(model, 'checkpoints/iccvw/' + config['name'])
+
+    name = config_1['name'] + '_' + str(it)
+
+    folder = '/data-local/127_2/results'
+    out_folder = f'{folder}/{name}'
+
+    os.makedirs(out_folder, exist_ok=True)
+
+    generate()

scripts/get_error.py

@@ -13,10 +13,10 @@ parser.add_argument('-g', '--gpu', type=str)
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.pyr_64_128_cape_f_2 import config
+    from configs.iccvw.pyr_64_128_cape_f_2 import config_1
 else:
     config_module = importlib.import_module(args['config'])
-    config = config_module.config
+    config_1 = config_module.config
 
 
 import time
@@ -57,21 +57,21 @@ def calculate():
 
 
 if __name__ == '__main__':
-    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+    skirt_modifier = '_skirt' if config_1['is_skirt'] else ''
 
-    name = config['name']
+    name = config_1['name']
-    if not config['is_cape']:
+    if not config_1['is_cape']:
-        frames_file = config['test_frames']
+        frames_file = config_1['test_frames']
-        with open(os.path.join(config['test_dir'], frames_file), 'r') as f:
+        with open(os.path.join(config_1['test_dir'], frames_file), 'r') as f:
             frames = f.readlines()
         frames = [Frame(frames.split()[0], int(frames.split()[1]), frames.split()[2]) for frames in frames]
     else:
         from data.dataset_cape import read_frames
-        frames = read_frames(config['test_dir'], shuffle=False, is_cape=True)
+        frames = read_frames(config_1['test_dir'], shuffle=False, is_cape=True)
 
     results_folder = '/data-local/127_2/results'
     if args['iteration'] == '':
-        folders = [f for f in os.listdir(results_folder) if os.path.isdir(os.path.join(results_folder, f)) and f.startswith(config['name'])]
+        folders = [f for f in os.listdir(results_folder) if os.path.isdir(os.path.join(results_folder, f)) and f.startswith(config_1['name'])]
         assert len(folders) == 1
         name = folders[0]
     else:

scripts/get_error_merged.py

@@ -0,0 +1,109 @@
+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('--config1', type=str)
+parser.add_argument('--config2', type=str, default='')
+parser.add_argument('--debug', action='store_true')
+parser.add_argument('--iteration1', type=str, default='')
+parser.add_argument('--iteration2', type=str, default='')
+parser.add_argument('-g', '--gpu', type=str)
+args = vars(parser.parse_args())
+
+
+if args['debug']:
+    from configs.iccvw.pyr_64_128_cape_f_2 import config_1
+else:
+    config_module = importlib.import_module(args['config1'])
+    config_1 = config_module.config
+    config_module = importlib.import_module(args['config2'] if args['config2'] != '' else args['config1'] + '_skirt')
+    config_2 = config_module.config
+
+import time
+import argparse
+import tqdm
+import random
+import numpy as np
+import plotly.graph_objects as go
+import plotly.express as ex
+from data.utils import Frame
+
+
+
+
+def calculate(name):
+    errors = dict()
+    errors_mean = dict()
+
+    with open(f'{results_folder}/errors/{name}.csv', 'r') as f:
+        lines = f.readlines()
+
+    for line in lines:
+        line = line.split(',')
+        category = line[2]
+        if category not in errors:
+            errors[category] = []
+        errors[category].append(float(line[3]))
+
+    for category in errors:
+        errors_mean[category] = np.mean(errors[category])
+
+    return errors_mean
+
+
+
+if __name__ == '__main__':
+    name = config_1['name']
+    if not config_1['is_cape']:
+        frames_file = config_1['test_frames']
+        with open(os.path.join(config_1['test_dir'], frames_file), 'r') as f:
+            frames = f.readlines()
+        frames = [Frame(frames.split()[0], int(frames.split()[1]), frames.split()[2]) for frames in frames]
+    else:
+        from data.dataset_cape import read_frames
+        frames = read_frames(config_1['test_dir'], shuffle=False, is_cape=True)
+
+    results_folder = '/data-local/127_2/results'
+    if args['iteration1'] == '':
+        folders = [f for f in os.listdir(results_folder) if os.path.isdir(os.path.join(results_folder, f)) and f.startswith(config_1['name'])]
+        if not config_1['is_skirt']:
+            folders = [f for f in folders if f.find('skirt') == -1]
+        assert len(folders) == 1
+        name = folders[0]
+    else:
+        name = f'{name}_{args["iteration1"]}'
+
+    non_skirt = calculate(name)
+
+    name = config_2['name']
+    if not config_2['is_cape']:
+        frames_file = config_2['test_frames']
+        with open(os.path.join(config_2['test_dir'], frames_file), 'r') as f:
+            frames = f.readlines()
+        frames = [Frame(frames.split()[0], int(frames.split()[1]), frames.split()[2]) for frames in frames]
+    else:
+        from data.dataset_cape import read_frames
+        frames = read_frames(config_2['test_dir'], shuffle=False, is_cape=True)
+
+    results_folder = '/data-local/127_2/results'
+    if args['iteration2'] == '':
+        folders = [f for f in os.listdir(results_folder) if os.path.isdir(os.path.join(results_folder, f)) and f.startswith(config_2['name'])]
+        if not config_2['is_skirt']:
+            folders = [f for f in folders if f.find('skirt') == -1]
+        assert len(folders) == 1
+        name = folders[0]
+    else:
+        name = f'{name}_{args["iteration2"]}'
+    
+    skirt = calculate(name)
+
+    all = {**non_skirt, **skirt}
+
+    print(all)
+    mean = np.mean(list(all.values()))
+
+    print(f'Final error: {mean}')

scripts/get_final_cape.py

@@ -0,0 +1,188 @@
+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('--iteration', type=str, default='')
+parser.add_argument('--debug', action='store_true')
+parser.add_argument('-g', '--gpu', type=str, default='1')
+args = vars(parser.parse_args())
+
+if args['debug']:
+    from configs.iccvw.pyr_64_128_256_512_cape_f_e2e import config as config_1
+else:
+    config_module = importlib.import_module(args['config'])
+    config_1 = config_module.config
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = args['gpu']
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+import time
+import argparse
+import wandb
+import tqdm
+import random
+import numpy as np
+import plotly.graph_objects as go
+import plotly.express as ex
+from PIL import Image
+
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+from tqdm import tqdm
+from data.dataset import Dataset
+from utils.checkpoints import CheckpointSaver, load_model
+from torch.utils.data import DataLoader
+from utils.utils import text_to_html
+from utils.postprocess import fix_collision
+from losses.losses import collision_loss, normal_loss, kl_loss
+from pytorch_msssim import MS_SSIM
+from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary_html, maskFacesBatch, get_summary_str, plot_3d_collision, group_inputs_for_levels
+from smpl.smpl_torch import SMPL
+from plotly.subplots import make_subplots
+from models.pyramid_model import Pyramid
+from models.cvae_model import CVAE
+from data.utils import Frame
+
+
+uv2mesh = UV2Mesh(res=256, is_skirt=config_1['is_skirt'], is_smpl_orig=config_1['is_cape'])
+
+
+def step_pyramid(batch):
+    model.eval()
+    with torch.inference_mode():
+        levels_input, conds = group_inputs_for_levels(batch, config_1)
+
+        out = []
+
+        out_levels = model([levels_input, conds])
+        for i in range(len(out_levels)):
+            pred, _, _ = out_levels[i]
+
+            mesh_pred = uv2mesh(pred)
+
+            if config_1['use_skinning']:
+                shape = batch['pose_shape'][:, 72:]
+                gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
+                shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+
+                mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
+                if not config_1['is_cape']:
+                    mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+                else:
+                    mesh_pred_posed = mesh_pred_shaped
+            else:
+                mesh_pred_posed = mesh_pred
+
+            out.append(mesh_pred_posed)
+
+    return out
+
+from smpl.smpl_torch import SMPL
+smpl = SMPL(path='smpl', orig_smpl=True, device='cpu')
+clothing_verts_idx = np.load(os.path.join('custom_data', 'clothing_verts_idx.npy'))
+cape_mask = np.zeros((6890), bool)
+cape_mask[clothing_verts_idx] = True
+
+def get_error(pred, frame: Frame):
+    gt = np.load(f'{config_1["test_dir"]}/mesh_unposed_no_shape/{frame.id}.{frame.frame_nr}.npz')['mesh_unposed_no_shape']
+    pose_shape = np.load(f'{config_1["test_dir"]}/pose_shape/{frame.id}.{frame.frame_nr}.npy').astype(np.float32)
+    pose, shape = pose_shape[:72], pose_shape[72:]
+    gender = int(shape[10] >= 0.5)
+    shape_offsets = torch.tensordot(torch.from_numpy(shape[:10][None, :]), smpl.shapedirs[gender].to('cpu'), dims=[[1], [2]])[0,smpl.to_keep,:].numpy()
+
+    gt_shaped = gt + shape_offsets 
+    
+    min_body = np.load(os.path.join('/data-local/127_2/cape_processed', 'minimal_body_shape', frame.id.split('_')[0], frame.id.split('_')[0] + '_minimal.npy'))
+    offsets_pred = pred[cape_mask[smpl.to_keep],:] - min_body[clothing_verts_idx,:]
+    offsets_orig =  gt_shaped[cape_mask[smpl.to_keep],:] - min_body[clothing_verts_idx,:]
+
+    error = np.linalg.norm(offsets_pred - offsets_orig, axis=1).mean()
+    return error
+
+
+def generate():
+    errors = {0: [], 1: [], 2: [], 3: []}
+    for batch in tqdm(test_dataloader):
+        with torch.inference_mode():
+            mesh_pred = step_pyramid(batch)
+
+
+        ids = batch['id']
+        frame_nrs = batch['frame_nr']
+        garments = batch['garment']
+
+        for i in range(len(mesh_pred)):
+            frame = Frame(ids[i], frame_nrs[i], garments[i])
+            for j in range(4):
+                error = get_error(mesh_pred[j][i].cpu().numpy(), frame)
+                errors[j].append(error)
+
+        # print(errors)
+    for j in range(4):
+        print(f'Level {j}: {np.mean(errors[j])}')
+
+
+        # import plotly.graph_objects as go
+        # mesh_pred = mesh_pred.cpu().numpy()
+        # mask = batch['mesh_mask'].cpu().numpy()
+        # i = 3
+        # fig = go.Figure(data=[go.Scatter3d(x=mesh_pred[i,mask[i],0], y=mesh_pred[i,mask[i],2], z=mesh_pred[i,mask[i],1], mode='markers', marker=dict(size=1))])
+        # fig.update_layout(scene_aspectmode='cube')
+        # fig.show()
+
+        # import plotly.graph_objects as go
+        # mesh_pred = mesh_pred.cpu().numpy()
+        # i = 3
+        # fig = go.Figure(data=[go.Scatter3d(x=mesh_pred[i,:,0], y=mesh_pred[i,:,2], z=mesh_pred[i,:,1], mode='markers', marker=dict(size=1))])
+        # fig.update_layout(scene_aspectmode='cube')
+        # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
+        # fig.show()
+
+
+
+
+
+
+if __name__ == '__main__':
+
+    smpl = SMPL(orig_smpl=config_1['is_cape'], path='smpl', device=config_1['device'])
+    skirt_modifier = '_skirt' if config_1['is_skirt'] else ''
+    dataset_resolutions = []
+    if 'pyramid' in config_1:
+        dataset_resolutions = config_1['pyramid']['resolutions'].copy()
+    if 512 not in dataset_resolutions:
+        dataset_resolutions.append(512)
+
+    # remove all values from config['signature] that start with mesh_
+    config_1['signature'] = set([x for x in config_1['signature'] if not x.startswith('mesh_')])
+
+    dataset_test = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file=config_1['test_frames'], use_cache=config_1['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
+    # dataset_test = Dataset(config['test_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
+    test_dataloader = DataLoader(dataset_test, batch_size=4, drop_last=True, shuffle=True, )
+
+    is_pyramid = config_1['name'].startswith('pyr')
+
+    model: Pyramid = Pyramid(config_1)
+
+    it = load_model(model, 'checkpoints/iccvw/' + config_1['name'], iteration=args['iteration'])
+    model.eval()
+    # it = load_model(model, 'checkpoints/iccvw/' + config['name'])
+
+    name = config_1['name'] + '_' + str(it)
+
+    folder = '/data-local/127_2/results'
+    out_folder = f'{folder}/{name}'
+
+    os.makedirs(out_folder, exist_ok=True)
+
+    generate()