diff --git a/calculate_error.py b/calculate_error.py
new file mode 100644
index 0000000..4358089
--- /dev/null
+++ b/calculate_error.py
@@ -0,0 +1,151 @@
+import os
+import sys
+from configs.config_abl_abs_skirt 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 csv
+import glob
+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
+
+# if config['is_skirt']:
+#     from losses import edge_loss_skirt as edge_loss
+# else:
+#     from losses import edge_loss
+
+
+
+# from postprocessing import fix_collision
+# from losses import collision_ratio, kl_loss #, concept_loss
+# # from losses import collision_penalty as collision_penalty
+# from losses import collision_loss as collision_penalty
+# from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
+# from custom_utils.smpl import SMPL
+# from custom_utils.supersmpl_io import writeOBJ
+
+
+
+
+# uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+
+    # smpl = SMPL(use_one_hot=config['use_one_hot'])
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+
+    folder = config['test_dir']
+
+    checkpoint_name = config['name']
+    iteration = '70000'
+
+    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
+    garment_type_table = np.load('custom_data/garment_type_table_test.npy', allow_pickle=True).item()
+
+
+    output_folder = '/home/workspace/data/results/' + checkpoint_name + '-' + iteration
+
+
+    faces_folder = '/home/workspace/data/results/support/faces' + ('_skirt' if config['is_skirt'] else '')
+    ids = os.listdir(faces_folder)
+    ids = [id.split('.')[0] for id in ids]
+    faces = {id : np.load(faces_folder + '/' + id + '.npy') for id in ids }
+
+    masks_folder = '/home/workspace/data/dataset/test' + ('_skirt' if config['is_skirt'] else '') + '_ready/mesh_mask'
+    ids = os.listdir(masks_folder)
+    ids = [id.split('.')[0] for id in ids]
+    masks = {id : np.load(masks_folder + '/' + id + '.npy') for id in ids }
+
+    frames = os.listdir(output_folder)
+    frames = [frame for frame in frames if not frame.endswith('_fixed_coll.npy')]
+
+    if config['remove_global_orientation']:
+        gt_name = 'mesh_posed_no_orient'
+    else:
+        gt_name = 'mesh_posed'
+
+    orig_folder = '/home/workspace/data/dataset/test' + ('_skirt' if config['is_skirt'] else '') + '_ready' + ('_2' if config['is_skirt'] else '') + '/' + gt_name
+
+    errors = dict()
+    if config['is_skirt']:
+        per_garment_errors = { 'Tshirt': [], 'Dress': [], 'Skirt': [], 'Top': [] }
+    else:
+        per_garment_errors = { 'Tshirt': [], 'Jumpsuit': [], 'Trousers': [], 'Top': [] }
+
+    for frame in tqdm.tqdm(frames):
+        id = frame.split('_')[0]
+        frame_nr = frame.split('.')[0].split('_')[1]
+        gt = np.load(os.path.join(orig_folder, id + '_' + frame_nr + '.npz'))[gt_name][masks[id], :].astype(np.float32)
+        pred = np.load(os.path.join(output_folder, id + '_' + frame_nr + '.npy'))
+        metric = np.mean(np.sqrt(np.sum((gt - pred) ** 2, axis=-1)))
+        errors[id + '_' + frame_nr] = metric
+        per_garment_errors[garment_type_table[id]].append(metric)
+
+
+    for key in per_garment_errors.keys():
+        print(key + ': ' + str(np.mean(per_garment_errors[key])))
+
+    with open('/home/workspace/data/results/errors/' + checkpoint_name + '-' + iteration + '.csv', mode='w') as f:
+        file_writer = csv.writer(f, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
+
+        for key in sorted(errors.keys()):
+            file_writer.writerow([key, errors[key]])
+
+        # x = gt
+        # y = pred
+        # fig = go.Figure(data=[
+        #         go.Mesh3d(
+        #             x=x[:,0],
+        #             y=x[:,1],
+        #             z=x[:,2],
+        #             i = faces[id][:,0],
+        #             j = faces[id][:,1],
+        #             k = faces[id][:,2],
+        #         ),
+        #         go.Mesh3d(
+        #             x=y[:,0],
+        #             y=y[:,1],
+        #             z=y[:,2],
+        #             i = faces[id][:,0],
+        #             j = faces[id][:,1],
+        #             k = faces[id][:,2],
+        #         ),
+        #     ])
+
+        # fig.update_layout(scene_aspectmode='manual',
+        #                 scene_aspectratio=dict(x=1, y=1, z=1))
+        # fig.update_layout(
+        #     scene = dict(
+        #         xaxis = dict(nticks=4, range=[-1,1],),
+        #                     yaxis = dict(nticks=4, range=[-1,1],),
+        #                     zaxis = dict(nticks=4, range=[-1,1],),)
+        #                     )
+
+        # fig.show()
+
+    #     print()
+
+    # print()
+
diff --git a/convert_to_uv_local_single.py b/convert_to_uv_local_single.py
index 06b9867..8858584 100644
--- a/convert_to_uv_local_single.py
+++ b/convert_to_uv_local_single.py
@@ -195,7 +195,8 @@ def convert_to_uv_batch(res, m2uv, uv_map, verts, neighbours=15, mask=None, meth
         if garment_uv_mask is not None:
             result *= garment_uv_mask[:,:,np.newaxis]
             
-        np.savez_compressed(os.path.join(local_folder, 'uv_unposed_no_shape_one_hot_2', id + '_' + str(i) + '.npz'), uv_unposed_no_shape=result)
+        np.savez_compressed(os.path.join(local_folder, 'uv_posed_no_orient', id + '_' + str(i) + '.npz'), uv_posed_no_orient=result)
+        # np.savez_compressed(os.path.join(local_folder, 'uv_unposed_no_shape_one_hot_2', id + '_' + str(i) + '.npz'), uv_unposed_no_shape=result)
         
 
 
@@ -383,17 +384,18 @@ def test(verts, faces, values):
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/local-data/train')
+    parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/local-data/valid')
     parser.add_argument("-s", "--start", help = "start of range", default=0)
-    parser.add_argument("-e", "--end", help = "end of range", default=10)
+    parser.add_argument("-e", "--end", help = "end of range", default=1500)
     parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
-    parser.add_argument("-sk", "--skirt", help = "end of range", default=0)
+    parser.add_argument("-sk", "--skirt", help = "end of range", default=1)
     args = parser.parse_args()
 
     # smpl = SMPL('FeatureLines/superSMPL')
 
     local_folder = args.folder 
-    os.makedirs(os.path.join(local_folder, 'uv_unposed_no_shape_one_hot_2'), exist_ok=True)
+    os.makedirs(os.path.join(local_folder, 'uv_posed_no_orient'), exist_ok=True)
+    # os.makedirs(os.path.join(local_folder, 'uv_unposed_no_shape_one_hot_2'), exist_ok=True)
 
     print("skirt" if args.skirt == 1 else "not skirt")
 
@@ -432,8 +434,11 @@ if __name__ == '__main__':
     # rbf_dest[:, [0,1]] = rbf_dest[:, [1, 0]]
     # rbf_dest = rbf_dest[unique_pairs(rbf_dest), :] # not good, unique_pairs doesn't consider coord order
 
-    ids = os.listdir(os.path.join(local_folder, 'mesh_static'))[int(args.start):int(args.end)]
-    ids = [id.split('.')[0] for id in ids]
+    # ids = os.listdir(os.path.join(local_folder, 'mesh_static'))[int(args.start):int(args.end)]
+    # ids = [id.split('.')[0] for id in ids]
+    ids = os.listdir(os.path.join(local_folder, 'pose_shape'))
+    ids = [id.split('_')[0] for id in ids]
+    ids = sorted(list(set(ids)))[int(args.start):int(args.end)]
 
     for id in tqdm.tqdm(ids):
 
@@ -463,8 +468,10 @@ if __name__ == '__main__':
         # mesh_unposed_no_shape = np.stack([mesh_unposed_no_shape[i]['mesh_unposed_no_shape'] for i in range(gt.shape[0])])
         # body_no_orient = [np.load(os.path.join(folder, id, 'mesh_posed_body_no_orient',  id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])]
         # body_no_orient = np.stack([body_no_orient[i]['mesh_posed_body_no_orient'] for i in range(gt.shape[0])])
-        mesh_unposed_no_shape_one_hot = [np.load(os.path.join(local_folder, 'mesh_unposed_no_shape_one_hot',  id + '_' + str(i) + '.npz')) for i in range(frames)]
-        mesh_unposed_no_shape_one_hot = [mesh_unposed_no_shape_one_hot[i]['mesh_unposed_no_shape'] for i in range(frames)]
+        # mesh_unposed_no_shape_one_hot = [np.load(os.path.join(local_folder, 'mesh_unposed_no_shape_one_hot',  id + '_' + str(i) + '.npz')) for i in range(frames)]
+        # mesh_unposed_no_shape_one_hot = [mesh_unposed_no_shape_one_hot[i]['mesh_unposed_no_shape'] for i in range(frames)]
+        mesh_posed_no_orient = [np.load(os.path.join(local_folder, 'mesh_posed_no_orient',  id + '_' + str(i) + '.npz')) for i in range(frames)]
+        mesh_posed_no_orient = [mesh_posed_no_orient[i]['mesh_posed_no_orient'] for i in range(frames)]
 
 
         # verts_dynamic = gt
@@ -487,7 +494,8 @@ if __name__ == '__main__':
         # uv_curvature = convert_to_uv_batch_label(uv_res, m2uv, uv_map, curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
         # uv_unposed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
         # uv_unposed_no_shape = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed_no_shape[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
-        uv_unposed_no_shape_one_hot = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed_no_shape_one_hot, neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border, id=id)
+        # uv_unposed_no_shape_one_hot = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed_no_shape_one_hot, neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border, id=id)
+        uv_posed_no_orient = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_posed_no_orient, neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border, id=id)
         # uv_shaped_body = convert_to_uv(uv_res, m2uv, uv_map, shaped_body, neighbours=3, mask=None, method='rbf', garment_uv_mask=None).astype(np.float32)
         print(time.perf_counter() - start)
 
diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py
index 00e76b2..77508c1 100644
--- a/custom_utils/custom_utils.py
+++ b/custom_utils/custom_utils.py
@@ -144,9 +144,25 @@ def maskFacesBatch(V, F, mask):
     return Vs, Fs
 
 
-def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
+def flip_axis(x, ind=None):
+    y = x.copy()
+    if ind is not None:
+        y[ind:,1] = x[ind:,2]
+        y[ind:,2] = x[ind:,1]
+    else:
+        y[:,1] = x[:,2]
+        y[:,2] = x[:,1]
+    return y
+
+
+def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color, flip_xy=False):
     fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ])
 
+    if flip_xy:
+        mesh_gt = flip_axis(mesh_gt)
+        mesh_pred_unposed = flip_axis(mesh_pred_unposed)
+        collision = flip_axis(collision, ind=14475)
+
     fig.add_trace(go.Mesh3d(
                 # 8 vertices of a cube
                 x=mesh_gt[:,0],
@@ -200,8 +216,24 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli
     return fig
 
 
-def plot_3d_collision(collision_1, collision_2, collision_faces, collision_color):
-    fig = make_subplots(rows=1, cols=2, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(2)] ])
+def plot_3d_collision(gt, collision_1, collision_2, collision_faces, collision_color, flip_xy=False):
+    fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ])
+
+    if flip_xy:
+        collision_1 = flip_axis(collision_1, ind=14475)
+        collision_2 = flip_axis(collision_2, ind=14475)
+        gt = flip_axis(gt, ind=14475)
+
+    fig.add_trace(go.Mesh3d(
+                x=gt[:,0],
+                y=gt[:,2],
+                z=gt[:,1],
+                i = collision_faces[:,0],
+                j = collision_faces[:,1],
+                k = collision_faces[:,2],
+                vertexcolor=collision_color
+            ),
+            row=1, col=1)
 
     fig.add_trace(go.Mesh3d(
                 x=collision_1[:,0],
@@ -212,7 +244,7 @@ def plot_3d_collision(collision_1, collision_2, collision_faces, collision_color
                 k = collision_faces[:,2],
                 vertexcolor=collision_color
             ),
-            row=1, col=1)
+            row=1, col=2)
 
     fig.add_trace(go.Mesh3d(
                 x=collision_2[:,0],
@@ -223,17 +255,19 @@ def plot_3d_collision(collision_1, collision_2, collision_faces, collision_color
                 k = collision_faces[:,2],
                 vertexcolor=collision_color
             ),
-            row=1, col=2)
+            row=1, col=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),
                     )
     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,
                         )
 
     return fig
@@ -258,7 +292,7 @@ def sparse_to_tensor(L):
     return tf.SparseTensor(idx, data, shape)
 
 
-def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False):
+def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False, iteration=''):
     ckpt = tf.train.Checkpoint(model=model)
     manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
     if not disable_assert:
@@ -267,7 +301,12 @@ def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False
     if assert_consumed_on:
         ckpt.restore(manager.latest_checkpoint).assert_consumed()
     else:
-        ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed()
+        if iteration != '':
+            name_base = manager.latest_checkpoint.split('-')[0]
+            ckpt_ind = [i for i in range(len(manager.checkpoints)) if manager.checkpoints[i].endswith(name_base + '-' + iteration)][0]
+            ckpt.restore(manager.checkpoints[ckpt_ind]).assert_existing_objects_matched()#.assert_consumed()
+        else:
+            ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed()
 
 
 def get_summary_str(model):
@@ -350,4 +389,17 @@ def replace_arrays(string):
 
     out += string[i+1:]
 
-    return out
\ No newline at end of file
+    return out
+
+
+def group_inputs_for_levels(batch, config):
+    levels_input = []
+    for res in config['pyramid']['resolutions']:
+        res_modifier = '_' + str(res) if res != 512 else ''
+        if config['use_skinning']:
+            uv_orig = batch['uv_unposed' + res_modifier]
+        else:
+            uv_orig = batch['uv_posed' + res_modifier]
+        levels_input.append([uv_orig, batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier], batch['uv_mask' + res_modifier], batch['uv_mask']])
+    conds = [batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature']]
+    return levels_input, conds
\ No newline at end of file
diff --git a/custom_utils/smpl.py b/custom_utils/smpl.py
index f888576..77d09c6 100644
--- a/custom_utils/smpl.py
+++ b/custom_utils/smpl.py
@@ -34,12 +34,14 @@ eval_joints = {
         
 class SMPL:
 
-    def __init__(self, orig_smpl=False, use_one_hot=False):
+    def __init__(self, orig_smpl=False, use_one_hot=False, is_skirt=False):
         self.J_regressor = []
         self.v_template = []
         self.shapedirs = []
         self.posedirs = []
         self.weights = []
+        self.use_one_hot = use_one_hot
+        self.is_skirt = is_skirt
         # Load SMPL models
         # path = os.path.abspath(os.path.dirname(__file__))
         with open(os.path.join('custom_utils/super_model_f.pkl'), 'rb') as f:
@@ -299,7 +301,22 @@ class SMPL:
 
     def run_garment(self, pose, betas, v_cloth, v_shaped=None):
         J = self.get_joints_template(betas, v_shaped)            
-        R = self.run_rodrigues(pose)
+
+        if self.use_one_hot and self.is_skirt:
+            # p_avg = (pose[:, 3:6] + pose[:, 6:9]) / 2
+            # pose_dress = tf.concat([pose[:, :3], p_avg, p_avg, pose[:, 9:]], axis=-1)
+            # R = self.run_rodrigues(pose_dress).numpy()
+            # for b in range(betas.shape[0]):
+            #     R[b][[4,5,7,8,10,11],:,:] = np.eye(3)
+            
+            pose_dress = pose
+            # p_avg = (pose_dress[:, 3:6] + pose_dress[:, 6:9]) / 2
+            # pose_dress[:,3:9] = np.concatenate([p_avg, p_avg], axis=-1)
+            R = self.run_rodrigues(pose_dress).numpy()
+            for b in range(betas.shape[0]):
+                R[b][[4,5,7,8,10,11],:,:] = np.eye(3)
+        else:
+            R = self.run_rodrigues(pose)
         
         # if filter is not None:
         #     weights = tf.boolean_mask(self.weights, filter)
diff --git a/cvae.py b/cvae.py
index 4776857..fedb790 100644
--- a/cvae.py
+++ b/cvae.py
@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.config_cvae_test import config
+from configs.config_abl_norm import config
 
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import tensorflow as tf
@@ -26,7 +26,6 @@ import numpy as np
 import plotly.graph_objects as go
 import plotly.express as ex
 from PIL import Image
-from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
 from postprocessing import fix_collision
 
 if config['is_skirt']:
@@ -36,30 +35,30 @@ else:
 
 
 
-from losses import collision_ratio, kl_loss #, concept_loss
+from losses import collision_ratio, kl_loss, normal_loss #, concept_loss
 # from losses import collision_penalty as collision_penalty
 from losses import collision_loss as collision_penalty
-from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d
+from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.smpl import SMPL
 from models.cvae_model import CVAE
 
 
 def init_wandb():
-    run = wandb.init(project='cvae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp')
+    run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp')
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae.py')
     artifact.add_file('models/cvae_model.py')
     artifact.add_file('losses.py')
-    artifact.add_file('configs/config_base.py')
-    artifact.add_file('configs/config_cvae.py')
+    # artifact.add_file('configs/config_base.py')
+    # artifact.add_file('configs/config_cvae.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_name = [key for key in sys.modules.keys()  if key.startswith('configs.')]
-    assert len(config_file_name) == 1
-    artifact.add_file('/'.join(config_file_name[0].split('.')) + '.py')
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    for config_file_name in config_file_names:
+        artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
     run.log_artifact(artifact)
 
 uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False)
@@ -76,28 +75,42 @@ def kl_loss(mean, logvar):
 def train_step(batch):
 
     with tf.GradientTape() as tape:
-        pred, mean, logvar = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
+        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']
+
+        pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
 
         mesh_pred = uv2mesh(pred)
-        mesh_orig = batch['mesh_unposed']
 
-        shape = batch['pose_shape'][:, 72:]
-        gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
-        shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
+        if config['use_skinning']:
+            shape = batch['pose_shape'][:, 72:]
+            gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
+            shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
 
-        mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
-        mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
+            mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
+            mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
+        else:
+            mesh_pred_posed = mesh_pred
 
         # is the output already masked?
 
         e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-        c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
-        l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
+        c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+        c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
+        l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
         l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
+        n_loss = normal_loss(mesh_pred, mesh_orig, batch['mesh_mask'], is_skirt=config['is_skirt'])
         # conc_loss = concept_loss(pred, batch['uv_unposed'])
         kl = kl_loss(mean, logvar)
 
-        reconstr_loss = (l2_loss + l2_mesh_loss) * config['l2_coeff'] + (e_loss + c_loss) * config['reg_loss_coeff']
+        reconstr_loss = l2_loss
+        if config['use_extra_losses']:
+            reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
+        if config['use_norm_loss'] and step > 2000:
+            reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
+
         loss = reconstr_loss + g_balancing_coeff * kl
 
         mean = tf.reduce_mean(mean)
@@ -105,41 +118,54 @@ def train_step(batch):
 
     # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
     metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
-    # c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
     gradients = tape.gradient(loss, model.trainable_variables)
     optimizer.apply_gradients(zip(gradients, model.trainable_variables))
     return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/euclid_error': metric,
-             'train/reconstr_loss': reconstr_loss, 
+             'train/reconstr_loss': reconstr_loss, 'train/collision_ratio': c_ratio,
              'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
-             'train/l2_loss': l2_loss, 'train/l2_mesh_loss': l2_mesh_loss 
-             # 'train/ms_ssim': ms_ssim, 'train/collision_ratio': c_ratio,
+             'train/l2_loss': l2_loss, 'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss,
+             # 'train/ms_ssim': ms_ssim, 
             }
 
 
 @tf.function
 def valid_step(batch):
-    pred, mean, logvar = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
+    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']
+
+    pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
 
     mesh_pred = uv2mesh(pred)
-    mesh_orig = batch['mesh_unposed']
 
-    shape = batch['pose_shape'][:, 72:]
-    gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
-    shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
+    if config['use_skinning']:
+        shape = batch['pose_shape'][:, 72:]
+        gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
+        shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
 
-    mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
-    mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
+        mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
+        mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
+    else:
+        mesh_pred_posed = mesh_pred
 
     # is the output already masked?
 
     e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-    c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
-    l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
+    c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+    c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
+    l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
     l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
+    n_loss = normal_loss(mesh_pred, mesh_orig, batch['mesh_mask'], is_skirt=config['is_skirt'])
     # conc_loss = concept_loss(pred, batch['uv_unposed'])
     kl = kl_loss(mean, logvar)
 
-    reconstr_loss = (l2_loss + l2_mesh_loss) * config['l2_coeff'] + (e_loss + c_loss) * config['reg_loss_coeff']
+    reconstr_loss = l2_loss
+    if config['use_extra_losses']:
+        reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
+    if config['use_norm_loss']:
+        reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
+
     loss = reconstr_loss + g_balancing_coeff * kl
 
     mean = tf.reduce_mean(mean)
@@ -150,49 +176,62 @@ def valid_step(batch):
     # c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
 
     return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/euclid_error': metric, 
-             'valid/reconstr_loss': reconstr_loss, 'valid/l2_mesh_loss': l2_mesh_loss,
+             'valid/reconstr_loss': reconstr_loss, 'valid/l2_mesh_loss': l2_mesh_loss, 'valid/collision_ratio': c_ratio,
              'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar,
-             # 'valid/ms_ssim': ms_ssim, 'valid/collision_ratio': c_ratio,
+              'valid/normal_loss': n_loss,# 'valid/ms_ssim': ms_ssim, 
               }
 
 
 @tf.function
 def test_step(batch):
-    pred, mean, logvar = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
+    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']
+    
+    pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
 
     mesh_pred = uv2mesh(pred)
-    mesh_orig = batch['mesh_unposed']
 
-    shape = batch['pose_shape'][:, 72:]
-    gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
-    shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
+    if config['use_skinning']:
+        shape = batch['pose_shape'][:, 72:]
+        gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
+        shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
 
-    mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
-    mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
+        mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
+        mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
+        mesh_orig_shaped = mesh_orig + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
+        mesh_orig_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_orig_shaped)
+    else:
+        mesh_pred_posed = mesh_pred
+        mesh_orig_posed = mesh_orig
+    
     mesh_pred_posed_fixed = fix_collision(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
-    mesh_orig_shaped = mesh_orig + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
-    mesh_orig_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_orig_shaped)
+
 
     # is the output already masked?
 
     e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-    c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
-    l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
+    c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+    c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
+    l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
     l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
     # conc_loss = concept_loss(pred, batch['uv_unposed'])
     kl = kl_loss(mean, logvar)
 
-    reconstr_loss = (l2_loss + l2_mesh_loss) * config['l2_coeff'] + (e_loss + c_loss) * config['reg_loss_coeff']
+    reconstr_loss = l2_loss
+    if config['use_extra_losses']:
+        reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
+            
     loss = reconstr_loss + g_balancing_coeff * kl
 
     mean = tf.reduce_mean(mean)
     logvar = tf.reduce_mean(logvar)
 
-    ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
+    ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, uv_orig + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
     metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
     euclid_posed = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_posed - mesh_pred_posed) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
     euclid_posed_fixed = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_posed - mesh_pred_posed_fixed) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
-    c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
 
     return { 'test/loss': loss, 'test/edge_loss': e_loss, 'test/collision_loss': c_loss, 'test/euclid_error': metric, 'test/ms_ssim': ms_ssim,
              'test/reconstr_loss': reconstr_loss, 'test/collision_ratio': c_ratio, 'test/l2_mesh_loss': l2_mesh_loss, 
@@ -201,32 +240,46 @@ def test_step(batch):
 
 
 def plot_everything(batch: dict, type: str):
-    pred, mean, logvar = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
+    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']
+
+    pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
     
     mesh_pred = uv2mesh(pred).numpy()
-    mesh_orig = batch['mesh_unposed']
     
-    orig_img = batch['uv_unposed']
+    orig_img = uv_orig
     uv_static_img = batch['uv_static']
     uv_body_posed_img = batch['uv_body_posed']
     pred_img = pred
     img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
     imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
     
-    shape = batch['pose_shape'][:, 72:]
-    gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
-    shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
-
     F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy")
     mask = batch['mesh_mask'].numpy().astype(bool)
-    garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
-    garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
-    gt = [mesh_orig[i].numpy()[mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(pred.shape[0])]
+
+    if config['use_skinning']:
+        shape = batch['pose_shape'][:, 72:]
+        gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
+        shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
+
+        garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
+        garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
+        gt = [mesh_orig[i].numpy()[mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(pred.shape[0])]
     
+        for i in range(mesh_orig.shape[0]):
+            mesh_pred[i, mask[i], :] += shape_offsets.numpy()[i, mask[i], :]
+        mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
+    else:
+        mesh_posed = mesh_pred
+        gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
+        garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
+
+    gt_posed = [batch['mesh_posed'].numpy()[i, mask[i], :] for i in range(mask.shape[0])]
+
+
     body = batch['mesh_body_posed'].numpy()
-    for i in range(body.shape[0]):
-        mesh_pred[i, mask[i], :] += shape_offsets.numpy()[i, mask[i], :]
-    mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
     mesh_posed_fixed = fix_collision(mesh_posed, body, mask).numpy()
 
     garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
@@ -236,9 +289,10 @@ def plot_everything(batch: dict, type: str):
 
     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(garment_posed))]
+    collision_gt = [np.concatenate((body[i], gt_posed[i]), axis=0) for i in range(len(garment_posed))]
 
-    figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
-    figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision[i], collision_fixed[i], collision_faces[i], collision_color[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], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
+    figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
 
     mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0), num_bins=model.latent_dim)
     logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0), num_bins=model.latent_dim)
@@ -246,7 +300,7 @@ def plot_everything(batch: dict, type: str):
     return { **figs, **figs_coll, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
 
 
-def sample(folder):
+def sample(folder): 
     # run = wandb.init(project='vae_plots', name=folder.split('/')[-1])
     load_model(model,folder)
     count = 0
@@ -417,7 +471,7 @@ def test():
 
 
 if __name__ == '__main__':
-    # init_wandb()
+    init_wandb()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)
     args = vars(parser.parse_args())
@@ -488,10 +542,10 @@ if __name__ == '__main__':
         if step == 0:
             to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
 
-        if train_loss['train/euclid_error'] < 0.019:
+        if train_loss['train/euclid_error'] < 0.017:
             if balancing_activated_step == 0:
                 balancing_activated_step = step
-            g_balancing_coeff = g_balancing_coeff + 1 / 10000      
+            g_balancing_coeff = g_balancing_coeff + 1 / 20000      
 
         to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } }
         
diff --git a/cvae_pyramid.py b/cvae_pyramid.py
index e4ba70b..7018d29 100644
--- a/cvae_pyramid.py
+++ b/cvae_pyramid.py
@@ -1,5 +1,7 @@
 import os
-from configs.config_pyramid import config
+import sys
+from configs.config_abl_pyr_64_128_cape_m import config
+
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import tensorflow as tf
 
@@ -22,7 +24,7 @@ import wandb
 import random
 import numpy as np
 import scipy as sp
-from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, get_summary_str
+from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, get_summary_str, plot_3d_collision, group_inputs_for_levels
 import plotly.graph_objects as go
 import plotly.express as ex
 from plotly.subplots import make_subplots
@@ -33,7 +35,6 @@ if config['is_skirt']:
 else:
     from losses import edge_loss
 
-from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
 from losses import uv_mask_loss, bend_loss, kl_loss, collision_ratio
 # from losses import collision_penalty as collision_penalty
 from losses import collision_loss as collision_penalty
@@ -41,24 +42,26 @@ from custom_utils.custom_utils import batch_to_dict, load_model
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.smpl import SMPL
 from models.pyramid_model_optimized import Pyramid
+from postprocessing import fix_collision
 
 
 def init_wandb():
-    run = wandb.init(project='cvae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp')
+    run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/wandb')
 
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae_pyramid.py')
     artifact.add_file('models/cvae_model.py')
-    artifact.add_file('models/pyramid_model.py')
+    artifact.add_file('models/pyramid_model_optimized.py')
     artifact.add_file('models/pyramid_level.py')
     artifact.add_file('models/upscaler.py')
     artifact.add_file('losses.py')
-    artifact.add_file('configs/config_base.py')
-    artifact.add_file('configs/config_pyramid.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:
+        artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
     run.log_artifact(artifact)
 
 uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
@@ -205,26 +208,18 @@ def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, col
     return fig
 
 
-def group_inputs_for_levels(batch):
-    levels_input = []
-    for res in config['pyramid']['resolutions']:
-        res_modifier = '_' + str(res) if res != 512 else ''
-        levels_input.append([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier], batch['uv_mask' + res_modifier], batch['uv_mask']])
-    conds = [batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature']]
-    return levels_input, conds
-
 
 @tf.function
 def train_step(batch):
-    levels_input, conds = group_inputs_for_levels(batch)
+    levels_input, conds = group_inputs_for_levels(batch, config)
 
     with tf.GradientTape() as tape:
         out_levels = model([levels_input, conds], training=True)
 
-        shape = batch['pose_shape'][:, 72:]
-        gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
-        shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
-        mesh_orig = batch['mesh_unposed']
+        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
@@ -237,21 +232,36 @@ def train_step(batch):
         pred, mean, logvar = out_levels[i]
 
         mesh_pred = uv2mesh(pred)
-        mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
-        mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
+
+        if config['use_skinning']:
+            shape = batch['pose_shape'][:, 72:]
+            gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
+            shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
+
+            mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
+            mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
+        else:
+            mesh_pred_posed = mesh_pred
 
         # is the output already masked?
 
         e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-        c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
-        l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
+        c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+        c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
+        l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
         # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
-        ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
+        # ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
         l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
         kl = kl_loss(mean, logvar)
 
-        reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
-        # reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
+        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'] 
+        else:
+            reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
+
+        
         kl_total += kl
 
         mean = tf.reduce_mean(mean)
@@ -260,9 +270,9 @@ def train_step(batch):
         metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
 
         output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,'train/euclid_error' + res_modifier: metric, 
-                                    'train/ms_ssim' + res_modifier: ms_ssim, '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/euclid_error': metric, 'train/l2_mesh_loss': l2_mesh_loss } }
+                                 '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/euclid_error': metric, 'train/l2_mesh_loss': l2_mesh_loss } } # 'train/ms_ssim' + res_modifier: ms_ssim,
 
         loss = reconstr_loss_total + balancing_coeff * kl_total
 
@@ -273,13 +283,14 @@ def train_step(batch):
 
 @tf.function
 def valid_step(batch):
-    levels_input, conds = group_inputs_for_levels(batch)
+    levels_input, conds = group_inputs_for_levels(batch, config)
     out_levels = model([levels_input, conds], training=False)
 
-    shape = batch['pose_shape'][:, 72:]
-    gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
-    shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
-    mesh_orig = batch['mesh_unposed']
+
+    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
@@ -292,21 +303,34 @@ def valid_step(batch):
     pred, mean, logvar = out_levels[i]
 
     mesh_pred = uv2mesh(pred)
-    mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
-    mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
 
+    if config['use_skinning']:
+        shape = batch['pose_shape'][:, 72:]
+        gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
+        shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
+
+        mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
+        mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
+    else:
+        mesh_pred_posed = mesh_pred
     # is the output already masked?
 
     e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-    c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
-    l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
+    c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+    c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
+    l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
     # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
-    ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
+    # ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
     l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
     kl = kl_loss(mean, logvar)
 
-    reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
-    # reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
+    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'] 
+    else:
+        reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
+    
     kl_total += kl
 
     mean = tf.reduce_mean(mean)
@@ -315,9 +339,9 @@ def valid_step(batch):
     metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
 
     output = { **output, **{ 'valid/edge_loss' + res_modifier: e_loss, 'valid/collision_loss' + res_modifier: c_loss,'valid/euclid_error' + res_modifier: metric, 
-                                'valid/ms_ssim' + res_modifier: ms_ssim, 'valid/mean' + res_modifier: mean, 'valid/logvar' + res_modifier: logvar, 'valid/kl_loss' + res_modifier: kl,
+                                 '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/euclid_error': metric, 'valid/l2_mesh_loss': l2_mesh_loss } }
+                                'valid/euclid_error': metric, 'valid/l2_mesh_loss': l2_mesh_loss } } # 'valid/ms_ssim' + res_modifier: ms_ssim,
 
     loss = reconstr_loss_total + balancing_coeff * kl_total
 
@@ -326,7 +350,7 @@ def valid_step(batch):
 
 @tf.function
 def test_step(batch):
-    levels_input, conds = group_inputs_for_levels(batch)
+    levels_input, conds = group_inputs_for_levels(batch, config)
     out_levels = model([levels_input, conds], training=False)
 
     shape = batch['pose_shape'][:, 72:]
@@ -351,14 +375,17 @@ def test_step(batch):
     # is the output already masked?
 
     e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-    c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
-    c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
+    c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+    c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
     l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
+    l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
     # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
-    ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
+    # ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
     kl = kl_loss(mean, logvar)
 
-    reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
+    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']
     # reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
     kl_total += kl
 
@@ -369,7 +396,7 @@ def test_step(batch):
     euclid_error_per_garment = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1)), axis=1) / tf.reduce_sum(batch['mesh_mask'], axis=1)
 
     output = { **output, **{ 'test/edge_loss' + res_modifier: e_loss, 'test/collision_loss' + res_modifier: c_loss,'test/euclid_error' + res_modifier: metric, 
-                                'test/ms_ssim' + res_modifier: ms_ssim, 'test/mean' + res_modifier: mean, 'test/logvar' + res_modifier: logvar, 'test/kl_loss' + res_modifier: kl,
+                                'test/mean' + res_modifier: mean, 'test/logvar' + res_modifier: logvar, 'test/kl_loss' + res_modifier: kl,
                                 'test/var' + res_modifier: np.e ** logvar, 'test/l2_loss' + res_modifier: l2_loss,
                                 'test/euclid_error': metric, 'test/collision_ratio' + res_modifier: c_ratio } }
 
@@ -380,15 +407,19 @@ def test_step(batch):
 
 
 def plot_everything(batch: dict, type: str):
-    levels_input = group_inputs_for_levels(batch)
+    if config['use_skinning']:
+        mesh_orig, uv_orig = batch['mesh_unposed'].numpy(), batch['uv_unposed'].numpy()
+    else:
+        mesh_orig, uv_orig = batch['mesh_posed'].numpy(), batch['uv_posed'].numpy()
+    
+    levels_input = group_inputs_for_levels(batch, config)
     output_levels = model(levels_input, training=False)
 
     mesh_pred_levels = [uv2mesh(output_levels[i][0]).numpy() for i in range(len(output_levels))]
-    mesh_orig = batch['mesh_unposed'].numpy()
     mesh_orig_posed = batch['mesh_posed'].numpy()
     body = batch['mesh_body_posed'].numpy()
     
-    orig_img = (batch['uv_unposed'].numpy() + 1) / 2
+    orig_img = (uv_orig + 1) / 2
     uv_static_img = (batch['uv_static'].numpy() + 1) / 2
     uv_body_posed_img = (batch['uv_body_posed'].numpy() + 1) / 2
     uv_curvature = (np.tile(batch['uv_curvature'].numpy().argmax(-1)[:,:,:,np.newaxis], [1,1,1,3]) / 6)
@@ -397,43 +428,62 @@ def plot_everything(batch: dict, type: str):
     img = tf.concat((uv_static_img, uv_body_posed_img, uv_curvature, orig_img, pred_img), axis=2)
     imgs = [wandb.Image(img[i]) for i in range(4)]
 
-    shape = batch['pose_shape'][:, 72:]
-    gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
-    shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]).numpy()
-
     F = np.load("custom_utils/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
     mask = batch['mesh_mask'].numpy().astype(bool)
-
     _, faces = maskFacesBatch(np.ones_like(mesh_orig), F, mask)
-    gt = [mesh_orig[i][mask[i], :] + shape_offsets[i, mask[i], :] for i in range(mask.shape[0])]
+    
+    if config['use_skinning']:
+        shape = batch['pose_shape'][:, 72:]
+        gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
+        shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]).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.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_unposed_shaped_full).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_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.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_unposed_shaped_full).numpy()
+    mesh_posed_fixed = fix_collision(mesh_posed, body, mask).numpy()
+
     mesh_posed = [mesh_posed[j, mask[j], :] for j in range(mask.shape[0])]
-
-    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))]
-
-    # body_2 = [b.copy() for b in body]
-    # for i in range(len(gt_posed)):
-    #     body_2[i][:, 1] = body[i][:,2].copy()
-    #     body_2[i][:, 2] = body[i][:,1].copy()
-    # gt_posed_with_body = [np.concatenate((body_2[i], gt_posed[i]), axis=0) for i in range(len(gt_posed))]
+    
+    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'].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) }
 
-    return { **figs, **{ type + '/images': imgs, } }
+    return { **figs, **figs_coll, **{ type + '/images': imgs, } }
 
 
 def sample(folder):
@@ -622,7 +672,6 @@ def test():
 
 if __name__ == '__main__':
     init_wandb()
-    allow_memory_growth()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)
     args = vars(parser.parse_args())
@@ -634,10 +683,10 @@ if __name__ == '__main__':
         dataset_resolutions.append(512)
 
     factory = DatasetFactory()
-    dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
-    static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
-    dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
-    static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
+    dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
+    static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
+    dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
+    static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
 
     model: Pyramid = Pyramid(config)
 
@@ -654,11 +703,18 @@ if __name__ == '__main__':
     ckpt = tf.train.Checkpoint(model=model)
     manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=20)
 
-    # load_model(model, config['pyramid']['pretrained_pyramid'])
-    model.static_enc.trainable = False
-    model.body_enc.trainable = False
-    model.curv_enc.trainable = False
-    # model.cond_encoder.trainable = False
+    # load_model(model, 'checkpoints/pyr_abl_32_skirt', iteration='180000')
+
+    if len(config['pyramid']['resolutions']) == 1:
+        model.static_enc.trainable = False
+        model.body_enc.trainable = False
+        model.curv_enc.trainable = False
+    else:
+        load_model(model, config['pyramid']['pretrained_pyramid'])
+        model.static_enc.trainable = False
+        model.body_enc.trainable = False
+        model.curv_enc.trainable = False
+        model.cond_encoder.trainable = False
 
 
     # load_model(model.levels[0].upscale_enhancer, config['pyramid']['levels'][0]['upscale_enhancer'])
@@ -669,14 +725,15 @@ if __name__ == '__main__':
     for i in range(len(model.levels)-1):
         load_model(model.levels[i].upscale_enhancer, config['pyramid']['levels'][i]['upscale_enhancer'])
         model.levels[i].upscale_enhancer.trainable = False
-    load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
-    model.upscale_enhancer.trainable = False
+    if config['pyramid']['resolutions'][-1] != 512:
+        load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
+        model.upscale_enhancer.trainable = False
 
     for i in range(len(model.levels)-1):
         model.levels[i].trainable = False
 
-    # load_model(model.cond_encoder, 'checkpoints/cvae_low_4_3_d32')
-    # load_model(model.levels[0], 'checkpoints/cvae_low_4_3_d32')
+    # load_model(model.cond_encoder, config['pyramid']['pretrained_pyramid'])
+    # load_model(model.levels[0], config['pyramid']['pretrained_pyramid'])
     # load_model(model.levels[-1], 'checkpoints/cvae_19')
 
 
@@ -701,10 +758,10 @@ if __name__ == '__main__':
         if step == 0:
             to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
 
-        if train_loss['train/euclid_error'] < 0.01:
+        if train_loss['train/euclid_error'] < 0.004:
             if balancing_activated_step == 0:
                 balancing_activated_step = step
-            balancing_coeff = balancing_coeff + 1 / 15000
+            balancing_coeff = balancing_coeff + 1 / 20000
 
 
         to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } }
diff --git a/dataset.py b/dataset.py
index ad7a937..7994b26 100644
--- a/dataset.py
+++ b/dataset.py
@@ -160,6 +160,8 @@ class DatasetFactory():
         for key in signature.keys():
             if key == 'id':
                 tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
+            elif key == 'frame_nr':
+                tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
             elif key == 'mesh_curvature':
                 tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.int32))
             else:
@@ -209,14 +211,13 @@ class DatasetFactory():
 
             if 'id' in signature.keys():
                 item['id'] = id
+            if 'frame_nr' in signature.keys():
+                item['frame_nr'] = frame_nr
             if 'pose_shape' in signature.keys():
                 pose_shape = np.load(os.path.join(folder, 'pose_shape',  id + separator + frame_nr + '.npy'))
                 if use_no_orient:
                     pose_shape[:3] = 0
                 item['pose_shape'] = pose_shape
-            if 'uv_posed' in signature.keys():
-                uv_posed = np.load(os.path.join(folder, 'uv_posed', id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']
-                item['uv_posed'] = uv_posed
             if 'mesh_body_posed' in signature.keys():
                 if use_no_orient:
                     mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + separator + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
@@ -225,8 +226,12 @@ class DatasetFactory():
                     mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
                     item['mesh_body_posed'] = mesh_body_posed
             if 'mesh_posed' in signature.keys():
-                mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + separator + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed']
-                item['mesh_posed'] = mesh_posed
+                if use_no_orient:
+                    mesh_posed = np.load(os.path.join(folder, 'mesh_posed_no_orient', id + separator + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed_no_orient'].astype(np.float32)
+                    item['mesh_posed'] = mesh_posed
+                else:
+                    mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + separator + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed'].astype(np.float32)
+                    item['mesh_posed'] = mesh_posed
             # if 'mesh_unposed' in signature.keys():
             #     mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
             #     item['mesh_unposed'] = mesh_unposed
@@ -269,7 +274,17 @@ class DatasetFactory():
                         item['uv_body_posed' + res_string] = uv_posed_body_no_orient
                     else:
                         uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']
-                        item['uv_body_posed'] = uv_body_posed
+                        item['uv_body_posed' + res_string] = uv_body_posed
+
+                if 'uv_posed' + res_string in signature.keys():
+                    if use_no_orient:
+                        uv_posed = np.load(os.path.join(folder, 'uv_posed_no_orient' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient']
+                        item['uv_posed' + res_string] = uv_posed
+                    else:
+                        uv_posed = np.load(os.path.join(folder, 'uv_posed' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']
+                        item['uv_posed' + res_string] = uv_posed
+
+            # print('here')
 
             # item = { **static_info, **item }
             as_list = []
@@ -330,7 +345,7 @@ if __name__ == '__main__':
     #     print(d)
 
     factory = DatasetFactory()
-    signature_dynamic = { 'id': (),
+    signature_dynamic = { 'id': (), 'frame_nr': (),
                         #   'pose_shape': (72 + 11,),
                         #   'uv_posed': (3, 512, 256,),
                         #   'uv_body_posed': (512, 256, 3),
@@ -351,16 +366,14 @@ if __name__ == '__main__':
                         #  'uv_shaped_body': (3, 512, 256,),
                           }
 
-    dataset = factory.get_dataset(folder='/home/workspace/local-data/train',
+    dataset = factory.get_dataset(folder='/home/workspace/local-data/train_skirt',
                                   batch_size=300,
                                   signature=signature_dynamic,
-                                  use_filled=False,
                                   resolutions=[512],
                                   do_repeat=False,
-                                  use_one_hot=True)
-    static_dataset = StaticDataset(folder='/home/workspace/local-data/train',
+                                  use_one_hot=False)
+    static_dataset = StaticDataset(folder='/home/workspace/local-data/train_skirt',
                                    signature=signature_static,
-                                   use_filled=False,
                                    resolutions=[512])
 
 
@@ -393,7 +406,7 @@ if __name__ == '__main__':
                 k = F[:,2]
             )]).show()
 
-    np.save('custom_data/train_no_shape_one_hot_avg_uv.npy', avg_garment)
+    np.save('custom_data/train_no_shape_skirt_avg_uv.npy', avg_garment)
 
     # print()
 
diff --git a/deorient.py b/deorient.py
new file mode 100644
index 0000000..88c29f6
--- /dev/null
+++ b/deorient.py
@@ -0,0 +1,104 @@
+
+
+import os
+import sys
+
+
+os.environ["CUDA_VISIBLE_DEVICES"] = '6'
+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 glob
+import tqdm
+import copy
+import numpy as np
+from custom_utils.smpl import SMPL
+
+# from custom_utils.custom_utils import UV2Mesh, _maskFaces
+import plotly.graph_objects as go
+# import plotly.express as ex
+from scipy.spatial.transform import Rotation as R
+
+
+if __name__ == '__main__':
+    folder = '/home/workspace/local-data/test_skirt_ready'
+    # folder = '/home/workspace/data/dataset/test_ready'
+    # uv2mesh = UV2Mesh(is_skirt=False)
+    smpl = SMPL(use_one_hot=False)
+
+    F = np.load("custom_utils/smpl_faces.npy")
+
+    ids = os.listdir(os.path.join(folder, 'pose_shape'))
+    ids = [id.split('_')[0] for id in ids]
+    ids = set(ids)
+    os.makedirs(os.path.join(folder, 'mesh_posed_no_orient'), exist_ok=True)
+
+    for id in tqdm.tqdm(ids):
+
+        frames = glob.glob(os.path.join(folder, 'pose_shape', id + '_*'))
+        frames = len(frames)
+
+        mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')).astype(bool)
+        mesh_posed = [np.load(os.path.join(folder, 'mesh_posed',  id + '_' + str(i) + '.npz')) for i in range(frames)]
+        mesh_posed = [mesh_posed[i]['mesh_posed'] for i in range(frames)]
+        # mesh_unposed = [np.load(os.path.join(folder, 'mesh_unposed',  id + '_' + str(i) + '.npz')) for i in range(frames)]
+        # mesh_unposed = [mesh_unposed[i]['mesh_unposed'] for i in range(frames)]
+        pose_shape = [np.load(os.path.join(folder, 'pose_shape',  id + '_' + str(i) + '.npy')).astype(np.float32) for i in range(frames)]
+        # pose_shape_no_orient = copy.deepcopy(pose_shape)
+        # for i in range(len(pose_shape)):
+        #     pose_shape_no_orient[i][:3] = 0
+
+        rodr = smpl.run_rodrigues(np.stack([x[:72] for x in pose_shape]).astype(np.float32)).numpy()
+
+        for frame in range(frames):
+            r = R.from_matrix(rodr[frame, 0, :, :])
+            # assert r.as_matrix() == rodr[frame, 0, :, :]
+            out = r.apply(mesh_posed[frame], inverse=True)
+
+            # x = smpl.run_garment(pose_shape_no_orient[frame][:72][np.newaxis,:], pose_shape_no_orient[frame][72:][np.newaxis,:], mesh_unposed[frame][np.newaxis,:,:])[0].numpy()[mask, :]
+            # y = out[mask, :]
+            # _, faces = _maskFaces(np.zeros((14475,3), np.float32), F, mask)
+            # fig = go.Figure(data=[
+            #         go.Mesh3d(
+            #             x=x[:,0],
+            #             y=x[:,1],
+            #             z=x[:,2],
+            #             i = faces[:,0],
+            #             j = faces[:,1],
+            #             k = faces[:,2],
+            #         ),
+            #         go.Mesh3d(
+            #             x=y[:,0],
+            #             y=y[:,1],
+            #             z=y[:,2],
+            #             i = faces[:,0],
+            #             j = faces[:,1],
+            #             k = faces[:,2],
+            #         ),
+            #     ])
+
+            # fig.update_layout(scene_aspectmode='manual',
+            #                 scene_aspectratio=dict(x=1, y=1, z=1))
+            # fig.update_layout(
+            #     scene = dict(
+            #         xaxis = dict(nticks=4, range=[-1,1],),
+            #                     yaxis = dict(nticks=4, range=[-1,1],),
+            #                     zaxis = dict(nticks=4, range=[-1,1],),)
+            #                     )
+
+            # fig.show()
+
+
+            np.savez_compressed(os.path.join(folder, 'mesh_posed_no_orient', id + '_' + str(frame) + '.npz'), mesh_posed_no_orient=out)
diff --git a/evaluate_ae.py b/evaluate_ae.py
new file mode 100644
index 0000000..709758b
--- /dev/null
+++ b/evaluate_ae.py
@@ -0,0 +1,152 @@
+import os
+import sys
+from configs.config_abl_skin_skirt 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 glob
+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
+from postprocessing import fix_collision
+
+if config['is_skirt']:
+    from losses import edge_loss_skirt as edge_loss
+else:
+    from losses import edge_loss
+
+
+
+from losses import collision_ratio, kl_loss #, concept_loss
+# from losses import collision_penalty as collision_penalty
+from losses import collision_loss as collision_penalty
+from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.smpl import SMPL
+from models.cvae_model import CVAE
+from custom_utils.supersmpl_io import writeOBJ
+
+
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False)
+
+@tf.function
+def step(batch):
+    if config['use_skinning']:
+        uv_orig = batch['uv_unposed']
+    else:
+        uv_orig =  batch['uv_posed']
+    
+    pred, _, _ = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
+
+    mesh_pred = uv2mesh(pred)
+
+    if config['use_skinning']:
+        shape = batch['pose_shape'][:, 72:]
+        gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
+        shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
+
+        mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
+        mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
+    else:
+        mesh_pred_posed = mesh_pred
+        
+    # mesh_pred_posed_fixed = fix_collision(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+
+    return mesh_pred_posed, None #mesh_pred_posed_fixed
+
+
+if __name__ == '__main__':
+    # init_wandb()
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+
+    smpl = SMPL(use_one_hot=config['use_one_hot'])
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+
+    folder = config['test_dir']
+    smpl = SMPL(use_one_hot=False)
+    model = CVAE(config)
+
+    # str_split = config['name'].split('_')
+    # test_idx = str_split.index('test')
+    # checkpoint_name = '_'.join(str_split[:test_idx])
+    checkpoint_name = config['name']
+
+    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
+
+    ids = os.listdir(os.path.join(folder, 'pose_shape'))
+    ids = [id.split('_')[0] for id in ids]
+    ids = set(ids)
+
+    iteration = '70000'
+
+    load_model(model, 'checkpoints/' + checkpoint_name, iteration=iteration)
+    # load_model(model, 'checkpoints/' + 'abl_all_120k')
+
+    output_folder = '/home/workspace/data/results/' + checkpoint_name + '-' + iteration
+
+
+    os.makedirs(output_folder, exist_ok=True)
+
+    if config['use_skinning']:
+        config['signature_dynamic'].pop('mesh_posed')
+        config['signature_dynamic'].pop('mesh_unposed')
+    config['signature_dynamic'].pop('mesh_body_posed')
+
+    factory = DatasetFactory()
+    dataset_test = factory.get_dataset(folder=config['test_dir'], 
+                                       batch_size=config['batch_size'], 
+                                       signature=config['signature_dynamic'], 
+                                       resolutions=[config['resolution']], 
+                                       in_memory=False, 
+                                       do_repeat=False,
+                                       use_one_hot=config['use_one_hot'],
+                                       use_no_orient=config['remove_global_orientation'])
+    static_dataset_test = StaticDataset(folder=config['test_dir'], 
+                                        signature=config['signature_static'], 
+                                        resolutions=[config['resolution']], 
+                                        in_memory=True)
+
+    # faces_folder = '/home/workspace/data/results/support/faces' + ('_skirt' if config['is_skirt'] else '')
+    # ids = os.listdir(faces_folder)
+    # ids = [id.split('.')[0] for id in ids]
+    # faces = {id : np.load(faces_folder + '/' + id + '.npy') for id in ids }
+
+    for batch in tqdm.tqdm(dataset_test):
+        batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
+
+        mesh_pred_posed, mesh_pred_posed_fixed = step(batch)
+        # mesh_pred_posed, mesh_pred_posed_fixed = mesh_pred_posed.numpy(), mesh_pred_posed_fixed.numpy()
+        mesh_pred_posed = mesh_pred_posed.numpy()
+
+        mask = batch['mesh_mask'].numpy().astype(bool)
+
+        for i in range(mask.shape[0]):
+            id, frame_nr = batch['id'][i].numpy().decode("utf-8"), batch['frame_nr'][i].numpy().decode("utf-8")
+            # writeOBJ(os.path.join(output_folder, id + '_' + frame_nr + '.obj'), mesh_pred_posed[i, mask[i], :], faces[id])
+            np.save(os.path.join(output_folder, id + '_' + frame_nr + '.npy'), mesh_pred_posed[i, mask[i], :])
+            # np.save(os.path.join(output_folder, id + '_' + frame_nr + '_fixed_coll.npy'), mesh_pred_posed_fixed[i, mask[i], :])
+            # writeOBJ(os.path.join(output_folder, id + '_' + frame_nr + '_fixed_coll.obj'), mesh_pred_posed_fixed[i, mask[i], :], faces[id])
+
+
+
diff --git a/evaluate_ae_pyramid.py b/evaluate_ae_pyramid.py
new file mode 100644
index 0000000..8f008ff
--- /dev/null
+++ b/evaluate_ae_pyramid.py
@@ -0,0 +1,186 @@
+import os
+import sys
+from configs.config_abl_pyr_32 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 glob
+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
+from models.pyramid_model_optimized import Pyramid
+from postprocessing import fix_collision
+
+if config['is_skirt']:
+    from losses import edge_loss_skirt as edge_loss
+else:
+    from losses import edge_loss
+
+
+
+from losses import collision_ratio, kl_loss #, concept_loss
+# from losses import collision_penalty as collision_penalty
+from losses import collision_loss as collision_penalty
+from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision, group_inputs_for_levels
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.smpl import SMPL
+from models.cvae_model import CVAE
+from custom_utils.supersmpl_io import writeOBJ
+
+
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False)
+
+@tf.function
+def step(batch):
+    levels_input, conds = group_inputs_for_levels(batch, config)
+
+    out_levels = model([levels_input, conds], training=True)
+
+    pred_final = out_levels[-1][0]
+    mesh_pred = uv2mesh(pred_final)
+
+    if config['use_skinning']:
+        shape = batch['pose_shape'][:, 72:]
+        gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
+        shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
+
+        mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
+        mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
+    else:
+        mesh_pred_posed = mesh_pred
+    # mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
+        
+    # mesh_pred_posed_fixed = fix_collision(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+
+    return mesh_pred_posed, None
+
+
+if __name__ == '__main__':
+    # init_wandb()
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+
+    smpl = SMPL()
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+
+    folder = config['test_dir']
+    smpl = SMPL(use_one_hot=False)
+    model: Pyramid = Pyramid(config)
+
+    # str_split = config['name'].split('_')
+    # test_idx = str_split.index('test')
+    # checkpoint_name = '_'.join(str_split[:test_idx])
+    checkpoint_name = config['name']
+
+    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
+
+    ids = os.listdir(os.path.join(folder, 'pose_shape'))
+    ids = [id.split('_')[0] for id in ids]
+    ids = set(ids)
+
+    iteration = '430000'
+
+    load_model(model, 'checkpoints/' + checkpoint_name, iteration=iteration)
+    # load_model(model, 'checkpoints/' + 'abl_all_120k')
+
+    output_folder = '/home/workspace/data/results/' + checkpoint_name + '-' + iteration
+
+
+    os.makedirs(output_folder, exist_ok=True)
+    dataset_resolutions = config['pyramid']['resolutions'].copy()
+    if 512 not in dataset_resolutions:
+        dataset_resolutions.append(512)
+
+    if config['use_skinning']:
+        config['signature_dynamic'].pop('mesh_posed')
+        config['signature_dynamic'].pop('mesh_unposed')
+    config['signature_dynamic'].pop('mesh_body_posed')
+
+    factory = DatasetFactory()
+    dataset_test = factory.get_dataset(folder=config['test_dir'], 
+                                       batch_size=config['batch_size'], 
+                                       signature=config['signature_dynamic'], 
+                                       resolutions=dataset_resolutions, 
+                                       in_memory=False, 
+                                       do_repeat=False,
+                                       use_no_orient=True)
+    static_dataset_test = StaticDataset(folder=config['test_dir'], 
+                                        signature=config['signature_static'], 
+                                        resolutions=dataset_resolutions, 
+                                        in_memory=True)
+
+    faces_folder = '/home/workspace/data/results/support/faces' + ('_skirt' if config['is_skirt'] else '')
+    ids = os.listdir(faces_folder)
+    ids = [id.split('.')[0] for id in ids]
+    faces = {id : np.load(faces_folder + '/' + id + '.npy') for id in ids }
+
+    for batch in tqdm.tqdm(dataset_test):
+        batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
+
+        mesh_pred_posed, _ = step(batch)
+        mesh_pred_posed = mesh_pred_posed.numpy()
+
+        mask = batch['mesh_mask'].numpy().astype(bool)
+
+        for i in range(mask.shape[0]):
+            id, frame_nr = batch['id'][i].numpy().decode("utf-8"), batch['frame_nr'][i].numpy().decode("utf-8")
+            # writeOBJ(os.path.join(output_folder, id + '_' + frame_nr + '.obj'), mesh_pred_posed[i, mask[i], :], faces[id])
+            np.save(os.path.join(output_folder, id + '_' + frame_nr + '.npy'), mesh_pred_posed[i, mask[i], :])
+            # np.save(os.path.join(output_folder, id + '_' + frame_nr + '_fixed_coll.npy'), mesh_pred_posed_fixed[i, mask[i], :])
+            # writeOBJ(os.path.join(output_folder, id + '_' + frame_nr + '_fixed_coll.obj'), mesh_pred_posed_fixed[i, mask[i], :], faces[id])
+
+            # x = batch['mesh_posed'].numpy()[i, mask[i], :]
+            # y = mesh_pred_posed[i, mask[i], :]
+            # fig = go.Figure(data=[
+            #         go.Mesh3d(
+            #             x=x[:,0],
+            #             y=x[:,1],
+            #             z=x[:,2],
+            #             i = faces[id][:,0],
+            #             j = faces[id][:,1],
+            #             k = faces[id][:,2],
+            #         ),
+            #         go.Mesh3d(
+            #             x=y[:,0],
+            #             y=y[:,1],
+            #             z=y[:,2],
+            #             i = faces[id][:,0],
+            #             j = faces[id][:,1],
+            #             k = faces[id][:,2],
+            #         ),
+            #     ])
+
+            # fig.update_layout(scene_aspectmode='manual',
+            #                 scene_aspectratio=dict(x=1, y=1, z=1))
+            # fig.update_layout(
+            #     scene = dict(
+            #         xaxis = dict(nticks=4, range=[-1,1],),
+            #                     yaxis = dict(nticks=4, range=[-1,1],),
+            #                     zaxis = dict(nticks=4, range=[-1,1],),)
+            #                     )
+
+            # fig.show()
+
+            # print()
+
+
diff --git a/evaluate_ae_support.py b/evaluate_ae_support.py
new file mode 100644
index 0000000..7628de4
--- /dev/null
+++ b/evaluate_ae_support.py
@@ -0,0 +1,110 @@
+import os
+import sys
+from configs.config_abl_all 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 glob
+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
+from postprocessing import fix_collision
+
+if config['is_skirt']:
+    from losses import edge_loss_skirt as edge_loss
+else:
+    from losses import edge_loss
+
+
+
+from losses import collision_ratio, kl_loss #, concept_loss
+# from losses import collision_penalty as collision_penalty
+from losses import collision_loss as collision_penalty
+from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.smpl import SMPL
+from models.cvae_model import CVAE
+from custom_utils.supersmpl_io import writeOBJ
+
+
+
+
+if __name__ == '__main__':
+    # init_wandb()
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+
+    smpl = SMPL(use_one_hot=config['use_one_hot'])
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+
+    folder = config['test_dir']
+
+
+    # str_split = config['name'].split('_')
+    # test_idx = str_split.index('test')
+    # checkpoint_name = '_'.join(str_split[:test_idx])
+    checkpoint_name = config['name']
+
+    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
+
+
+    output_folder = '/home/workspace/data/results/support'
+
+
+    os.makedirs(output_folder + '/bodies', exist_ok=True)
+    os.makedirs(output_folder + '/faces', exist_ok=True)
+
+    config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                            'mesh_body_posed': (14475, 3,),
+                            }
+    config['signature_static'] = { 
+                                'mesh_mask': (14475,),
+                                }
+
+    factory = DatasetFactory()
+    dataset_test = factory.get_dataset(folder=config['test_dir'], 
+                                       batch_size=config['batch_size'], 
+                                       signature=config['signature_dynamic'], 
+                                       resolutions=[config['resolution']], 
+                                       in_memory=False, 
+                                       do_repeat=False,
+                                       use_one_hot=config['use_one_hot'],
+                                       use_no_orient=config['remove_global_orientation'])
+    static_dataset_test = StaticDataset(folder=config['test_dir'], 
+                                        signature=config['signature_static'], 
+                                        resolutions=[config['resolution']], 
+                                        in_memory=False)
+   
+    for batch in tqdm.tqdm(dataset_test):
+        batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
+
+        # mask = batch['mesh_mask'].numpy().astype(bool)
+        body = batch['mesh_body_posed'].numpy()
+        # _, faces = maskFacesBatch(np.zeros((body.shape[0],14475,3)), F, mask)
+
+        for i in range(body.shape[0]):
+            id, frame_nr = batch['id'][i].numpy().decode("utf-8"), batch['frame_nr'][i].numpy().decode("utf-8")
+            writeOBJ(os.path.join(output_folder, 'bodies', id + '_' + frame_nr + '_body_no_orient.obj'), body[i], F)
+            # np.save(os.path.join(output_folder, 'faces', id + '_' + frame_nr + '.npy'), faces[i].astype(np.int32))
+
+
diff --git a/evaluate_cape.py b/evaluate_cape.py
index 0336497..8daa0af 100644
--- a/evaluate_cape.py
+++ b/evaluate_cape.py
@@ -1,5 +1,5 @@
 import os
-from configs.config_pyramid import config
+from configs.config_abl_pyr_64_128_256_cape_f import config
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import tensorflow as tf
 
@@ -44,7 +44,7 @@ from models.pyramid_model_optimized import Pyramid
 
 
 def init_wandb():
-    run = wandb.init(project='cvae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp')
+    run = wandb.init(project='results', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp')
 
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae_pyramid.py')
@@ -220,45 +220,46 @@ def test_step(batch):
     levels_input, conds = group_inputs_for_levels(batch)
     out_levels = model([levels_input, conds], training=False)
 
-    shape = batch['pose_shape'][:, 72:]
-    gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
-    shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
-    mesh_orig = batch['mesh_unposed']
+    # shape = batch['pose_shape'][:, 72:]
+    # gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
+    # shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
+    # mesh_orig = batch['mesh_unposed']
 
-    output = dict()
-    reconstr_loss_total = 0
-    kl_total = 0
+    # output = dict()
+    # reconstr_loss_total = 0
+    # kl_total = 0
 
-    #for i, res in enumerate(config['pyramid']['resolutions']):
-    i = -1
-    res = config['pyramid']['resolutions'][-1]
-    res_modifier = '_' + str(res) if res != 512 else ''
-    pred, mean, logvar = out_levels[i]
+    # #for i, res in enumerate(config['pyramid']['resolutions']):
+    # i = -1
+    # res = config['pyramid']['resolutions'][-1]
+    # res_modifier = '_' + str(res) if res != 512 else ''
+    # pred, mean, logvar = out_levels[i]
 
-    mesh_pred = uv2mesh(pred)
-    mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
-    mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
+    # mesh_pred = uv2mesh(pred)
+    # mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
+    # mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
 
-    # is the output already masked?
+    # # is the output already masked?
 
-    e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-    c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
-    c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
-    l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
-    # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
-    ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
+    # e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+    # c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
+    # c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
+    # l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
+    # # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
+    # ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
 
-    reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
+    # reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
 
-    # metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    # # metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
 
-    output = { **output, **{ 'test/edge_loss' + res_modifier: e_loss, 'test/collision_loss' + res_modifier: c_loss,
-                                'test/ms_ssim' + res_modifier: ms_ssim, 'test/l2_loss' + res_modifier: l2_loss,
-                                'test/collision_ratio' + res_modifier: c_ratio } }
+    # output = { **output, **{ 'test/edge_loss' + res_modifier: e_loss, 'test/collision_loss' + res_modifier: c_loss,
+    #                             'test/ms_ssim' + res_modifier: ms_ssim, 'test/l2_loss' + res_modifier: l2_loss,
+    #                             'test/collision_ratio' + res_modifier: c_ratio } }
 
-    loss = reconstr_loss_total + balancing_coeff * kl_total
+    # loss = reconstr_loss_total + balancing_coeff * kl_total
 
-    return { **{ 'test/loss': loss, 'test/reconstr_loss_total': reconstr_loss_total }, **output }, out_levels
+    # return { **{ 'test/loss': loss, 'test/reconstr_loss_total': reconstr_loss_total }, **output }, out_levels
+    return dict(), out_levels
 
 
 
@@ -322,13 +323,10 @@ def plot_everything(batch: dict, type: str):
 
 
 def test():
-    str_split = config['name'].split('_')
-    test_idx = str_split.index('test')
-    checkpoint_name = '_'.join(str_split[:test_idx])
-    load_model(model, 'checkpoints/' + checkpoint_name)
+    load_model(model, 'checkpoints/' + config['name'])
 
-    dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=False, do_repeat=False, is_cape=config['is_cape'])
-    static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=False, is_cape=config['is_cape'])
+    dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], resolutions=dataset_resolutions, in_memory=False, do_repeat=False, is_cape=config['is_cape'])
+    static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], resolutions=dataset_resolutions, in_memory=True, is_cape=config['is_cape'])
    
     step = 0
     losses = []
@@ -337,7 +335,7 @@ def test():
     for batch in tqdm.tqdm(dataset_test):
         batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
 
-        test_loss, out_levels = test_step(batch)
+        _, out_levels = test_step(batch)
         all_errors = []
 
         errors = []
@@ -350,15 +348,15 @@ def test():
             error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((offsets_orig - offsets_pred) ** 2, axis=-1)))
             errors.append(error)
 
-        test_loss = { **{ 'test/error': tf.reduce_mean(errors) }, **test_loss }
-        losses.append(test_loss)
+        # test_loss = { **{ 'test/error': tf.reduce_mean(errors) }, **test_loss }
+        # losses.append(test_loss)
         all_errors.extend(errors)
-        to_log = { **{ 'current_step': step, 'test/error': np.asarray(errors).mean() }, **test_loss }
-        if step % 1000 == 0:
-            figs = plot_everything(batch, 'test')
-            to_log = { **to_log, **figs }
+        # to_log = { **{ 'current_step': step, 'test/error': np.asarray(errors).mean() }, **test_loss }
+        # if step % 1000 == 0:
+        #     figs = plot_everything(batch, 'test')
+        #     to_log = { **to_log, **figs }
 
-        wandb.log(to_log, step = step)
+        # wandb.log(to_log, step = step)
         
         # print(step)
         step += 1
@@ -371,25 +369,25 @@ def test():
     print('min: ' + str(errors.min()))
     print('max: ' + str(errors.max()))
     print('median: ' + str(np.median(errors)))
-    np.save('temp/cape_f_pyr_64_128_256_error.npy', errors)
+    np.save('temp/error_'+ config['name'] +'.npy', errors)
 
-    histograms = dict()
-    for key in test_loss.keys():
-        accumulated = [losses[i][key].numpy() for i in range(len(losses))]
-        fig = ex.histogram(accumulated, title='Average: ' + str(np.asarray(accumulated).mean()))
-        fig.write_html(path_to_plotly_html, auto_play = False)
+    # histograms = dict()
+    # for key in test_loss.keys():
+    #     accumulated = [losses[i][key].numpy() for i in range(len(losses))]
+    #     fig = ex.histogram(accumulated, title='Average: ' + str(np.asarray(accumulated).mean()))
+    #     fig.write_html(path_to_plotly_html, auto_play = False)
 
-        histograms['all/' + key.split('/')[-1] + '_hist'] = wandb.Html(path_to_plotly_html)
-        # averages['all/' + key.split('/')[-1] + '_avg'] = np.asarray(accumulated).mean()
+    #     histograms['all/' + key.split('/')[-1] + '_hist'] = wandb.Html(path_to_plotly_html)
+    #     # averages['all/' + key.split('/')[-1] + '_avg'] = np.asarray(accumulated).mean()
 
-    wandb.log({ **histograms }, step = step)
-    wandb.finish()
+    # wandb.log({ **histograms }, step = step)
+    # wandb.finish()
     pass
 
 
 
 if __name__ == '__main__':
-    init_wandb()
+    # init_wandb()
     allow_memory_growth()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)
diff --git a/get_error.py b/get_error.py
new file mode 100644
index 0000000..4d85d81
--- /dev/null
+++ b/get_error.py
@@ -0,0 +1,60 @@
+import os
+import csv
+import numpy as np
+
+
+if __name__ == '__main__':
+    garment_type_table = np.load('custom_data/garment_type_table_test.npy', allow_pickle=True).item()
+    folder = '/home/workspace/data/results/errors'
+
+    checkpoint_name = 'abl_all'
+    iteration = '70000'
+
+    errors = dict()
+    errors_skirt = dict()
+    per_garment_errors = { 'Tshirt': [], 'Dress': [], 'Skirt': [], 'Top': [] }
+    per_garment_errors_skirt = { 'Tshirt': [], 'Jumpsuit': [], 'Trousers': [], 'Top': [] }
+
+
+    with open(os.path.join(folder, checkpoint_name + '-' + iteration + '.csv'), 'r') as f:
+        csv_reader = csv.reader(f, delimiter=',')
+        for row in csv_reader:
+            per_garment_errors[garment_type_table[row[0].split('_')[0]]].append(row[1])
+
+    with open(os.path.join(folder, checkpoint_name + '_skirt' + '-' + iteration + '.csv'), 'r') as f:
+        csv_reader = csv.reader(os.path.join(folder, checkpoint_name + '_skirt' + '-' + iteration + '.csv'), delimiter=',')
+        for row in csv_reader:
+            per_garment_errors_skirt[garment_type_table[row[0]]].append(row[1])
+
+    print('Non-skirt:')
+    for key in per_garment_errors.keys():
+        avg = np.mean(per_garment_errors[key])
+        errors.append(avg)
+        print(key + ': ' + str(avg))
+    print('Skirt:')
+    for key in per_garment_errors_skirt.keys():
+        avg = np.mean(per_garment_errors_skirt[key])
+        errors_skirt.append(avg)
+        print(key + ': ' + str(avg))
+
+    print('Avg non-skirt: ' + str(np.mean(errors)))
+    print('Avg skirt: ' + str(np.mean(errors_skirt)))
+
+    cum_avg_shirt = np.mean(per_garment_errors['Tshirt'] + per_garment_errors_skirt['Tshirt'])
+    cum_avg_top = np.mean(per_garment_errors['Top'] + per_garment_errors_skirt['Top'])
+
+    print('Global tshirt avg: ' + str(cum_avg_shirt))
+    print('Global tshirt avg: ' + str(cum_avg_top))
+
+    print('Globl per garment: ')
+
+    global_errors = []
+    per_garment_errors['Tshirt'].extend(per_garment_errors_skirt['Tshirt'])
+    per_garment_errors['Top'].extend(per_garment_errors_skirt['Top'])
+    for key in per_garment_errors.keys():
+        avg = np.mean(per_garment_errors[key])
+        global_errors.append(avg)
+        print(key + ': ' + str(avg))
+
+    print('Global avg: ' + str(np.mean(global_errors)))
+    
\ No newline at end of file
diff --git a/losses.py b/losses.py
index f66fefb..8575648 100644
--- a/losses.py
+++ b/losses.py
@@ -25,107 +25,14 @@ from scipy.spatial import cKDTree
 from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals
 from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
 
-from custom_utils.custom_utils import maskFacesBatch, _maskFaces
-from custom_utils.graph_utils import sparse_to_tensor
-
-
-def d_logistic(real_images, generator, discriminator, z_dim, conditions):
-    batch_size = tf.shape(real_images)[0]
-    z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
-
-    # forward pass
-    fake_images = generator([z, conditions], training=True)
-    real_scores = discriminator(real_images, training=True)
-    fake_scores = discriminator(fake_images, training=True)
-
-    # gan loss
-    d_loss = tf.math.softplus(fake_scores)
-    d_loss += tf.math.softplus(-real_scores)
-
-
-    return d_loss
-
-
-def d_logistic_r1_reg(real_images, generator, discriminator, z_dim, conditions):
-    batch_size = tf.shape(real_images)[0]
-    z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
-
-    # forward pass
-    fake_images = generator([z, conditions], training=True)
-    real_scores = discriminator(real_images, training=True)
-    fake_scores = discriminator(fake_images, training=True)
-
-    # gan loss
-    d_loss = tf.math.softplus(fake_scores)
-    d_loss = d_loss + tf.math.softplus(-real_scores)
-
-    # gradient penalty
-    with tf.GradientTape() as r1_tape:
-        r1_tape.watch([real_images, conditions])
-        real_scores = discriminator(real_images, training=True)
-        real_loss = tf.reduce_sum(real_scores)
-
-    real_grads = r1_tape.gradient(real_loss, real_images)
-    r1_penalty = tf.reduce_sum(tf.math.square(real_grads), axis=[1, 2, 3])
-    r1_penalty = tf.expand_dims(r1_penalty, axis=1)
-    return d_loss, r1_penalty
-
-
-def g_logistic_non_saturating(real_images, generator, discriminator, z_dim, conditions):
-    batch_size = tf.shape(real_images)[0]
-    z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
-
-    # forward pass
-    fake_images = generator([z, conditions], training=True)
-    fake_scores = discriminator(fake_images, training=True)
-
-    # gan loss
-    g_loss = tf.math.softplus(-fake_scores)
-    return fake_images, g_loss
-
-
-def g_logistic_ns_pathreg(real_images, generator, discriminator, z_dim,
-                          pl_mean, pl_minibatch_shrink, pl_denorm, pl_decay,
-                          conditions):
-    batch_size = tf.shape(real_images)[0]
-    z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
-
-    pl_minibatch = tf.maximum(1, tf.math.floordiv(batch_size, pl_minibatch_shrink))
-    pl_z = tf.random.normal(shape=[pl_minibatch, z_dim], dtype=tf.float32)
-    pl_conditions = { key: conditions[key][:pl_minibatch] for key in conditions.keys() }
-    # for key in conditions.keys():
-    #     pl_conditions[key] = conditions[key][:pl_minibatch]
-    # pl_conditions = (conditions[0][:pl_minibatch], conditions[1][:pl_minibatch], conditions[2][:pl_minibatch], conditions[3][:pl_minibatch], conditions[4][:pl_minibatch])
-
-    # forward pass
-    fake_images = generator([z, conditions], ret_w_broadcasted=False, training=True)
-    fake_scores = discriminator(fake_images, training=True)
-    g_loss = tf.math.softplus(-fake_scores)
-
-    # Evaluate the regularization term using a smaller minibatch to conserve memory.
-    with tf.GradientTape() as pl_tape:
-        pl_tape.watch([pl_z, pl_conditions])
-        pl_fake_images, pl_w_broadcasted = generator([pl_z, pl_conditions], ret_w_broadcasted=True, training=True)
-
-        pl_noise = tf.random.normal(tf.shape(pl_fake_images)) * pl_denorm
-        pl_noise_applied = tf.reduce_sum(pl_fake_images * pl_noise)
-
-    pl_grads = pl_tape.gradient(pl_noise_applied, pl_w_broadcasted)
-    pl_lengths = tf.math.sqrt(tf.reduce_mean(tf.reduce_sum(tf.math.square(pl_grads), axis=2), axis=1))
-
-    # Track exponential moving average of |J*y|.
-    pl_mean_val = pl_mean + pl_decay * (tf.reduce_mean(pl_lengths) - pl_mean)
-    pl_mean.assign(pl_mean_val)
-
-    # Calculate (|J*y|-a)^2.
-    pl_penalty = tf.square(pl_lengths - pl_mean)
-    return fake_images, g_loss, pl_penalty
 
 
 # from hbertiche/PBNS
 """ COLLISION FUNCTIONS START """
-F = np.load("custom_utils/smpl_faces.npy")
-
+F_coll = np.load("custom_utils/smpl_faces.npy")
+F_coll_skirt = np.load("custom_utils/smpl_faces_skirt.npy")
+F_norm = np.tile(F_coll[np.newaxis,:,:], [4,1,1])
+F_norm_skirt = np.tile(F_coll_skirt[np.newaxis,:,:], [4,1,1])
 
 def _nearest_neighbour(V, B):
 	tree = cKDTree(B.numpy())
@@ -137,7 +44,7 @@ def _tf_nn_parallel(V, B):
 
 # @tf.function
 def collision_loss(V, B, mask, thr=.004):
-    B_vn = vertex_normals(B, tf.tile(F[None], [B.shape[0], 1, 1]))
+    B_vn = vertex_normals(B, tf.tile(F_coll[None], [B.shape[0], 1, 1]))
    
     # compute correspondences
     idx = _tf_nn_parallel(V, B)
@@ -149,9 +56,11 @@ def collision_loss(V, B, mask, thr=.004):
     D = V - tf.gather_nd(B, idx)
     b_vn = tf.gather_nd(B_vn, idx)
     dot = tf.einsum('abc,abc->ab', D, b_vn)
-    loss = tf.reduce_sum((tf.minimum(dot - thr, 0)  * mask) ** 2, axis=-1) # / tf.reduce_sum(mask, axis=-1)
+    thresholded = tf.minimum(dot - thr, 0)  * mask
+    loss = tf.reduce_sum(thresholded ** 2, axis=-1)
+    ratio = tf.reduce_sum(tf.cast(thresholded < 0, dtype=tf.float32)) / tf.reduce_sum(mask)
 
-    return loss
+    return loss, ratio
 
 
 def collision_ratio(V, B, mask, thr=.004):
@@ -327,6 +236,11 @@ def kl_loss(mean, logvar):
     return kl_loss
 
 
+def normal_loss(pred, gt, mask, is_skirt=False):
+    pred_norm = vertex_normals(pred, F_norm_skirt if is_skirt else F_norm)
+    gt_norm = vertex_normals(gt, F_norm_skirt if is_skirt else F_norm)
+    return tf.reduce_mean(tf.reduce_mean((tf.reduce_sum(tf.math.abs(pred_norm - gt_norm), axis=-1) * mask), axis=-1))
+
 # vgg = tf.keras.applications.vgg16.VGG16(include_top=False, input_shape=(512,256,3))
 # vgg_block_1 = tf.keras.Sequential(vgg.layers[:4])
 # vgg_block_2 = tf.keras.Sequential(vgg.layers[4:7])
diff --git a/models/cvae_model.py b/models/cvae_model.py
index 7770d5b..7b8424e 100644
--- a/models/cvae_model.py
+++ b/models/cvae_model.py
@@ -26,14 +26,16 @@ class CVAE(tf.keras.Model):
         if config['body_model'] != '':
             load_model(self.body_enc, config['body_model'] + '_enc')
 
-        self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
-        if config['curvature_model'] != '':
-            load_model(self.curv_enc, config['curvature_model'] + '_enc')
+        if config['use_curvature']:
+            self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
+            if config['curvature_model'] != '':
+                load_model(self.curv_enc, config['curvature_model'] + '_enc')
 
         if config['freeze_encoders']:
             self.static_enc.trainable = False
             self.body_enc.trainable = False
-            self.curv_enc.trainable = False
+            if config['use_curvature']:
+                self.curv_enc.trainable = False
 
         self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
                                                  tf.keras.layers.Dense(256, activation='relu'),
@@ -50,7 +52,13 @@ class CVAE(tf.keras.Model):
 
         self.add_noise = config['cvae']['add_noise']
         self.noise_amount = config['cvae']['noise_amount']
+        self.use_curvature = config['use_curvature']
+        self.use_offsets_over_body = not config['use_skinning']
 
+        if self.use_offsets_over_body:
+            assert self.avg_garment is None, 'When we are not using skinning, we can\'t use avg_garment!'
+        if self.avg_garment is not None:
+            assert not self.use_offsets_over_body, 'When avg_garment is given, we must use skinning!'
 
     @tf.function
     def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None):
@@ -84,27 +92,24 @@ class CVAE(tf.keras.Model):
     def encode_conds(self, uv_template, uv_body_posed, uv_curv, training=None):
         z_body = self.body_enc(uv_body_posed, training=training)
         z_static = self.static_enc(uv_template, training=training)
-        z_curv = self.curv_enc(uv_curv, training=training)
+        if self.use_curvature:
+            z_curv = self.curv_enc(uv_curv, training=training)
+            z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
+        else:
+            z_cond = tf.concat([z_static, z_body], axis=-1)
 
-        z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
         z_cond = self.cond_encoder(z_cond, training=training)
 
         return z_cond
 
 
     def call(self, inputs, training=None):
-
-        if self.size == 512:
-            uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
-            uv_template_high, uv_body_posed_high, uv_curv_high = uv_template, uv_body_posed, uv_curv
-        else:
-            uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_template_high, uv_body_posed_high, uv_curv_high, uv_mask_high = inputs
-
+        uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
 
         x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
         mean, logvar = self.encoder(x, training=training)
         z = self.reparameterize(mean, logvar)
-        z_cond = self.encode_conds(uv_template_high, uv_body_posed_high, uv_curv_high, training=training)
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=training)
 
         if self.add_noise and training == True:
             z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
@@ -115,7 +120,10 @@ class CVAE(tf.keras.Model):
 
         if self.avg_garment is not None:
             x = x + self.avg_garment[tf.newaxis,:,:,:]
- 
+    
+        if self.use_offsets_over_body:
+            x = x + uv_body_posed
+
         x = x * uv_mask[:,:,:,tf.newaxis]
 
         return x, mean, logvar
diff --git a/models/pyramid_model_optimized.py b/models/pyramid_model_optimized.py
index 0bec6a4..51330e0 100644
--- a/models/pyramid_model_optimized.py
+++ b/models/pyramid_model_optimized.py
@@ -42,6 +42,28 @@ class Pyramid(tf.keras.Model):
 
         self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae': config['cond_ae']} }) for i in range(len(config['pyramid']['levels']))]
 
+
+        # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+        # self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+        # self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+
+        if len(config['pyramid']['resolutions']) == 2:
+            self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10)])
+        if len(config['pyramid']['resolutions']) == 3:
+            self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+            self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10)])
+        if len(config['pyramid']['resolutions']) == 4:
+            self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
+            self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+            self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10)])
+        if len(config['pyramid']['resolutions']) == 5:
+            self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
+            self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
+            self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+            self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10)])
+
+        self.add_noise = config['pyramid']['add_noise']
+        self.use_offsets = config['use_offsets']
  
         self.upscale_enhancer = None
         if config['pyramid']['final_upscale_enhancer'] != '':
@@ -66,42 +88,64 @@ class Pyramid(tf.keras.Model):
     def call(self, inputs, training=None):
         levels_input, conds = inputs
         uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
+
         z_cond = self.encode_conds(uv_static, uv_body, uv_curv)
 
-        # out_levels = []
-        # pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
-        # if 1 != len(self.levels):
-        #     pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training)  # levels_input[i][-2]=uv_mask_res
-        #     pred = tf.image.resize(pred, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
-        #     out_levels.append([pred, mean, logvar])
-        #     pred = pred * levels_input[1][-2][:,:,:,tf.newaxis]
-        # else:
-        #     pred = self.enhance(pred, levels_input[0][-2], training=training)
-        #     pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
-        #     pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
-        #     out_levels.append([pred, mean, logvar])
+        if self.add_noise:
+            z_cond = z_cond + tf.random.uniform(z_cond.shape, maxval=10)
 
-        # for i in range(1, len(self.levels)):
-        #     pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond, training=training)
-        #     pred = pred + pred_level
-        #     if i != len(self.levels) - 1:
-        #         pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
-        #         out_levels.append([pred, mean, logvar])
-        #         pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.LANCZOS5)
-        #         pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
-        #     else:
-        #         if self.upscale_enhancer is not None:
-        #             pred = self.enhance(pred, levels_input[-1][-2], training=training)
-        #             pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
-        #         pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
-        #         out_levels.append([pred, mean, logvar])
+        out_levels = []
+        pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
+        if self.use_offsets:
+            pred = pred + levels_input[0][2] # uv_body_posed_32
+        if 1 != len(self.levels):
+            pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training)  # levels_input[i][-2]=uv_mask_res
+            pred = tf.image.resize(pred, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
+            out_levels.append([pred, mean, logvar])
+            pred = pred * levels_input[1][-2][:,:,:,tf.newaxis]
+        else:
+            pred = self.enhance(pred, levels_input[0][-2], training=training)
+            pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+            pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
+            out_levels.append([pred, mean, logvar])
 
-        pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
-        pred_0 = self.enhance(pred_0, levels_input[0][-2], training=training)
-        pred_0 = tf.image.resize(pred_0, (512, 256), tf.image.ResizeMethod.LANCZOS5)
-        pred_0 = pred_0 * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
+        for i in range(1, len(self.levels)):
+            pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond, training=training)
+            pred = pred + pred_level
+            if i != len(self.levels) - 1:
+                pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
+                out_levels.append([pred, mean, logvar])
+                pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.LANCZOS5)
+                pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
+            else:
+                if self.upscale_enhancer is not None:
+                    pred = self.enhance(pred, levels_input[-1][-2], training=training)
+                    pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+                pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
+                out_levels.append([pred, mean, logvar])
 
-        return [ [pred_0, mean_0, logvar_0] ]
+        return out_levels
+
+        # # 32
+        # pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
+        # pred_0 = self.enhance(pred_0, levels_input[0][-2], training=training)
+        # pred_0 = tf.image.resize(pred_0, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+        # pred_0 = pred_0 * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
+        # return [ [pred_0, mean_0, logvar_0] ]
+
+        # # 32 64
+        # pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
+        # pred_0 = self.levels[0].enhance(pred_0, levels_input[0][-2], training=training)  # levels_input[i][-2]=uv_mask_res
+        # pred_0 = tf.image.resize(pred_0, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
+        # pred = pred_0 * levels_input[1][-2][:,:,:,tf.newaxis]
+
+        # pred_1, mean_1, logvar_1 = self.levels[1](levels_input[1], z_cond, training=training)
+        # pred_1 = pred + pred_1
+        # pred_1 = self.enhance(pred_1, levels_input[-1][-2], training=training)
+        # pred_1 = tf.image.resize(pred_1, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+        # pred_1 = pred_1 * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
+
+        # return [ [pred_0, mean_0, logvar_0], [pred_1, mean_1, logvar_1] ]
 
 
 # class Pyramid(tf.keras.Model):
diff --git a/verify_cape_skinning.py b/verify_cape_skinning.py
index d6773b8..0715f8c 100644
--- a/verify_cape_skinning.py
+++ b/verify_cape_skinning.py
@@ -38,7 +38,7 @@ if __name__ == '__main__':
     F = np.load("custom_utils/smpl_faces_orig.npy")
 
     factory = DatasetFactory()
-    signature_dynamic = { 'id': (),
+    signature_dynamic = { 'id': (), 'frame_nr': (),
                           'pose_shape': (72 + 11,),
                         #   'uv_posed': (3, 512, 256,),
                         #   'uv_body_posed': (3, 512, 256,),
diff --git a/verify_poses.py b/verify_poses.py
index 865b233..90c4ad7 100644
--- a/verify_poses.py
+++ b/verify_poses.py
@@ -1,10 +1,27 @@
 import os
+
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = '5'
+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 sys
 import time
 import cv2
 import random
 import numpy as np
-import tensorflow as tf
 import cmapy as cm
 import plotly.graph_objects as go
 import plotly.express as ex
@@ -14,7 +31,7 @@ from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
 from custom_utils.smpl import SMPL
 
 if __name__ == '__main__':
-    os.environ["CUDA_VISIBLE_DEVICES"] = "6"
+    os.environ["CUDA_VISIBLE_DEVICES"] = "4"
 
     uv2mesh = UV2Mesh()
     smpl = SMPL()
@@ -22,7 +39,7 @@ if __name__ == '__main__':
     F = np.load("custom_utils/smpl_faces.npy")
 
     factory = DatasetFactory()
-    signature_dynamic = { 'id': (),
+    signature_dynamic = { 'id': (), 'frame_nr': (),
                           'pose_shape': (72 + 11,),
                         #   'uv_posed': (3, 512, 256,),
                         #   'uv_body_posed': (3, 512, 256,),
@@ -41,11 +58,13 @@ if __name__ == '__main__':
                         #  'uv_shaped_body': (3, 512, 256,),
                           }
 
-    dataset = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_test/train',
+    dataset = factory.get_dataset(folder='/home/workspace/local-data/train',
                                   batch_size=1,
-                                  signature=signature_dynamic)
-    static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_test/train',
-                                   signature=signature_static)
+                                  signature=signature_dynamic,
+                                  resolutions=[512])
+    static_dataset = StaticDataset(folder='/home/workspace/local-data/train',
+                                   signature=signature_static,
+                                   resolutions=[512])
 
 
     uv_res = 256
@@ -59,102 +78,82 @@ if __name__ == '__main__':
         mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
         mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
         pose, shape = batch['pose_shape'][:, :72].numpy(), batch['pose_shape'][:, 72:].numpy()
-        mesh_reposed = smpl.run_garment(pose, shape, batch['mesh_unposed']).numpy()[0]
+        gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
+        shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]).numpy()
+        mesh_unposed_shaped = batch['mesh_unposed'] + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
+        mesh_reposed = smpl.run_garment(pose, shape, mesh_unposed_shaped).numpy()[0]
         mesh_posed = np.squeeze(batch['mesh_posed'].numpy())
         mesh_body_posed = np.squeeze(batch['mesh_body_posed'].numpy())
-        mesh_posed_body_no_orient = np.squeeze(batch['mesh_posed_body_no_orient'].numpy())
 
-        # x = mesh_posed[mask, :]
-        R = smpl.run_rodrigues(pose)
-        # shape[:, -1] = 1
-        x = smpl.run(R, shape)[0][0]
+        # x = smpl.run(smpl.run_rodrigues(pose), shape)[0][0]
+        x = mesh_reposed[mask, :]
         y = mesh_body_posed
-        z = mesh_posed[mask, :]
         _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
         go.Figure(data=[
                 go.Mesh3d(
-                    # 8 vertices of a cube
                     x=x[:,0],
                     y=x[:,1],
                     z=x[:,2],
-                    # i, j and k give the vertices of triangles
-                    i = F[:,0],
-                    j = F[:,1],
-                    k = F[:,2]
+                    i = faces[:,0],
+                    j = faces[:,1],
+                    k = faces[:,2]
                 ),
                 go.Mesh3d(
-                    # 8 vertices of a cube
                     x=y[:,0],
                     y=y[:,1],
                     z=y[:,2],
-                    # i, j and k give the vertices of triangles
                     i = F[:,0],
                     j = F[:,1],
                     k = F[:,2]
-                ),
-                # go.Mesh3d(
-                #     # 8 vertices of a cube
-                #     x=z[:,0],
-                #     y=z[:,1],
-                #     z=z[:,2],
-                #     # i, j and k give the vertices of triangles
-                #     i = faces[:,0],
-                #     j = faces[:,1],
-                #     k = faces[:,2]
-                # ),
+                )
                 ]).show()
 
-        pose_2 = pose.copy()
-        pose_2[:, :3] = 0
-        R = smpl.run_rodrigues(pose_2)
-        x = smpl.run(R, shape)[0][0]
-        y = mesh_posed_body_no_orient
-        _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
-        go.Figure(data=[
-                go.Mesh3d(
-                    # 8 vertices of a cube
-                    x=x[:,0],
-                    y=x[:,1],
-                    z=x[:,2],
-                    # i, j and k give the vertices of triangles
-                    i = F[:,0],
-                    j = F[:,1],
-                    k = F[:,2]
-                ),
-                go.Mesh3d(
-                    # 8 vertices of a cube
-                    x=y[:,0],
-                    y=y[:,1],
-                    z=y[:,2],
-                    # i, j and k give the vertices of triangles
-                    i = F[:,0],
-                    j = F[:,1],
-                    k = F[:,2]
-                )]).show()
+        # pose_2 = pose.copy()
+        # pose_2[:, :3] = 0
+        # # x = smpl.run(smpl.run_rodrigues(pose_2), shape)[0][0]
+        # x = mesh_unposed
+        # y = mesh_posed_body_no_orient
+        # _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
+        # go.Figure(data=[
+        #         go.Mesh3d(
+        #             # 8 vertices of a cube
+        #             x=x[:,0],
+        #             y=x[:,1],
+        #             z=x[:,2],
+        #             # i, j and k give the vertices of triangles
+        #             i = F[:,0],
+        #             j = F[:,1],
+        #             k = F[:,2]
+        #         ),
+        #         go.Mesh3d(
+        #             # 8 vertices of a cube
+        #             x=y[:,0],
+        #             y=y[:,1],
+        #             z=y[:,2],
+        #             i = F[:,0],
+        #             j = F[:,1],
+        #             k = F[:,2]
+        #         )]).show()
 
-        x = mesh_posed[mask, :]
-        y = mesh_reposed[mask, :]
-        _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
-        go.Figure(data=[
-                go.Mesh3d(
-                    # 8 vertices of a cube
-                    x=x[:,0],
-                    y=x[:,1],
-                    z=x[:,2],
-                    # 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=y[:,0],
-                    y=y[:,1],
-                    z=y[:,2],
-                    # i, j and k give the vertices of triangles
-                    i = faces[:,0],
-                    j = faces[:,1],
-                    k = faces[:,2]
-                )]).show()
+        # x = mesh_posed[mask, :]
+        # y = mesh_reposed[mask, :]
+        # _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
+        # go.Figure(data=[
+        #         go.Mesh3d(
+        #             x=x[:,0],
+        #             y=x[:,1],
+        #             z=x[:,2],
+        #             i = faces[:,0],
+        #             j = faces[:,1],
+        #             k = faces[:,2]
+        #         ),
+        #         go.Mesh3d(
+        #             x=y[:,0],
+        #             y=y[:,1],
+        #             z=y[:,2],
+        #             i = faces[:,0],
+        #             j = faces[:,1],
+        #             k = faces[:,2]
+        #         )]).show()
 
         print()
\ No newline at end of file