From 534fbfcde1976dbe16910f974d64ffcd7944b9bb Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Thu, 1 Jun 2023 08:41:53 +0000 Subject: [PATCH] replace curvature with normals --- configs/abl_norm.py | 192 +++++++++++++++++++++++++++++++++++++++++++ models/cvae_model.py | 18 ++-- run.py | 2 +- scripts/cvae.py | 22 ++--- 4 files changed, 214 insertions(+), 20 deletions(-) create mode 100644 configs/abl_norm.py diff --git a/configs/abl_norm.py b/configs/abl_norm.py new file mode 100644 index 0000000..31ca571 --- /dev/null +++ b/configs/abl_norm.py @@ -0,0 +1,192 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '0' +config['name'] = 'abl_torch_norm' +config['notes'] = '' +config['tags'] = ['iccv', 'abl_norm', 'torch'] +config['script'] = 'scripts/cvae.py' + +config['pretrained_model'] = '' + +config['use_skinning'] = True # and deshaping +config['remove_global_orientation'] = True +config['use_normals'] = 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_normals'], 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_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'] = False +config['is_cape'] = False +config['device'] = 'cuda' + +config['train_dir'] = '/data-local/127/train' +config['valid_dir'] = '/data-local/127/valid' +config['test_dir'] = '/data-net/dataset/test_ready' +config['train_frames'] = None +config['valid_frames'] = None + +config['static_model'] = 'checkpoints/iccv/static_torch_orig_2' +config['body_model'] = 'checkpoints/iccv/body_torch_1' +config['normals_model'] = 'checkpoints/iccv/torch/norm_torch_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['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']['is_label'] = True +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 a92cd2a..d71c651 100644 --- a/models/cvae_model.py +++ b/models/cvae_model.py @@ -27,16 +27,16 @@ class CVAE(nn.Module): if config['body_model'] != '': load_model(self.body_enc, config['body_model'] + '_enc') - if config['use_curvature']: - self.curv_enc : AEEncoder = AEEncoder(config['cond_ae_curv']['encoder_block_settings'], config['cond_ae_curv']['kernel_size'], config['cond_ae_curv']['activation'], device=config['device'], is_label=True, nr_of_labels=7) - if config['curvature_model'] != '': - load_model(self.curv_enc, config['curvature_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']) + if config['normals_model'] != '': + load_model(self.norm_enc, config['normals_model'] + '_enc') if config['freeze_encoders']: for param in self.static_enc.parameters(): param.requires_grad = False for param in self.body_enc.parameters(): param.requires_grad = False - if config['use_curvature']: - for param in self.curv_enc.parameters(): param.requires_grad = False + if config['use_normals']: + for param in self.norm_enc.parameters(): param.requires_grad = False self.encoder = VariationalEncoder(config['cvae']['encoder_block_settings'], config['cvae']['kernel_size'], config['cvae']['input_features'], device=config['device']) self.decoder : Decoder = Decoder(config['cvae']['decoder_block_settings'], kernel_size=config['cvae']['kernel_size'], device=config['device']) @@ -48,7 +48,7 @@ class CVAE(nn.Module): else: self.avg_garment = None - self.use_curvature = config['use_curvature'] + self.use_normals = config['use_normals'] self.use_offsets_over_body = not config['use_skinning'] if self.use_offsets_over_body: @@ -98,8 +98,8 @@ class CVAE(nn.Module): def encode_conds(self, uv_template, uv_body_posed, uv_curv): z_body = self.body_enc(uv_body_posed) z_static = self.static_enc(uv_template) - if self.use_curvature: - z_curv = self.curv_enc(uv_curv) + if self.use_normals: + z_curv = self.norm_enc(uv_curv) z_cond = torch.concat([z_static, z_body, z_curv], dim=1) else: z_cond = torch.concat([z_static, z_body], dim=1) diff --git a/run.py b/run.py index e1d7d52..88c7e24 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.abl_norm import config import argparse # parser = argparse.ArgumentParser(description='') diff --git a/scripts/cvae.py b/scripts/cvae.py index 87fff2d..46c7261 100644 --- a/scripts/cvae.py +++ b/scripts/cvae.py @@ -13,7 +13,8 @@ args = vars(parser.parse_args()) config_module = importlib.import_module(args['config']) config = config_module.config if args['debug']: - pass + from configs.abl_norm import config + os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] if args['debug']: @@ -83,7 +84,7 @@ def train_step(batch): mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed'] model.train() - pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']]) + pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']]) mesh_pred = uv2mesh(pred) if config['use_skinning']: @@ -139,7 +140,7 @@ def valid_step(batch): model.eval() with torch.inference_mode(): - pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']]) + pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']]) mesh_pred = uv2mesh(pred) if config['use_skinning']: @@ -188,15 +189,16 @@ def plot_everything(batch: dict, type: str): model.eval() with torch.inference_mode(): - pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']]) + pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']]) mesh_pred = uv2mesh(pred).detach().cpu().numpy() orig_img = torch.permute(uv_orig, [0,2,3,1]) uv_static_img = torch.permute(batch['uv_static'], [0,2,3,1]) uv_body_posed_img = torch.permute(batch['uv_body_posed'], [0,2,3,1]) + uv_normals_img = (torch.permute(batch['uv_normals'], [0,2,3,1]) + 1.1) / 2.2 pred_img = torch.permute(pred, [0,2,3,1]) - img = torch.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2) + img = torch.concat((uv_static_img, uv_body_posed_img, uv_normals_img, orig_img, pred_img), axis=2) imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)] F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy") @@ -262,10 +264,10 @@ def sample(folder): count = 0 for batch in dataset_train: batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) - uv_curvature = tf.one_hot(batch['uv_curvature'], 7) + uv_normals = tf.one_hot(batch['uv_normals'], 7) nr_of_samples=2 - pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples) + pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples) pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:] pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1]) @@ -281,7 +283,7 @@ def sample(folder): mask = batch['mesh_mask'].numpy().astype(bool)[0] _, faces = _maskFaces(np.ones((config['smpl_res'], 3)), F, mask) - pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True) + pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True) mesh_pred_zero = uv2mesh(pred_zero) mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy() @@ -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 @@ -491,6 +493,6 @@ if __name__ == '__main__': 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() \ No newline at end of file