train baseline

configs/iccvw/abl_all.py

@@ -0,0 +1,195 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '2'
+config['name'] = 'abl_iccvw_all_1'
+config['notes'] = ''
+config['tags'] = ['iccv', 'abl_all', 'torch']
+config['script'] = 'scripts/cvae.py'
+
+config['pretrained_model'] = ''
+
+config['use_skinning'] =              True # and deshaping
+config['use_normals'] =             True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+
+config['use_norm_loss'] =             True
+
+ablation_count = [config['use_skinning'], config['use_normals'], config['use_extra_losses']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['avg_garment_path'] = 'custom_data/avg/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['avg_garment_path'] != '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['device'] = 'cuda'
+
+config['train_dir'] = '/data-local/127/cloth3d_processed'
+config['valid_dir'] = '/data-local/127/cloth3d_processed'
+config['test_dir'] = '/data-net/dataset/test_ready'
+
+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['static_model'] = 'checkpoints/iccvw/static_iccvw_4'
+config['body_model'] = 'checkpoints/iccvw/body_iccvw_1'
+config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.0001
+config['balancing_activation_limit'] = 0.01
+
+config['signature']: set = { 'id',
+                            'frame_nr',
+                            'pose_shape',
+                            'uv_body_posed',
+                            'mesh_body_posed',
+                            'mesh_posed',
+                            'uv_normals',
+                            'uv_static',
+                            'mesh_mask',
+                            'uv_mask',
+                        }
+
+if config['use_skinning']:
+  config['signature'].add('uv_unposed')
+  config['signature'].add('mesh_unposed')
+else:
+  config['signature'].add('uv_posed')
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+config['bottleneck_scale'] = 2
+
+config['cvae'] = dict()
+config['cvae']['device'] = config['device']
+config['cvae']['is_label'] = False
+config['cvae']['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
+config['cvae']['kernel_size'] = 7
+config['cvae']['bottleneck_size'] = 64 * config['bottleneck_scale']
+config['cvae']['latent_dim'] = config['cvae']['bottleneck_size']
+
+block_setting = [
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 3),
+    CNBlockConfig(128, 256, 3),
+    CNBlockConfig(256, config['cvae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cvae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cvae']['bottleneck_size'] * 2, 256, 3),
+    CNBlockConfig(256, 128, 9),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 3),
+    CNBlockConfig(32, 3, 3),
+]
+
+config['cvae']['encoder_block_settings'] = block_setting
+config['cvae']['decoder_block_settings'] = block_setting_reverse
+
+
+
+####################### cond_ae_norm ############################
+config['cond_ae_norm'] = dict()
+config['cond_ae_norm']['device'] = config['device']
+config['cond_ae_norm']['latent_dim'] = config['latent_dim']
+config['cond_ae_norm']['kernel_size'] = 7
+config['cond_ae_norm']['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_norm']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_norm']['encoder_block_settings'] = block_setting
+config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_norm']['activation'] = 'tanh'
+
+####################### cond_ae_body ############################
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['device'] = config['device']
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['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_body']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae_body']['encoder_block_settings'] = block_setting
+config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_body']['activation'] = 'tanh'
+
+####################### cond_ae_static ############################
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['device'] = config['device']
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['latent_dim'] = config['latent_dim']
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['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_static']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+]
+
+config['cond_ae_static']['encoder_block_settings'] = block_setting
+config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae_static']['activation'] = 'tanh'

configs/iccvw/ae_norm_1.py

@@ -58,7 +58,6 @@ block_setting_reverse = [
     CNBlockConfig(64, 32, 9),
     CNBlockConfig(32, 16, 3),
     CNBlockConfig(16, 3, 3),
-    # CNBlockConfig(3, 3, 3),
 ]
 
 config['cond_ae']['encoder_block_settings'] = block_setting

data/preprocessing.py

@@ -679,7 +679,8 @@ def check_irregular_poses(folder, folder_reg, out_folder):
     # mesh_static = np.load(f'{out_folder}/mesh_static/{frame.id}_{frame.garment}.npy')
 
     mesh_mask = np.load(f'{out_folder}/mesh_mask/{frame.id}_{frame.garment}.npy')
-    mesh_body_posed = readPC2Frame(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), frame.frame_nr, True)
+    mesh_body_posed = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), True)['V']
+    mesh_body_posed = mesh_body_posed[int(mesh_body_posed.shape[0]/2)+frame.frame_nr]
 
     F = np.load('smpl/smpl_faces_skirt.npy')
     _, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask)
@@ -694,6 +695,46 @@ def check_irregular_poses(folder, folder_reg, out_folder):
     fig.show()
     pass
 
+
+def calculate_average(folder, out_folder, frames_file='frames_train_skirt.txt'):
+    from custom_utils.custom_utils import UV2Mesh
+    import plotly.express as ex
+    import plotly.graph_objects as go
+
+    frames = []
+    with open(os.path.join(out_folder, frames_file), 'r') as f:
+        for line in f:
+            line = line.strip()
+            id, frame_nr, garment = line.split(' ')
+            frames.append(Frame(id, int(frame_nr), garment))
+
+    uv_sum = np.zeros((512, 256, 3), dtype=np.float32)
+    mask_sum = np.zeros((512, 256), dtype=np.float32)
+
+    for frame in tqdm(frames):
+        uv = np.load(os.path.join(out_folder, 'uv_unposed_no_shape', f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))['uv_unposed']
+        mask = np.load(os.path.join(out_folder, 'uv_mask', f'{frame.id}_{frame.garment}.npz'))['uv_mask']
+        uv_sum += uv
+        mask_sum += mask
+
+    mask_sum[mask_sum == 0] = 1
+    avg_garment = uv_sum / mask_sum[:, :, None]
+
+    avg_mesh = UV2Mesh()(torch.from_numpy(avg_garment.transpose(2,0,1))[None,:,:,:]).numpy()[0]
+    F = np.load('smpl/smpl_faces.npy')
+
+    fig = go.Figure(data=[go.Mesh3d(x=avg_mesh[:,0], y=avg_mesh[:,2], z=avg_mesh[:,1], i=F[:,0], j=F[:,1], k=F[:,2], color='lightpink', opacity=1)])
+    fig.update_layout(scene = dict(
+        xaxis = dict(range=[-1,1],),
+                    yaxis = dict(range=[-1,1],),
+                    zaxis = dict(range=[-1,1],),),)
+    fig.update_layout(scene_aspectmode='cube')
+    fig.show()
+
+    np.save('custom_data/avg/train_skirt_no_shape_avg_uv.npy', avg_garment)
+
+
+
 if __name__ == '__main__':
 
 
@@ -721,9 +762,10 @@ if __name__ == '__main__':
     # create_mesh_body_posed_no_orient(folder, out_folder)
     # create_mesh_normals(folder, out_folder)
     # create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end)
+    calculate_average(folder, out_folder)
 
     # create_frames_file(folder, out_folder)
-    create_frames_file_partial(folder, out_folder, part=0.5)
+    # create_frames_file_partial(folder, out_folder, part=0.5)
     
     # folder = '/data-local/127_2/cape_processed/test/female'
     # create_mesh_normals_cape(folder, None)

models/cvae_model.py

@@ -18,14 +18,13 @@ class CVAE(nn.Module):
         super(CVAE, self).__init__()
         self.size = config['resolution']
 
-        if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
-            self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'], device=config['device'])
-            if config['static_model'] != '':
-                load_model(self.static_enc, config['static_model'] + '_enc')
+        self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'], device=config['device'])
+        if config['static_model'] != '':
+            load_model(self.static_enc, config['static_model'] + '_enc')
 
-            self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], config['cond_ae_body']['kernel_size'], config['cond_ae_body']['activation'], device=config['device'])
-            if config['body_model'] != '':
-                load_model(self.body_enc, config['body_model'] + '_enc')
+        self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], config['cond_ae_body']['kernel_size'], config['cond_ae_body']['activation'], device=config['device'])
+        if config['body_model'] != '':
+            load_model(self.body_enc, config['body_model'] + '_enc')
 
         if config['use_normals']:
             self.norm_enc : AEEncoder = AEEncoder(config['cond_ae_norm']['encoder_block_settings'], config['cond_ae_norm']['kernel_size'], config['cond_ae_norm']['activation'], device=config['device'])

scripts/ae_norm_train.py

@@ -205,15 +205,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=10000, log_graph=True)
+        wandb.watch(model, log='all', log_freq=20000, log_graph=True)
 
     train()
 

scripts/cvae.py

@@ -10,10 +10,12 @@ parser.add_argument('--config', type=str)
 parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
-config_module = importlib.import_module(args['config'])
-config = config_module.config
+
 if args['debug']:
-    from configs.abl_norm import config
+    from configs.iccvw.abl_all import config
+else:
+    config_module = importlib.import_module(args['config'])
+    config = config_module.config
 
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
@@ -42,6 +44,7 @@ from data.dataset import Dataset
 from models.cond_encoder import AutoEncoder
 from utils.checkpoints import CheckpointSaver, load_model
 from utils.utils import InfiniteDataLoader as DataLoader
+from utils.utils import text_to_html
 from utils.postprocess import fix_collision
 from losses.losses import collision_loss, normal_loss
 from pytorch_msssim import MS_SSIM
@@ -51,7 +54,8 @@ from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary
 
 
 def init_wandb():
-    run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
+    wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
+    run = wandb.init(project='iccvw', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('scripts/cvae.py')
     artifact.add_file('losses/losses.py')
@@ -255,7 +259,14 @@ def plot_everything(batch: dict, type: str):
     mean_hist = wandb.Histogram(sequence=mean.cpu().numpy().mean(axis=0))
     logvar_hist = wandb.Histogram(sequence=logvar.cpu().numpy().mean(axis=0))
 
-    return { **figs, **figs_coll, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
+    batch_id = batch['id']
+    batch_frame_nr = batch['frame_nr'].cpu().numpy()
+    batch_garment = batch['garment']
+    batch_info = [f'batch_info: id: {batch_id[i]}, frame_nr: {batch_frame_nr[i]}, garment: {batch_garment[i]}' for i in range(len(batch_id))]
+    batch_info = text_to_html(batch_info)
+    batch_info = wandb.Html(batch_info)
+
+    return { **figs, **figs_coll, **{ f'{type}/batch_info': batch_info }, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
 
 
 def sample(folder): 
@@ -455,10 +466,10 @@ def train():
             valid_accumulator.reset()
             print(step)
 
-        if train_loss['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
+        if train_loss['train/euclid_error'] < config['balancing_activation_limit']:
             g_balancing_coeff += 1 / 10000 
 
-        if step % 10000 == 0:
+        if step % 25000 == 0:
             ckpt_saver(model, step, optimizer, None)
 
         step += 1
@@ -485,7 +496,7 @@ if __name__ == '__main__':
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
     g_balancing_coeff = config['starting_balancing_coeff']
 
-    ckpt_saver = CheckpointSaver('checkpoints/iccv/' + config['name'], top_n=5)
+    ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=10)
 
     step = 0
     if config['pretrained_model'] != '':

utils/checkpoints.py

@@ -39,5 +39,6 @@ def load_model(model: torch.nn.Module, checkpoint, optimizer: torch.nn.Module=No
     model.load_state_dict(checkpoint['model'])
     if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer'].state_dict())
     if scaler is not None and checkpoint['scaler'] is not None: scaler.load_state_dict(checkpoint['scaler'])
+    print(f'Loaded model from {checkpoints[-1]} at iteration {iteration}')
     return int(iteration) if iteration != '' else 0
 

utils/utils.py

@@ -42,4 +42,10 @@ class InfiniteDataLoader(DataLoader):
             # Dataset exhausted, use a new fresh iterator.
             self.dataset_iterator = super().__iter__()
             batch = next(self.dataset_iterator)
-        return batch
+        return batch
+
+def text_to_html(lines):  
+    html = ''
+    for line in lines:
+        html += "<pre>" + line + "</pre>"
+    return html