update docker

configs/ae_body_cape_m.py

@@ -2,7 +2,7 @@ from utils.convnext_utils import CNBlockConfig
 
 config = dict()
 
-config['gpu'] = '4'
+config['gpu'] = '0'
 config['name'] = 'body_torch_cape_m'
 config['notes'] = ' ' 
 config['tags'] = ['torch', 'body'] 
@@ -11,9 +11,9 @@ config['script'] = 'scripts/ae_body_train.py'
 config['pretrained_model'] = ''
 
 config['batch_size'] = 4
-config['train_dir'] = '/data-local/127/cape_processed/train/male'
-config['valid_dir'] = '/data-local/127/cape_processed/test/male'
-config['test_dir'] = '/data-local/127/cape_processed/test/male'
+config['train_dir'] = '/data-local/128/cape_processed/train/male'
+config['valid_dir'] = '/data-local/128/cape_processed/test/male'
+config['test_dir'] = '/data-local/128/cape_processed/test/male'
 config['train_frames'] = None
 config['valid_frames'] = None
 

configs/iccvw/ae_body_1.py

@@ -0,0 +1,78 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '2'
+config['name'] = 'body_iccvw_1'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'body'] 
+config['script'] = 'scripts/ae_body_train.py'
+
+config['pretrained_model'] = ''
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/cloth3d_processed'
+config['valid_dir'] = '/data-local/127/cloth3d_processed'
+config['test_dir'] = '/data-net/dataset/test_ready'
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+if config['is_skirt']:
+  config['train_frames'] = 'frames_train_skirt.txt'
+  config['valid_frames'] = 'frames_valid_skirt.txt'
+else:
+  config['train_frames'] = 'frames_train.txt'
+  config['valid_frames'] = 'frames_valid.txt'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 16 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.0001
+config['cond_ae']['upscale_output'] = False
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'pose_shape',
+                        'uv_body_posed',
+                        'mesh_mask',
+                        'mesh_body_posed',
+                        }

configs/iccvw/ae_body_cape_f.py

@@ -0,0 +1,73 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '0'
+config['name'] = 'body_iccvw_cape_f'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'body'] 
+config['script'] = 'scripts/ae_body_train.py'
+
+config['pretrained_model'] = ''
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/128/cape_processed/train/female'
+config['valid_dir'] = '/data-local/128/cape_processed/test/female'
+config['test_dir'] = '/data-local/128/cape_processed/test/female'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = True
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 16 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.0001
+config['cond_ae']['upscale_output'] = False
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'pose_shape',
+                        'uv_body_posed',
+                        'mesh_mask',
+                        'mesh_body_posed',
+                        }

configs/iccvw/ae_body_cape_m.py

@@ -0,0 +1,73 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '0'
+config['name'] = 'body_torch_cape_m'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'body'] 
+config['script'] = 'scripts/ae_body_train.py'
+
+config['pretrained_model'] = ''
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/128/cape_processed/train/male'
+config['valid_dir'] = '/data-local/128/cape_processed/test/male'
+config['test_dir'] = '/data-local/128/cape_processed/test/male'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = True
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 16 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.0001
+config['cond_ae']['upscale_output'] = False
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'pose_shape',
+                        'uv_body_posed',
+                        'mesh_mask',
+                        'mesh_body_posed',
+                        }

configs/iccvw/ae_norm_1.py

@@ -0,0 +1,78 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '1'
+config['name'] = 'norm_iccvw_1'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'norm'] 
+config['script'] = 'scripts/ae_norm_train.py'
+
+config['pretrained_model'] = ''
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/cloth3d_processed'
+config['valid_dir'] = '/data-local/127/cloth3d_processed'
+config['test_dir'] = '/data-net/dataset/test_ready'
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+if config['is_skirt']:
+  config['train_frames'] = 'frames_train_skirt.txt'
+  config['valid_frames'] = 'frames_valid_skirt.txt'
+else:
+  config['train_frames'] = 'frames_train.txt'
+  config['valid_frames'] = 'frames_valid.txt'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 40 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.0001
+config['cond_ae']['upscale_output'] = False
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'uv_mask',
+                        'uv_normals',
+                        'mesh_mask',
+                        'mesh_normals',
+                        }

configs/iccvw/ae_static_1.py

@@ -0,0 +1,78 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '3'
+config['name'] = 'static_iccvw_2'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'static'] 
+config['script'] = 'scripts/ae_static_train.py'
+
+config['pretrained_model'] = 'checkpoints/iccvw/static_iccvw_1'
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/cloth3d_processed'
+config['valid_dir'] = '/data-local/127/cloth3d_processed'
+config['test_dir'] = '/data-net/dataset/test_ready'
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = True
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+if config['is_skirt']:
+  config['train_frames'] = 'frames_train_skirt.txt'
+  config['valid_frames'] = 'frames_valid_skirt.txt'
+else:
+  config['train_frames'] = 'frames_train.txt'
+  config['valid_frames'] = 'frames_valid.txt'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 8 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(8, 16, 3),
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.0005
+config['cond_ae']['upscale_output'] = False
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'uv_static',
+                        'mesh_static',
+                        'mesh_mask',
+                        'uv_mask',
+                        }

configs/iccvw/ae_static_cape_f.py

@@ -0,0 +1,74 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '4'
+config['name'] = 'static_iccvw_cape_f'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'static'] 
+config['script'] = 'scripts/ae_static_train.py'
+
+config['pretrained_model'] = ''
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127_2/cape_processed/train/female'
+config['valid_dir'] = '/data-local/127_2/cape_processed/test/female'
+config['test_dir'] = '/data-local/127_2/cape_processed/test/female'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = True
+config['is_cape'] = True
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 8 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(8, 16, 3),
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.0005
+config['cond_ae']['upscale_output'] = False
+config['pretrained_model'] = ''
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'uv_static',
+                        'mesh_static',
+                        'mesh_mask',
+                        'uv_mask',
+                        }

configs/iccvw/ae_static_cape_m.py

@@ -0,0 +1,74 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '5'
+config['name'] = 'static_iccvw_cape_m'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'static'] 
+config['script'] = 'scripts/ae_static_train.py'
+
+config['pretrained_model'] = ''
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127_2/cape_processed/train/male'
+config['valid_dir'] = '/data-local/127_2/cape_processed/test/male'
+config['test_dir'] = '/data-local/127_2/cape_processed/test/male'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = True
+config['is_cape'] = True
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 8 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(8, 16, 3),
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.0005
+config['cond_ae']['upscale_output'] = False
+config['pretrained_model'] = ''
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'uv_static',
+                        'mesh_static',
+                        'mesh_mask',
+                        'uv_mask',
+                        }

data/dataset.py

@@ -1,7 +1,7 @@
 import os
 import sys
 
-# os.environ["CUDA_VISIBLE_DEVICES"] = '6'
+# os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'
 
 import math
 import time
@@ -14,7 +14,6 @@ import scipy.io as sio
 import plotly.express as ex
 import torchvision
 
-from torch.utils.data import Dataset
 from torchvision.transforms import Resize
 from torchvision.transforms.functional import InterpolationMode
 
@@ -22,6 +21,7 @@ sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 
 from data.utils import Frame
 from data.cloth3d_io import readPC2, readPC2Frame
+from data.dataset_cape import DatasetCape
 
 
 def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=None):
@@ -131,7 +131,7 @@ def read_frames_file(frames_file, shuffle=False):
     return frames
 
 
-class Dataset(Dataset):
+class DatasetCloth3D(torch.utils.data.Dataset):
     def __init__(self, folder, signature, resolutions=[512], frames_file='', device='cuda', use_cache=False, is_cape=False):
         self.folder = folder
         self.device = device
@@ -229,6 +229,14 @@ class Dataset(Dataset):
         data = { **data, **self.static_cache[frame] }
         return data
 
+def Dataset(folder, signature, resolutions=[512], frames_file='', device='cuda', use_cache=False, is_cape=False):
+    if not is_cape:
+        return DatasetCloth3D(folder, signature, resolutions=resolutions, frames_file=frames_file, device=device, use_cache=use_cache, is_cape=False)
+    else:
+        return DatasetCape(folder, signature, resolutions=resolutions, frames_file=frames_file, device=device, use_cache=use_cache, is_cape=True)
+
+
+
 
 if __name__ == '__main__':
     from torch.utils.data import DataLoader

data/dataset_cape.py

@@ -0,0 +1,306 @@
+import os
+import sys
+
+# os.environ["CUDA_VISIBLE_DEVICES"] = '6'
+
+import math
+import time
+import torch
+import random
+import pytorch3d
+import torchvision as tv
+import numpy as np
+import scipy.io as sio
+import plotly.express as ex
+import torchvision
+
+from torchvision.transforms import Resize
+from torchvision.transforms.functional import InterpolationMode
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+from data.utils import Frame
+
+
+def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=None):
+    if frames_file is None:
+        uv_static_folder = os.path.join(folder, 'pose_shape') 
+        frames = sorted(os.listdir(uv_static_folder))
+
+        if is_cape:
+            frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0]
+            frames = [Frame(frame.split('.')[0], frame.split('.')[1]) for frame in frames]
+        else:
+            frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
+            frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames]
+
+        # if testing:
+        #     frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') ))
+        #     frames = [frame.split('.')[0] + '_4' for frame in frames]
+
+
+    if frames_file is not None:
+        frames = []
+        with open(os.path.join(folder, frames_file), 'r') as f:
+            lines = f.readlines()
+            lines = [line.strip() for line in lines]
+        frames = [Frame(line.split('_')[0], line.split('_')[1]) for line in lines]
+
+    if shuffle:
+        random.shuffle(frames)
+
+    return frames
+
+
+class StaticCache():
+    def __init__(self, folder, signature, resolutions=[512], is_cape=False, device='cuda') -> None:
+        self.store = dict()
+        self.signature = signature
+        self.folder = folder
+        assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
+        self.resolutions = resolutions
+        self.filled = ''
+        self.is_cape = is_cape
+        self.device = device
+        
+        if is_cape:
+            self.uv_masks = dict()
+            uv_mask = torch.tensor(np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32), device=device)
+            self.uv_masks[512] = uv_mask.to(torch.float32).cpu().numpy()
+            for res in [32, 64, 128, 256]:
+                down = Resize(size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST_EXACT)(uv_mask[None,:,:])
+                down = torch.squeeze(down).to(torch.float32)
+                self.uv_masks[res] = down.cpu().numpy()
+
+
+    def read(self, id):
+        data = dict()
+
+        for res in self.resolutions:
+            res_string = ''
+            if res != 512:
+                res_string = '_' + str(res)
+
+            if 'uv_static' + res_string in self.signature:
+                uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id +  '.npz')); uv_static = uv_static['uv_static' + self.filled]; data['uv_static' + res_string] = uv_static
+            if 'uv_mask' + res_string in self.signature:
+                if self.is_cape:
+                    data['uv_mask' + res_string] = self.uv_masks[res]
+                else:
+                    uv_mask = np.load(os.path.join(self.folder, 'uv_mask' + res_string, id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask
+
+        if 'mesh_static' in self.signature:
+            mesh_static = np.load(os.path.join(self.folder, 'mesh_static', id +  '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static.astype(np.float32)
+        if 'mesh_mask' in self.signature:
+            if self.is_cape:
+                data['mesh_mask'] = np.ones((3638), np.float32)
+            else:
+                mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask
+        # if 'uv_shaped_body' in self.signature:
+        #     uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; data['uv_shaped_body'] = uv_shaped_body
+
+        return data
+
+
+    def __getitem__(self, id) -> dict:
+        if id not in self.store:
+            item = self.read(id)
+            stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()}
+            for key, value in stored_data.items():
+                if key.startswith('uv_') and value.shape[-1] == 3:
+                    stored_data[key] = torch.permute(value, [2,0,1])
+            self.store[id] = stored_data
+        
+        data = { key: value.to(self.device) for key, value in self.store[id].items()}
+        return data
+
+
+class DatasetCape(torch.utils.data.Dataset):
+    def __init__(self, folder, signature, resolutions=[512], frames_file='frames_all.txt', device='cuda', use_cache=False, is_cape=False):
+        self.folder = folder
+        self.device = device
+        self.is_cape = is_cape
+        self.signature = signature
+        self.resolutions = resolutions
+
+        self.transforms = tv.transforms.ToTensor()
+        self.frames = read_frames(folder, frames_file=frames_file, is_cape=is_cape, shuffle=True)
+
+        self.body_faces = np.load('smpl/smpl_faces.npy')
+        self.static_cache = StaticCache(folder, signature, is_cape=is_cape, resolutions=resolutions, device=device)
+
+        self.cache = dict() if use_cache else None
+
+
+    def __len__(self):
+        return len(self.frames)
+
+
+    def __getitem__(self, idx):
+        frame : Frame = self.frames[idx]
+        if self.cache is not None:
+            if frame in self.cache.keys():
+                data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in self.cache[frame].items()}
+                data = { **data, **self.static_cache[frame.id] }
+                return data
+
+        separator = '_'
+        if self.is_cape:
+            separator = '.'
+
+        uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient_2'
+        if self.is_cape:
+            uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient'
+
+        item = dict()
+
+        if 'id' in self.signature:
+            item['id'] = frame.id
+        if 'frame_nr' in self.signature:
+            item['frame_nr'] = frame.frame_nr
+        if 'pose_shape' in self.signature:
+            pose_shape = np.load(os.path.join(self.folder, 'pose_shape', frame.id + separator + frame.frame_nr + '.npy')).astype(np.float32)
+            pose_shape[:3] = 0
+            item['pose_shape'] = pose_shape
+        if 'mesh_body_posed' in self.signature:
+            mesh_posed_body_no_orient = np.load(os.path.join(self.folder, 'mesh_posed_body_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
+            item['mesh_body_posed'] = mesh_posed_body_no_orient
+        if 'mesh_posed' in self.signature:
+            mesh_posed = np.load(os.path.join(self.folder, 'mesh_posed_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed_no_orient'].astype(np.float32)
+            item['mesh_posed'] = mesh_posed
+        if 'mesh_unposed' in self.signature:
+            mesh_unposed_no_shape = np.load(os.path.join(self.folder, 'mesh_unposed_no_shape', frame.id + separator + frame.frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
+            item['mesh_unposed'] = mesh_unposed_no_shape
+        if 'mesh_curvature' in self.signature:
+            mesh_curvature = np.load(os.path.join(self.folder, 'mesh_curvature', frame.id + separator + frame.frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
+            item['mesh_curvature'] = mesh_curvature.astype(np.int32)
+        if 'mesh_normals' in self.signature:
+            mesh_normals = np.load(os.path.join(self.folder, 'mesh_normals', frame.id + separator + frame.frame_nr + '.npz')); mesh_normals = mesh_normals['mesh_normals']
+            item['mesh_normals'] = mesh_normals
+
+        for res in self.resolutions:
+            res_string = ''
+            if res != 512:
+                res_string = '_' + str(res)
+
+            if 'uv_curvature' + res_string in self.signature:
+                uv_curvature = np.load(os.path.join(self.folder, 'uv_curvature' + res_string, frame.id + separator + frame.frame_nr + '.npz')); 
+                if res_string != '':
+                    uv_curvature = uv_curvature['uv_curvature'] 
+                else:
+                    uv_curvature = uv_curvature['curvature_uv']
+                uv_curvature = torch.nn.functional.one_hot(torch.tensor(uv_curvature, device=self.device, dtype=torch.long), 7).float().to(self.device)
+                item['uv_curvature' + res_string] = uv_curvature
+
+            if 'uv_unposed' + res_string in self.signature:
+                uv_unposed_no_shape = np.load(os.path.join(self.folder, 'uv_unposed_no_shape' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']
+                item['uv_unposed' + res_string] = uv_unposed_no_shape
+
+            if 'uv_body_posed' + res_string in self.signature:
+                uv_posed_body_no_orient = np.load(os.path.join(self.folder, uv_posed_body_no_orient_folder_name + res_string, frame.id + separator + frame.frame_nr + '.npz'))
+                uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']
+                item['uv_body_posed' + res_string] = uv_posed_body_no_orient
+
+            if 'uv_posed' + res_string in self.signature:
+                uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient']
+                item['uv_posed' + res_string] = uv_posed
+
+            if 'uv_normals' + res_string in self.signature:
+                uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_normals = uv_normals['uv_normals']
+                item['uv_normals' + res_string] = uv_normals
+
+
+        data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in item.items()}
+
+        for key, value in data.items():
+            if key.startswith('uv_') and (value.shape[-1] == 3 or value.shape[-1] == 7):
+                data[key] = torch.permute(value, [2,0,1])
+
+        if self.cache is not None:
+            self.cache[frame] = data
+        
+        data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in data.items()}
+        data = { **data, **self.static_cache[frame.id] }
+        return data
+
+
+if __name__ == '__main__':
+    from torch.utils.data import DataLoader
+    
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    signature = { 'id', 
+                  'frame_nr',
+                  'pose_shape',
+                  'uv_posed',
+                  'uv_body_posed',
+                  'mesh_body_posed',
+                  'mesh_posed',
+                  'mesh_unposed',
+                  'uv_unposed',
+                  'uv_curvature',
+                  'mesh_curvature',
+                  'uv_static',
+                  'mesh_static',
+                  'mesh_mask',
+                  'uv_mask',
+                  'uv_shaped_body',
+                  }
+    
+    dataset = Dataset('/data-local/127/train', signature, resolutions=[512], device=device, frames_file='frames_jumpsuit.txt')
+    train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
+
+    for batch in train_dataloader:
+        batch
+        pass
+
+    pass 
+
+
+
+
+
+
+
+
+
+
+# if __name__ == '__main__':
+#     os.environ["CUDA_VISIBLE_DEVICES"] = "6"
+#     from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
+#     import plotly.graph_objects as go
+#     import plotly.express as ex
+#     # import cmapy as cm
+#     uv2mesh = UV2Mesh(is_skirt=True)
+
+#     F = np.load("custom_utils/smpl_faces.npy")
+
+#     sum_uv = np.zeros((512, 256, 3), dtype=np.float32)
+#     sum_mask = np.zeros((512, 256), dtype=np.float32)
+#     for i, batch in enumerate(dataset):
+#         print(i)
+#         batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
+#         ex.imshow(batch['uv_curvature'][0].numpy().argmax(-1) / 6.0).show()
+#         sum_uv += batch['uv_unposed'].numpy().sum(axis=0)
+#         sum_mask += batch['uv_mask'].numpy().sum(axis=0)
+
+#     sum_mask[sum_mask == 0] = 1
+#     avg_garment = sum_uv / sum_mask[:, :, np.newaxis]
+
+#     # x = np.transpose(avg_garment, [1,2,0])
+#     x = uv2mesh(avg_garment[tf.newaxis,:,:,:])[0]
+#     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]
+#             )]).show()
+
+#     np.save('custom_data/train_dress_30k_no_shape_avg_uv.npy', avg_garment)
+
+#     print()

docker/Dockerfile2

@@ -1,7 +1,7 @@
 # Use an existing image as the base
-FROM nvcr.io/nvidia/pytorch:23.04-py3
+# FROM nvcr.io/nvidia/pytorch:23.04-py3
 # FROM pytorch/pytorch:1.13.1-cuda11.6-cudnn8-devel
-# FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-devel
+FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel
 # 
 # Add the current user to the image
 ARG USER_NAME=lhunor
@@ -10,7 +10,7 @@ ARG GROUP_NAME=guest
 ARG GROUP_ID=11000
 
 # Update the system and install required packages
-RUN apt-get update && apt-get upgrade -y && apt-get install -y sudo git apt-utils htop
+RUN apt-get update && apt-get install -y sudo git apt-utils htop
 
 # Create a new group with the same GID as the host
 RUN groupadd --gid $GROUP_ID $GROUP_NAME
@@ -25,6 +25,12 @@ RUN echo '%sudo ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers
 # Add the new user to the sudo group
 RUN usermod -aG sudo $USER_NAME
 
+# install conda stuff still as root
+RUN conda update -n base conda
+RUN conda install -c fvcore -c iopath -c conda-forge fvcore iopath -y
+RUN conda install -c bottler nvidiacub -y
+RUN conda install pytorch3d -c pytorch3d -y
+
 # Set the default user to the new user
 USER $USER_NAME
 
@@ -33,8 +39,8 @@ WORKDIR /workspace
 ENV PATH="$PATH:/home/$USER_NAME/.local/bin"
 RUN pip install --upgrade pip
 COPY requirements.txt /tmp/
-RUN pip install --requirement /tmp/requirements.txt
 RUN pip install torch-cluster -f https://data.pyg.org/whl/torch-2.0.0+cu117.html
+RUN pip install --requirement /tmp/requirements.txt
 
 
 # defining the volumes that will be mounted

docker/docker-compose.yml

@@ -26,6 +26,28 @@ services:
     container_name: lhunor_torch_nv23.04
     ipc: host
     tty: true
+    ulimits:
+      memlock: -1
+      stack: 67108864
+    volumes:
+      - '/home-net/lhunor:/home-net'
+      - '/data-net/guest/lhunor/data:/data-net'
+      - '/data-fast/126-data1/lhunor/data:/data-local/126'
+      - '/data-fast/127-data1/lhunor/data:/data-local/127'
+      - '/data-fast/127-data2/lhunor/data:/data-local/127_2'
+      - '/data-fast/128-data1/lhunor/data:/data-local/128'
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - capabilities: [gpu]
+
+
+  lhunor_iccvw:
+    image: lhunor_iccvw:pt2.0.1
+    container_name: lhunor_iccvw_pt2.0.1
+    ipc: host
+    tty: true
     ulimits:
       memlock: -1
       stack: 67108864

docker/requirements.txt

@@ -1,13 +1,10 @@
-torch==2.0.0    
 numpy
-matplotlib
 opencv-python
 wandb
 scipy
 plotly
 pandas
 cmapy
-pytorch3d
 tables
 numpy-stl
 torchmetrics

scripts/ae_body_train.py

@@ -14,7 +14,7 @@ args = vars(parser.parse_args())
 config_module = importlib.import_module(args['config'])
 config = config_module.config
 if args['debug']:
-    from configs.ae_body_cape_m import config
+    from configs.iccvw.ae_body_cape_m import config
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 if args['debug']:
@@ -49,6 +49,7 @@ from pytorch_msssim import MS_SSIM
 
 
 def init_wandb():
+    wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
     run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('scripts/ae_body_train.py')
@@ -198,15 +199,15 @@ if __name__ == '__main__':
         load_model(model.encoder, config['pretrained_model'] + '_enc')
         load_model(model.decoder, config['pretrained_model'] + '_dec')
 
-    ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_enc', top_n=5)
-    ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_dec', top_n=5)
+    ckpt_saver_enc = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_enc', top_n=5)
+    ckpt_saver_dec = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_dec', top_n=5)
 
     run_name = config['name']
     print(f'started training: {run_name}')
 
     if not args['debug']:
         # model = torch.compile(model)
-        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
+        wandb.watch(model, log='all', log_freq=10000, log_graph=True)
 
     train()
 

scripts/ae_norm_train.py

@@ -49,7 +49,8 @@ from pytorch_msssim import MS_SSIM
 
 
 def init_wandb():
-    run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127_2/temp', tags=config['tags'])
+    wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
+    run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('scripts/ae_norm_train.py')
     artifact.add_file('losses/losses.py')
@@ -70,9 +71,9 @@ uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_
 
 def train_step(batch):
     model.train()
-    pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.1
+    pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.4
 
-    ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_normals'] + 1.1))
+    ms_ssim = torch.mean(_MSSSIM(pred + 1.4, batch['uv_normals'] + 1.4))
     loss = 1e3 * (1 - ms_ssim)
 
     optimizer.zero_grad()
@@ -93,13 +94,13 @@ def valid_step(batch):
     model.eval()
 
     with torch.inference_mode():
-        pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.1
+        pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.4
 
         mesh_pred = uv2mesh(pred)
         mesh_orig = batch['mesh_normals']
 
         l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_normals'] - pred)), axis=1))
-        ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_normals'] + 1.1))
+        ms_ssim = torch.mean(_MSSSIM(pred + 1.4, batch['uv_normals'] + 1.4))
         loss = 1e3 * (1 - ms_ssim)
 
         metric = torch.mean(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1)))
@@ -109,7 +110,7 @@ def valid_step(batch):
 
 def plot_everything(batch: dict, type: str):
     model.eval()
-    pred = model([batch['uv_normals'], batch['uv_mask']]).detach() * 1.1
+    pred = model([batch['uv_normals'], batch['uv_mask']]).detach() * 1.4
     
     # mesh_pred = uv2mesh(pred)
     # mesh_orig = batch['mesh_normals']
@@ -117,7 +118,7 @@ def plot_everything(batch: dict, type: str):
     orig_img = torch.permute(batch['uv_normals'], [0,2,3,1])
     pred_img = torch.permute(pred, [0,2,3,1])
     img = torch.concat((orig_img, pred_img), axis=2)
-    imgs = [wandb.Image((img[i].cpu().numpy() + 1.1) / 2.2) for i in range(4)]
+    imgs = [wandb.Image((img[i].cpu().numpy() + 1.4) / 2.8) for i in range(4)]
 
     # F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
     # garment_unposed, faces = mesh_pred.cpu().numpy(), F
@@ -196,7 +197,7 @@ if __name__ == '__main__':
     optimizer = optim.Adam(model.parameters(), lr=config['lr'])
 
     # _MSSSIM = MSSSIM(data_range=2).to(config['device'])
-    _MSSSIM = MS_SSIM(data_range=2.2, size_average=True, channel=3)
+    _MSSSIM = MS_SSIM(data_range=2.8, size_average=True, channel=3)
 
     if config['pretrained_model'] != '':
         load_model(model.encoder, config['pretrained_model'] + '_enc')

scripts/ae_static_train.py

@@ -13,8 +13,8 @@ args = vars(parser.parse_args())
 config_module = importlib.import_module(args['config'])
 config = config_module.config
 if args['debug']:
-    from configs.ae_static_cape_m import config
-    # from configs.ae_static_orig_2 import config
+    # from configs.ae_static_cape_m import config
+    from configs.iccvw.ae_static_1 import config
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 if args['debug']:
@@ -49,6 +49,7 @@ from pytorch_msssim import MS_SSIM
 
 
 def init_wandb():
+    wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
     run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('scripts/ae_static_train.py')
@@ -208,15 +209,15 @@ if __name__ == '__main__':
         load_model(model.encoder, config['pretrained_model'] + '_enc')
         load_model(model.decoder, config['pretrained_model'] + '_dec')
 
-    ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_enc', top_n=5)
-    ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_dec', top_n=5)
+    ckpt_saver_enc = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_enc', top_n=5)
+    ckpt_saver_dec = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_dec', top_n=5)
 
     run_name = config['name']
     print(f'started training: {run_name}')
 
     if not args['debug']:
         # model = torch.compile(model)
-        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
+        wandb.watch(model, log='all', log_freq=10000, log_graph=True)
 
     train()