diff --git a/configs/cape_00032_SL.py b/configs/cape_00032_SL.py new file mode 100644 index 0000000..fa3ce44 --- /dev/null +++ b/configs/cape_00032_SL.py @@ -0,0 +1,193 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '0' +config['name'] = 'cape_00032_SL' +config['notes'] = '' +config['tags'] = ['iccv', 'cape_00032_SL', 'torch'] +config['script'] = 'scripts/cvae.py' + +config['pretrained_model'] = '' + +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 + +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['avg_garment_path'] = 'custom_data/avg/train_cape_f_03223_SL_2_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['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'] = False +config['dataset_in_memory'] = True +config['is_cape'] = True +config['device'] = 'cuda' + +config['train_dir'] = '/data-local/127/cape_processed/train/male' +config['valid_dir'] = '/data-local/127/cape_processed/train/male' +config['test_dir'] = '/data-local/127/cape_processed/train/male' +config['train_frames'] = 'frames_00032_SL_train_2.txt' +config['valid_frames'] = 'frames_00032_SL_test_2.txt' +config['test_frames'] = 'frames_00032_SL_test_2.txt' + +config['static_model'] = 'checkpoints/iccv/torch/static_torch_cape_m' +config['body_model'] = 'checkpoints/iccv/torch/body_torch_cape_m' +config['curvature_model'] = 'checkpoints/iccv/torch/curv_torch_cape_m' + +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['signature']: set = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'uv_curvature', + '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+7 if config['use_curvature'] 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_curv ############################ +config['cond_ae_curv'] = dict() +config['cond_ae_curv']['device'] = config['device'] +config['cond_ae_curv']['is_label'] = True +config['cond_ae_curv']['latent_dim'] = config['latent_dim'] +config['cond_ae_curv']['kernel_size'] = 7 +config['cond_ae_curv']['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_curv']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_curv']['encoder_block_settings'] = block_setting +config['cond_ae_curv']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_curv']['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' \ No newline at end of file diff --git a/configs/cape_03223_SL.py b/configs/cape_03223_SL.py new file mode 100644 index 0000000..4fde966 --- /dev/null +++ b/configs/cape_03223_SL.py @@ -0,0 +1,193 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '4' +config['name'] = 'cape_03223_SL' +config['notes'] = '' +config['tags'] = ['iccv', 'cape_03223_SL', 'torch'] +config['script'] = 'scripts/cvae.py' + +config['pretrained_model'] = '' + +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 + +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['avg_garment_path'] = 'custom_data/avg/train_cape_f_03223_SL_2_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['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'] = False +config['dataset_in_memory'] = True +config['is_cape'] = True +config['device'] = 'cuda' + +config['train_dir'] = '/data-local/127/cape_processed/train/female' +config['valid_dir'] = '/data-local/127/cape_processed/train/female' +config['test_dir'] = '/data-local/127/cape_processed/train/female' +config['train_frames'] = 'frames_03223_SL_train_2.txt' +config['valid_frames'] = 'frames_03223_SL_test_2.txt' +config['test_frames'] = 'frames_03223_SL_test_2.txt' + +config['static_model'] = 'checkpoints/iccv/torch/static_torch_cape_f' +config['body_model'] = 'checkpoints/iccv/torch/body_torch_cape_f' +config['curvature_model'] = 'checkpoints/iccv/torch/curv_torch_cape_f' + +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['signature']: set = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'uv_curvature', + '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+7 if config['use_curvature'] 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_curv ############################ +config['cond_ae_curv'] = dict() +config['cond_ae_curv']['device'] = config['device'] +config['cond_ae_curv']['is_label'] = True +config['cond_ae_curv']['latent_dim'] = config['latent_dim'] +config['cond_ae_curv']['kernel_size'] = 7 +config['cond_ae_curv']['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_curv']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_curv']['encoder_block_settings'] = block_setting +config['cond_ae_curv']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_curv']['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' \ No newline at end of file diff --git a/configs/cape_03223_SS.py b/configs/cape_03223_SS.py new file mode 100644 index 0000000..bec63b7 --- /dev/null +++ b/configs/cape_03223_SS.py @@ -0,0 +1,193 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '4' +config['name'] = 'cape_03223_SS' +config['notes'] = '' +config['tags'] = ['iccv', 'cape_03223_SS', 'torch'] +config['script'] = 'scripts/cvae.py' + +config['pretrained_model'] = '' + +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 + +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['avg_garment_path'] = 'custom_data/avg/train_cape_f_03223_SL_2_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['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'] = False +config['dataset_in_memory'] = True +config['is_cape'] = True +config['device'] = 'cuda' + +config['train_dir'] = '/data-local/127/cape_processed/train/female' +config['valid_dir'] = '/data-local/127/cape_processed/train/female' +config['test_dir'] = '/data-local/127/cape_processed/train/female' +config['train_frames'] = 'frames_03223_SS_train_2.txt' +config['valid_frames'] = 'frames_03223_SS_test_2.txt' +config['test_frames'] = 'frames_03223_SS_test_2.txt' + +config['static_model'] = 'checkpoints/iccv/torch/static_torch_cape_f' +config['body_model'] = 'checkpoints/iccv/torch/body_torch_cape_f' +config['curvature_model'] = 'checkpoints/iccv/torch/curv_torch_cape_f' + +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['signature']: set = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'uv_curvature', + '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+7 if config['use_curvature'] 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_curv ############################ +config['cond_ae_curv'] = dict() +config['cond_ae_curv']['device'] = config['device'] +config['cond_ae_curv']['is_label'] = True +config['cond_ae_curv']['latent_dim'] = config['latent_dim'] +config['cond_ae_curv']['kernel_size'] = 7 +config['cond_ae_curv']['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_curv']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_curv']['encoder_block_settings'] = block_setting +config['cond_ae_curv']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_curv']['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' \ No newline at end of file diff --git a/custom_data/avg/train_cape_f_03223_SL_2_no_shape_avg_uv.npy b/custom_data/avg/train_cape_f_03223_SL_2_no_shape_avg_uv.npy new file mode 100644 index 0000000..0a1c0f4 Binary files /dev/null and b/custom_data/avg/train_cape_f_03223_SL_2_no_shape_avg_uv.npy differ diff --git a/data/dataset.py b/data/dataset.py index f85c5ad..749abe1 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -45,7 +45,10 @@ def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file= 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 is_cape: + frames = [Frame(line.split('.')[0], line.split('.')[1]) for line in lines] + else: + frames = [Frame(line.split('_')[0], line.split('_')[1]) for line in lines] if shuffle: random.shuffle(frames) diff --git a/run.py b/run.py index e1d7d52..ea58dcf 100644 --- a/run.py +++ b/run.py @@ -1,6 +1,6 @@ import os import sys -from configs.pyr_64_cape_m import config +from configs.cape_03223_SS import config import argparse # parser = argparse.ArgumentParser(description='') @@ -10,9 +10,9 @@ import argparse if __name__ == '__main__': - import wandb + # import wandb # wandb.login(key='******************************************') - wandb.login(host='https://wandb.lnrweb.com', key='local-40a577c9eccab708f1aa81027f21a995596a5838') + # wandb.login(host='https://wandb.lnrweb.com', key='local-40a577c9eccab708f1aa81027f21a995596a5838') name = config['name'] script = config['script'] config = [key for key in sys.modules.keys() if key.startswith('configs.')][-1] diff --git a/scripts/cvae.py b/scripts/cvae.py index 87fff2d..a3d61d6 100644 --- a/scripts/cvae.py +++ b/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']: - pass + from configs.cape_00032_SL import config +else: + config_module = importlib.import_module(args['config']) + config = config_module.config + os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] if args['debug']: @@ -50,7 +52,7 @@ 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='/tmp', tags=config['tags']) + run = wandb.init(project='iccv', 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') @@ -457,7 +459,7 @@ def train(): g_balancing_coeff += 1 / 10000 if step % 10000 == 0: - ckpt_saver(model, step, optimizer) + ckpt_saver(model, step, optimizer, None) step += 1 @@ -469,8 +471,8 @@ if __name__ == '__main__': smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device']) skirt_modifier = '_skirt' if config['is_skirt'] else '' - dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory']) - dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory']) + dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape']) + dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape']) train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True, ) valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True) train_accumulator = Accumulator() @@ -478,12 +480,11 @@ if __name__ == '__main__': model = CVAE(config) - optimizer = optim.Adam(model.parameters(), lr=config['lr']) _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3) g_balancing_coeff = config['starting_balancing_coeff'] - ckpt_saver = CheckpointSaver('checkpoints/iccv/' + config['name'], top_n=5) + ckpt_saver = CheckpointSaver('checkpoints/iccv/torch/' + config['name'], top_n=10) step = 0 if config['pretrained_model'] != '':