From 60844f8011b09d42509664f508dd883b29564277 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Mon, 22 May 2023 12:44:43 +0200 Subject: [PATCH] pyramid ready for cape --- configs/pyr_64_cape_m.py | 296 +++++++++++++++++++++++++++++++++++++++ data/dataset.py | 4 +- losses/losses.py | 17 ++- scripts/cvae_pyramid.py | 2 +- utils/postprocess.py | 11 +- 5 files changed, 321 insertions(+), 9 deletions(-) create mode 100644 configs/pyr_64_cape_m.py diff --git a/configs/pyr_64_cape_m.py b/configs/pyr_64_cape_m.py new file mode 100644 index 0000000..acf949a --- /dev/null +++ b/configs/pyr_64_cape_m.py @@ -0,0 +1,296 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '5' +config['name'] = 'pyr_torch_64_cape_m' +config['notes'] = '' +config['tags'] = ['pyr', 'iccv', 'torch'] +config['script'] = 'scripts/cvae_pyramid.py' + +config['batch_size'] = 4 +config['is_skirt'] = False +config['is_cape'] = True +config['device'] = 'cuda' + + +config['use_skinning'] = True # and deshaping +config['remove_global_orientation'] = True +config['use_extra_losses'] = True # edge, collision, 3d mesh +config['use_offsets'] = False +config['use_norm_loss'] = True +config['use_curvature'] = True + +config['pretrained_model'] = '' +config['pretrained_model_iteration'] = '' + +config['pyramid'] = dict() +config['pyramid']['resolutions'] = [64] +config['pyramid']['pretrained_pyramid'] = '' +config['pyramid']['pretrained_pyramid_iteration'] = '' +config['pyramid']['final_upscale_enhancer'] = '' + +config['train_dir'] = '/data-local/127/cape_processed/train/male' +config['valid_dir'] = '/data-local/127/cape_processed/test/male' +config['test_dir'] = '/data-local/127/cape_processed/test/male' +config['train_frames'] = None +config['valid_frames'] = None + +config['dataset_in_memory'] = False +config['resolution'] = None + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['lr'] = 0.0001 +config['starting_balancing_coeff'] = 0.001 +config['balancing_activation_limit'] = 0.015 +config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision +config['avg_garment_path'] = '' +config['bottleneck_scale'] = 2 + +config['signature']: set = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'mesh_unposed', + 'uv_unposed', + 'uv_curvature', + 'uv_static', + 'mesh_mask', + 'uv_mask', + } + +config['pyramid']['is_label'] = False +# config['pyramid']['add_noise'] = False +# config['pyramid']['noise_amount'] = 0.0 + +config['pyramid']['loss_weights'] = [1,1,1,1,1] + +for res in config['pyramid']['resolutions']: + if res != 512: + res_modifier = '_' + str(res) + config['signature'].add('uv_body_posed' + res_modifier) + config['signature'].add('uv_unposed' + res_modifier) + config['signature'].add('uv_curvature' + res_modifier) + config['signature'].add('uv_static' + res_modifier) + config['signature'].add('uv_mask' + res_modifier) + + +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['freeze_encoders'] = True + + +config['pyramid']['levels'] = [] +for res in config['pyramid']['resolutions']: + if res == 512: + level = dict() + level['resolution'] = 512 + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['kernel_size'] = 7 + level['upscale_enhancer'] = '' + level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3 + + block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 3), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 256: + level = dict() + level['resolution'] = 256 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccv/torch/iccv_256_cape_m' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, 128, 3), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 3), + CNBlockConfig(128, 128, 9), + CNBlockConfig(128, 65, 3), + CNBlockConfig(64, 3, 3), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 128: + level = dict() + level['resolution'] = 128 + level['kernel_size'] = 5 + level['upscale_enhancer'] = 'checkpoints/iccv/torch/iccv_128_cape_m' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 128, 9), + CNBlockConfig(128, 128, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 64: + level = dict() + level['resolution'] = 64 + level['kernel_size'] = 3 + level['upscale_enhancer'] = 'checkpoints/iccv/torch/iccv_64_cape_m' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3 + + block_setting = [ + CNBlockConfig(64, level['bottleneck_size'], 9), + CNBlockConfig(level['bottleneck_size'], None, 9), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 64, 9), + CNBlockConfig(64, 3, 9), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + elif res == 32: + level = dict() + level['resolution'] = 32 + level['kernel_size'] = 3 + level['upscale_enhancer'] = 'checkpoints/iccv/torch/iccv_32_cape_m' + level['upscale_enhancer_freeze'] = True + level['bottleneck_size'] = 64 * config['bottleneck_scale'] + level['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3 + + block_setting = [ + CNBlockConfig(level['bottleneck_size'], None, 16), + ] + + block_setting_reverse = [ + CNBlockConfig(level['bottleneck_size'] * 2, 3, 16), + ] + + level['encoder_block_settings'] = block_setting + level['decoder_block_settings'] = block_setting_reverse + config['pyramid']['levels'].append(level) + else: + raise NotImplementedError + + +####################### 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/data/dataset.py b/data/dataset.py index e64dde9..f85c5ad 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -1,7 +1,7 @@ import os import sys -os.environ["CUDA_VISIBLE_DEVICES"] = '6' +# os.environ["CUDA_VISIBLE_DEVICES"] = '6' import math import time @@ -69,7 +69,7 @@ class StaticCache(): uv_mask = torch.tensor(np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32), device=device) self.uv_masks[512] = uv_mask.to(torch.float32).cpu().numpy() for res in [32, 64, 128, 256]: - down = Resize(size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST)(uv_mask[None,:,:]) + down = Resize(size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST_EXACT)(uv_mask[None,:,:]) down = torch.squeeze(down).to(torch.float32) self.uv_masks[res] = down.cpu().numpy() diff --git a/losses/losses.py b/losses/losses.py index 9fdc46c..4f47a09 100644 --- a/losses/losses.py +++ b/losses/losses.py @@ -15,6 +15,9 @@ from torch_cluster import knn # from hbertiche/PBNS """ COLLISION FUNCTIONS START """ +smpl_to_keep_orig = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1) +smpl_to_keep_orig = smpl_to_keep_orig[smpl_to_keep_orig < 6890] + to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1) F_coll = torch.tensor(np.load("smpl/smpl_faces.npy")[to_keep,:], device='cuda') F_coll_skirt = torch.tensor(np.load("smpl/smpl_faces_skirt.npy")[to_keep,:], device='cuda') @@ -22,23 +25,29 @@ F_coll_orig = torch.tensor(np.load("smpl/smpl_faces_orig.npy")[:6890,:], device= F_norm = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda') F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda') F_norm_orig = torch.tensor(np.tile(np.load("smpl/smpl_faces_orig.npy")[:6890,:][None,:,:], [4,1,1]), device='cuda') +F_norm_orig = F_norm_orig[:,smpl_to_keep_orig,:] def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor): # flatten V and B along the batch dimension, create a batch_x and batch_y variable to restore the batch dimension after the nearest neighbour search using pytorch_geomtry knn function batch_size = V.shape[0] + vertices_count = V.shape[1] batch_x = torch.arange(batch_size, device=V.get_device()).repeat_interleave(V.shape[1]) batch_y = torch.arange(batch_size, device=B.get_device()).repeat_interleave(B.shape[1]) V = V.reshape(-1, V.shape[-1]) B = B.reshape(-1, B.shape[-1]) idx = knn(V, B.to(V.dtype), k=1, batch_x=batch_x, batch_y=batch_y)[1] - idx = (idx - batch_x * 14475).reshape(batch_size, 14475) + idx = (idx - batch_x * vertices_count).reshape(batch_size, vertices_count) return idx # @tf.function def collision_loss(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: float = .004): - mesh = Meshes(B, F_norm) + if B.shape[1] == 14475: + F_norm_local = F_norm + elif B.shape[1] == 3638: + F_norm_local = F_norm_orig + mesh = Meshes(B, F_norm_local) B_vn = torch.stack(mesh.verts_normals_list()) # compute correspondences @@ -64,9 +73,7 @@ edges: torch.Tensor = torch.tensor(np.stack(np.nonzero(laplacian_loss.todense()) laplacian_loss_skirt = sio.loadmat("custom_utils/ssmpl_skirt0.mat")['L'].tocoo() edges_skirt: torch.Tensor = torch.tensor(np.stack(np.nonzero(laplacian_loss.todense())).transpose(), device='cuda') -smpl_to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1) -smpl_to_keep = smpl_to_keep[smpl_to_keep < 6890] -laplacian_loss_orig = laplacian_loss.todense()[:6890,:6890][smpl_to_keep,:][:,smpl_to_keep] +laplacian_loss_orig = laplacian_loss.todense()[:6890,:6890][smpl_to_keep_orig,:][:,smpl_to_keep_orig] edges_cape = torch.tensor(np.stack(np.nonzero(laplacian_loss_orig)).transpose(), device='cuda') def get_edge_mask(edges: torch.Tensor, masks: torch.Tensor): diff --git a/scripts/cvae_pyramid.py b/scripts/cvae_pyramid.py index deb3e2c..56bd3c9 100644 --- a/scripts/cvae_pyramid.py +++ b/scripts/cvae_pyramid.py @@ -13,7 +13,7 @@ args = vars(parser.parse_args()) config_module = importlib.import_module(args['config']) config = config_module.config if args['debug']: - from configs.pyr_32 import config + from configs.pyr_64_cape_m import config os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] if args['debug']: diff --git a/utils/postprocess.py b/utils/postprocess.py index 53e3c15..c2d2272 100644 --- a/utils/postprocess.py +++ b/utils/postprocess.py @@ -23,10 +23,19 @@ from losses.losses import _nearest_neighbour # F = np.load("smpl/smpl_faces.npy") to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1) F_fix_coll = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda') +smpl_to_keep_orig_post = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1) +smpl_to_keep_orig_post = smpl_to_keep_orig_post[smpl_to_keep_orig_post < 6890] +F_fix_coll_orig = torch.tensor(np.tile(np.load("smpl/smpl_faces_orig.npy")[:6890,:][None,:,:], [4,1,1]), device='cuda')[:,smpl_to_keep_orig_post,:] def fix_collision(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr=.004): - mesh = Meshes(B, F_fix_coll) + if V.shape[1] == 14475: + F_fix_coll_local = F_fix_coll + elif V.shape[1] == 3638: + F_fix_coll_local = F_fix_coll_orig + else: + raise ValueError("V must have 14475 or 3638 vertices") + mesh = Meshes(B, F_fix_coll_local) B_vn = torch.stack(mesh.verts_normals_list()) # compute correspondences