From 0543657e4e2b921a7ee1537e3c2ee46c912218c5 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sat, 28 Jan 2023 12:40:28 +0000 Subject: [PATCH] prepare to train --- configs/config_abl_all_10p_curv_10_r2.py | 129 +++++++++++++++++++++++ cvae.py | 8 +- dataset.py | 7 +- 3 files changed, 138 insertions(+), 6 deletions(-) create mode 100644 configs/config_abl_all_10p_curv_10_r2.py diff --git a/configs/config_abl_all_10p_curv_10_r2.py b/configs/config_abl_all_10p_curv_10_r2.py new file mode 100644 index 0000000..710053a --- /dev/null +++ b/configs/config_abl_all_10p_curv_10_r2.py @@ -0,0 +1,129 @@ +config = dict() + +# DON'T FORGET TO RESET SSIM LOSS USAGE +# DON'T FORGET TO CHECK MODEL LOADING +# DON'T FORGET TO CHECK IF STORE IN MEMORY +# DON'T FORGET TO CHECK REG LOSS COEFF +# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512 +# DON'T FORGET TO CHECK AVG GARMENT USED + +config['use_skinning'] = True # and deshaping +config['remove_global_orientation'] = True +config['use_curvature'] = True +config['use_extra_losses'] = True # edge, collision, 3d mesh +config['use_offsets'] = True + +config['use_norm_loss'] = True +config['only_curvatures'] = False + +config['curv_neighbourhood'] = 10 + +ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False) +assert ablation_count == 1 or ablation_count == 0 + +config['gpu'] = '6' +config['name'] = 'abl_all_10p_curv_10_r2' +config['notes'] = '' +config['tags'] = ['review_2'] + +config['avg_garment_path'] = 'custom_data/train_no_shape_skirt_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['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and 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'] = True +config['dataset_in_memory'] = True +config['is_cape'] = False + +config['train_dir'] = '/home/workspace/data/dataset/local-cache/train_skirt' +config['valid_dir'] = '/home/workspace/local-data/valid_skirt' +config['test_dir'] = '/home/workspace/data/dataset/test_skirt_ready_2' + +config['frames_file'] = 'frames_skirt_10p.txt' +config['frames_file_valid'] = 'frames_skirt_10p.txt' +# config['frames_file_test'] = 'frames_dress_30k.txt' + +config['static_model'] = 'checkpoints/static_1_skirt' +config['body_model'] = 'checkpoints/body_1_skirt_2_2' +config['curvature_model'] = 'checkpoints/curv_1_skirt' + + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + + +config['pretrained_model'] = '' + +############### +config['lr'] = 0.0001# 0.0001 +############### +config['starting_balancing_coeff'] = 0.001 + +config['signature_dynamic'] = { 'id': (), 'frame_nr': (), + 'pose_shape': (72 + 11,), + 'uv_body_posed': (512, 256, 3), + 'mesh_body_posed': (config['smpl_res'], 3,), + 'mesh_posed': (config['smpl_res'], 3,), + 'uv_curvature': (512, 256, 7), + } +config['signature_static'] = { 'uv_static': (512, 256, 3), + 'mesh_mask': (config['smpl_res'],), + 'uv_mask': (512, 256,), + } + +if config['use_skinning']: + config['signature_dynamic']['uv_unposed'] = (512, 256, 3) + config['signature_dynamic']['mesh_unposed'] = (config['smpl_res'], 3,) +else: + config['signature_dynamic']['uv_posed'] = (512, 256, 3) + + +config['upscale_enhancer'] = '' +config['freeze_encoders'] = True +config['use_one_hot'] = False + +config['cvae'] = dict() +config['cvae']['is_label'] = False +config['cvae']['add_noise'] = True +config['cvae']['noise_amount'] = 0.001 + +config['cvae']['kernel_size'] = 7 +config['cvae']['encoder_depths'] = [3, 3, 3, 9, 3] +config['cvae']['encoder_dims'] = [32, 64, 128, 256, 64] +config['cvae']['decoder_depths'] = config['cvae']['encoder_depths'] +config['cvae']['decoder_dims'] = config['cvae']['encoder_dims'] + + +config['cond_ae_static'] = dict() +config['cond_ae_static']['is_label'] = False +config['cond_ae_static']['kernel_size'] = 7 +config['cond_ae_static']['bottleneck_size'] = 8 +config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 9, 3] +config['cond_ae_static']['encoder_dims'] = [16, 32, 64, 64, config['cond_ae_static']['bottleneck_size']] +config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths'] +config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims'] +config['cond_ae_static']['activation'] = 'tanh' + +config['cond_ae_curv'] = dict() +config['cond_ae_curv']['is_label'] = True +config['cond_ae_curv']['kernel_size'] = 7 +config['cond_ae_curv']['bottleneck_size'] = 40 +config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3] +config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']] +config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths'] +config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims'] +config['cond_ae_curv']['activation'] = 'tanh' + +config['cond_ae_body'] = dict() +config['cond_ae_body']['is_label'] = False +config['cond_ae_body']['kernel_size'] = 7 +config['cond_ae_body']['bottleneck_size'] = 16 +config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3] +config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']] +config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths'] +config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims'] +config['cond_ae_body']['activation'] = 'tanh' \ No newline at end of file diff --git a/cvae.py b/cvae.py index 59d7c75..d25acf6 100644 --- a/cvae.py +++ b/cvae.py @@ -1,6 +1,6 @@ import os import sys -from configs.config_abl_all_post import config +from configs.config_abl_all_10p_curv_10_r2 import config os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] import tensorflow as tf @@ -494,7 +494,8 @@ if __name__ == '__main__': use_one_hot=config['use_one_hot'], use_no_orient=config['remove_global_orientation'], frames_file=config['frames_file'] if 'frames_file' in config.keys() else None, - is_cape=config['is_cape']) + is_cape=config['is_cape'], + curv_neighbourhood=config['curv_neighbourhood']) static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], resolutions=[config['resolution']], @@ -509,7 +510,8 @@ if __name__ == '__main__': use_one_hot=config['use_one_hot'], use_no_orient=config['remove_global_orientation'], frames_file=config['frames_file_valid'] if 'frames_file_valid' in config.keys() else None, - is_cape=config['is_cape']) + is_cape=config['is_cape'], + curv_neighbourhood=config['curv_neighbourhood']) static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], resolutions=[config['resolution']], diff --git a/dataset.py b/dataset.py index d2ee42d..d27b7de 100644 --- a/dataset.py +++ b/dataset.py @@ -194,7 +194,8 @@ class DatasetFactory(): use_one_hot=False, use_no_orient=True, testing=False, - frames_file=None): + frames_file=None, + curv_neighbourhood=25): assert set(resolutions).issubset(set([32, 64, 128, 256, 512])) dataset_idx = DatasetIdx(folder, shuffle, is_cape=is_cape, testing=testing, frames_file=frames_file) @@ -260,7 +261,7 @@ class DatasetFactory(): if 'mesh_curvature' in signature.keys(): - mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', id + separator + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature'] + mesh_curvature = np.load(os.path.join(folder, f'mesh_curvature_{curv_neighbourhood}', id + separator + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature'] item['mesh_curvature'] = mesh_curvature.astype(np.int32) @@ -270,7 +271,7 @@ class DatasetFactory(): res_string = '_' + str(res) if 'uv_curvature' + res_string in signature.keys(): - uv_curvature = np.load(os.path.join(folder, 'uv_curvature' + res_string, id + separator + frame_nr + '.npz')); + uv_curvature = np.load(os.path.join(folder, f'uv_curvature_{curv_neighbourhood}' + res_string, id + separator + frame_nr + '.npz')); if res_string != '': uv_curvature = uv_curvature['uv_curvature'] else: