diff --git a/configs/iccvw/abl_all.py b/configs/iccvw/abl_all.py new file mode 100644 index 0000000..b9d9639 --- /dev/null +++ b/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' \ No newline at end of file diff --git a/models/cvae_model.py b/models/cvae_model.py index b81ee6a..788f31e 100644 --- a/models/cvae_model.py +++ b/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']) diff --git a/scripts/cvae.py b/scripts/cvae.py index c485b5e..2ee71e9 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']: - 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'] != '': diff --git a/utils/utils.py b/utils/utils.py index d1d247a..5b5fe02 100644 --- a/utils/utils.py +++ b/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 \ No newline at end of file + return batch + +def text_to_html(lines): + html = '' + for line in lines: + html += "
" + line + "
" + return html \ No newline at end of file