luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

pyramid ready for cape

Browse Source
This commit is contained in:
luni
2023-05-22 12:44:43 +02:00
parent 6b81c0d1a6
commit 60844f8011
5 changed files with 321 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
configs/pyr_64_cape_m.py
+296
View File
@@ -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'
data/dataset.py
+2 -2
View File
@@ -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()
losses/losses.py
+12 -5
View File
@@ -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):
scripts/cvae_pyramid.py
+1 -1
View File
@@ -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']:
utils/postprocess.py
+10 -1
View File
@@ -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