Merge remote-tracking branch 'origin/main' into main
This commit is contained in:
@@ -2,7 +2,7 @@ from utils.convnext_utils import CNBlockConfig
|
||||
|
||||
config = dict()
|
||||
|
||||
config['gpu'] = '4'
|
||||
config['gpu'] = '0'
|
||||
config['name'] = 'body_torch_cape_m'
|
||||
config['notes'] = ' '
|
||||
config['tags'] = ['torch', 'body']
|
||||
@@ -11,9 +11,9 @@ config['script'] = 'scripts/ae_body_train.py'
|
||||
config['pretrained_model'] = ''
|
||||
|
||||
config['batch_size'] = 4
|
||||
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_dir'] = '/data-local/128/cape_processed/train/male'
|
||||
config['valid_dir'] = '/data-local/128/cape_processed/test/male'
|
||||
config['test_dir'] = '/data-local/128/cape_processed/test/male'
|
||||
config['train_frames'] = None
|
||||
config['valid_frames'] = None
|
||||
|
||||
@@ -0,0 +1,78 @@
|
||||
from utils.convnext_utils import CNBlockConfig
|
||||
|
||||
config = dict()
|
||||
|
||||
config['gpu'] = '2'
|
||||
config['name'] = 'body_iccvw_1'
|
||||
config['notes'] = ' '
|
||||
config['tags'] = ['torch', 'body']
|
||||
config['script'] = 'scripts/ae_body_train.py'
|
||||
|
||||
config['pretrained_model'] = ''
|
||||
|
||||
config['batch_size'] = 4
|
||||
config['train_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['valid_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['test_dir'] = '/data-net/dataset/test_ready'
|
||||
|
||||
config['resolution'] = 512
|
||||
config['is_skirt'] = False
|
||||
config['dataset_in_memory'] = False
|
||||
config['is_cape'] = False
|
||||
config['remove_global_orientation'] = True
|
||||
config['device'] = 'cuda'
|
||||
|
||||
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['smpl_res'] = 14475
|
||||
if config['is_cape']:
|
||||
config['smpl_res'] = 3638
|
||||
|
||||
config['latent_dim'] = 128
|
||||
|
||||
config['bottleneck_scale'] = 2
|
||||
|
||||
config['cond_ae'] = dict()
|
||||
config['cond_ae']['device'] = config['device']
|
||||
config['cond_ae']['is_label'] = False
|
||||
config['cond_ae']['latent_dim'] = config['latent_dim']
|
||||
config['cond_ae']['kernel_size'] = 7
|
||||
config['cond_ae']['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']['bottleneck_size'], 9),
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
|
||||
]
|
||||
|
||||
block_setting_reverse = [
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
|
||||
CNBlockConfig(128, 64, 9),
|
||||
CNBlockConfig(64, 32, 9),
|
||||
CNBlockConfig(32, 16, 3),
|
||||
CNBlockConfig(16, 3, 3),
|
||||
# CNBlockConfig(3, 3, 3),
|
||||
]
|
||||
|
||||
config['cond_ae']['encoder_block_settings'] = block_setting
|
||||
config['cond_ae']['decoder_block_settings'] = block_setting_reverse
|
||||
config['cond_ae']['activation'] = 'tanh'
|
||||
|
||||
|
||||
config['lr'] = 0.0001
|
||||
config['cond_ae']['upscale_output'] = False
|
||||
|
||||
config['signature'] = { 'id',
|
||||
'frame_nr',
|
||||
'pose_shape',
|
||||
'uv_body_posed',
|
||||
'mesh_mask',
|
||||
'mesh_body_posed',
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
from utils.convnext_utils import CNBlockConfig
|
||||
|
||||
config = dict()
|
||||
|
||||
config['gpu'] = '5'
|
||||
config['name'] = 'body_iccvw_1_10p'
|
||||
config['notes'] = ' '
|
||||
config['tags'] = ['torch', 'body']
|
||||
config['script'] = 'scripts/ae_body_train.py'
|
||||
|
||||
config['pretrained_model'] = ''
|
||||
|
||||
config['batch_size'] = 4
|
||||
config['train_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['valid_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['test_dir'] = '/data-net/dataset/test_ready'
|
||||
|
||||
config['resolution'] = 512
|
||||
config['is_skirt'] = False
|
||||
config['dataset_in_memory'] = False
|
||||
config['is_cape'] = False
|
||||
config['remove_global_orientation'] = True
|
||||
config['device'] = 'cuda'
|
||||
|
||||
if config['is_skirt']:
|
||||
config['train_frames'] = 'frames_train_skirt_10p.txt'
|
||||
config['valid_frames'] = 'frames_valid_skirt.txt'
|
||||
else:
|
||||
config['train_frames'] = 'frames_train_10p.txt'
|
||||
config['valid_frames'] = 'frames_valid.txt'
|
||||
|
||||
config['smpl_res'] = 14475
|
||||
if config['is_cape']:
|
||||
config['smpl_res'] = 3638
|
||||
|
||||
config['latent_dim'] = 128
|
||||
|
||||
config['bottleneck_scale'] = 2
|
||||
|
||||
config['cond_ae'] = dict()
|
||||
config['cond_ae']['device'] = config['device']
|
||||
config['cond_ae']['is_label'] = False
|
||||
config['cond_ae']['latent_dim'] = config['latent_dim']
|
||||
config['cond_ae']['kernel_size'] = 7
|
||||
config['cond_ae']['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']['bottleneck_size'], 9),
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
|
||||
]
|
||||
|
||||
block_setting_reverse = [
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
|
||||
CNBlockConfig(128, 64, 9),
|
||||
CNBlockConfig(64, 32, 9),
|
||||
CNBlockConfig(32, 16, 3),
|
||||
CNBlockConfig(16, 3, 3),
|
||||
# CNBlockConfig(3, 3, 3),
|
||||
]
|
||||
|
||||
config['cond_ae']['encoder_block_settings'] = block_setting
|
||||
config['cond_ae']['decoder_block_settings'] = block_setting_reverse
|
||||
config['cond_ae']['activation'] = 'tanh'
|
||||
|
||||
|
||||
config['lr'] = 0.0001
|
||||
config['cond_ae']['upscale_output'] = False
|
||||
|
||||
config['signature'] = { 'id',
|
||||
'frame_nr',
|
||||
'pose_shape',
|
||||
'uv_body_posed',
|
||||
'mesh_mask',
|
||||
'mesh_body_posed',
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
from utils.convnext_utils import CNBlockConfig
|
||||
|
||||
config = dict()
|
||||
|
||||
config['gpu'] = '5'
|
||||
config['name'] = 'body_iccvw_1_25p'
|
||||
config['notes'] = ' '
|
||||
config['tags'] = ['torch', 'body']
|
||||
config['script'] = 'scripts/ae_body_train.py'
|
||||
|
||||
config['pretrained_model'] = ''
|
||||
|
||||
config['batch_size'] = 4
|
||||
config['train_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['valid_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['test_dir'] = '/data-net/dataset/test_ready'
|
||||
|
||||
config['resolution'] = 512
|
||||
config['is_skirt'] = False
|
||||
config['dataset_in_memory'] = False
|
||||
config['is_cape'] = False
|
||||
config['remove_global_orientation'] = True
|
||||
config['device'] = 'cuda'
|
||||
|
||||
if config['is_skirt']:
|
||||
config['train_frames'] = 'frames_train_skirt_25p.txt'
|
||||
config['valid_frames'] = 'frames_valid_skirt.txt'
|
||||
else:
|
||||
config['train_frames'] = 'frames_train_25p.txt'
|
||||
config['valid_frames'] = 'frames_valid.txt'
|
||||
|
||||
config['smpl_res'] = 14475
|
||||
if config['is_cape']:
|
||||
config['smpl_res'] = 3638
|
||||
|
||||
config['latent_dim'] = 128
|
||||
|
||||
config['bottleneck_scale'] = 2
|
||||
|
||||
config['cond_ae'] = dict()
|
||||
config['cond_ae']['device'] = config['device']
|
||||
config['cond_ae']['is_label'] = False
|
||||
config['cond_ae']['latent_dim'] = config['latent_dim']
|
||||
config['cond_ae']['kernel_size'] = 7
|
||||
config['cond_ae']['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']['bottleneck_size'], 9),
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
|
||||
]
|
||||
|
||||
block_setting_reverse = [
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
|
||||
CNBlockConfig(128, 64, 9),
|
||||
CNBlockConfig(64, 32, 9),
|
||||
CNBlockConfig(32, 16, 3),
|
||||
CNBlockConfig(16, 3, 3),
|
||||
# CNBlockConfig(3, 3, 3),
|
||||
]
|
||||
|
||||
config['cond_ae']['encoder_block_settings'] = block_setting
|
||||
config['cond_ae']['decoder_block_settings'] = block_setting_reverse
|
||||
config['cond_ae']['activation'] = 'tanh'
|
||||
|
||||
|
||||
config['lr'] = 0.0001
|
||||
config['cond_ae']['upscale_output'] = False
|
||||
|
||||
config['signature'] = { 'id',
|
||||
'frame_nr',
|
||||
'pose_shape',
|
||||
'uv_body_posed',
|
||||
'mesh_mask',
|
||||
'mesh_body_posed',
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
from utils.convnext_utils import CNBlockConfig
|
||||
|
||||
config = dict()
|
||||
|
||||
config['gpu'] = '5'
|
||||
config['name'] = 'body_iccvw_1_50p'
|
||||
config['notes'] = ' '
|
||||
config['tags'] = ['torch', 'body']
|
||||
config['script'] = 'scripts/ae_body_train.py'
|
||||
|
||||
config['pretrained_model'] = ''
|
||||
|
||||
config['batch_size'] = 4
|
||||
config['train_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['valid_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['test_dir'] = '/data-net/dataset/test_ready'
|
||||
|
||||
config['resolution'] = 512
|
||||
config['is_skirt'] = False
|
||||
config['dataset_in_memory'] = False
|
||||
config['is_cape'] = False
|
||||
config['remove_global_orientation'] = True
|
||||
config['device'] = 'cuda'
|
||||
|
||||
if config['is_skirt']:
|
||||
config['train_frames'] = 'frames_train_skirt_50p.txt'
|
||||
config['valid_frames'] = 'frames_valid_skirt.txt'
|
||||
else:
|
||||
config['train_frames'] = 'frames_train_50p.txt'
|
||||
config['valid_frames'] = 'frames_valid.txt'
|
||||
|
||||
config['smpl_res'] = 14475
|
||||
if config['is_cape']:
|
||||
config['smpl_res'] = 3638
|
||||
|
||||
config['latent_dim'] = 128
|
||||
|
||||
config['bottleneck_scale'] = 2
|
||||
|
||||
config['cond_ae'] = dict()
|
||||
config['cond_ae']['device'] = config['device']
|
||||
config['cond_ae']['is_label'] = False
|
||||
config['cond_ae']['latent_dim'] = config['latent_dim']
|
||||
config['cond_ae']['kernel_size'] = 7
|
||||
config['cond_ae']['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']['bottleneck_size'], 9),
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
|
||||
]
|
||||
|
||||
block_setting_reverse = [
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
|
||||
CNBlockConfig(128, 64, 9),
|
||||
CNBlockConfig(64, 32, 9),
|
||||
CNBlockConfig(32, 16, 3),
|
||||
CNBlockConfig(16, 3, 3),
|
||||
# CNBlockConfig(3, 3, 3),
|
||||
]
|
||||
|
||||
config['cond_ae']['encoder_block_settings'] = block_setting
|
||||
config['cond_ae']['decoder_block_settings'] = block_setting_reverse
|
||||
config['cond_ae']['activation'] = 'tanh'
|
||||
|
||||
|
||||
config['lr'] = 0.0001
|
||||
config['cond_ae']['upscale_output'] = False
|
||||
|
||||
config['signature'] = { 'id',
|
||||
'frame_nr',
|
||||
'pose_shape',
|
||||
'uv_body_posed',
|
||||
'mesh_mask',
|
||||
'mesh_body_posed',
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
from utils.convnext_utils import CNBlockConfig
|
||||
|
||||
config = dict()
|
||||
|
||||
config['gpu'] = '1'
|
||||
config['name'] = 'body_iccvw_1_skirt'
|
||||
config['notes'] = ' '
|
||||
config['tags'] = ['torch', 'body']
|
||||
config['script'] = 'scripts/ae_body_train.py'
|
||||
|
||||
config['pretrained_model'] = ''
|
||||
|
||||
config['batch_size'] = 4
|
||||
config['train_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['valid_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['test_dir'] = '/data-net/dataset/test_ready'
|
||||
|
||||
config['resolution'] = 512
|
||||
config['is_skirt'] = True
|
||||
config['dataset_in_memory'] = False
|
||||
config['is_cape'] = False
|
||||
config['remove_global_orientation'] = True
|
||||
config['device'] = 'cuda'
|
||||
|
||||
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['smpl_res'] = 14475
|
||||
if config['is_cape']:
|
||||
config['smpl_res'] = 3638
|
||||
|
||||
config['latent_dim'] = 128
|
||||
|
||||
config['bottleneck_scale'] = 2
|
||||
|
||||
config['cond_ae'] = dict()
|
||||
config['cond_ae']['device'] = config['device']
|
||||
config['cond_ae']['is_label'] = False
|
||||
config['cond_ae']['latent_dim'] = config['latent_dim']
|
||||
config['cond_ae']['kernel_size'] = 7
|
||||
config['cond_ae']['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']['bottleneck_size'], 9),
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
|
||||
]
|
||||
|
||||
block_setting_reverse = [
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
|
||||
CNBlockConfig(128, 64, 9),
|
||||
CNBlockConfig(64, 32, 9),
|
||||
CNBlockConfig(32, 16, 3),
|
||||
CNBlockConfig(16, 3, 3),
|
||||
# CNBlockConfig(3, 3, 3),
|
||||
]
|
||||
|
||||
config['cond_ae']['encoder_block_settings'] = block_setting
|
||||
config['cond_ae']['decoder_block_settings'] = block_setting_reverse
|
||||
config['cond_ae']['activation'] = 'tanh'
|
||||
|
||||
|
||||
config['lr'] = 0.0001
|
||||
config['cond_ae']['upscale_output'] = False
|
||||
|
||||
config['signature'] = { 'id',
|
||||
'frame_nr',
|
||||
'pose_shape',
|
||||
'uv_body_posed',
|
||||
'mesh_mask',
|
||||
'mesh_body_posed',
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
from utils.convnext_utils import CNBlockConfig
|
||||
|
||||
config = dict()
|
||||
|
||||
config['gpu'] = '0'
|
||||
config['name'] = 'body_iccvw_cape_f'
|
||||
config['notes'] = ' '
|
||||
config['tags'] = ['torch', 'body']
|
||||
config['script'] = 'scripts/ae_body_train.py'
|
||||
|
||||
config['pretrained_model'] = ''
|
||||
|
||||
config['batch_size'] = 4
|
||||
config['train_dir'] = '/data-local/128/cape_processed/train/female'
|
||||
config['valid_dir'] = '/data-local/128/cape_processed/test/female'
|
||||
config['test_dir'] = '/data-local/128/cape_processed/test/female'
|
||||
config['train_frames'] = None
|
||||
config['valid_frames'] = None
|
||||
|
||||
config['resolution'] = 512
|
||||
config['is_skirt'] = False
|
||||
config['dataset_in_memory'] = False
|
||||
config['is_cape'] = True
|
||||
config['remove_global_orientation'] = True
|
||||
config['device'] = 'cuda'
|
||||
|
||||
config['smpl_res'] = 14475
|
||||
if config['is_cape']:
|
||||
config['smpl_res'] = 3638
|
||||
|
||||
config['latent_dim'] = 128
|
||||
|
||||
config['bottleneck_scale'] = 2
|
||||
|
||||
config['cond_ae'] = dict()
|
||||
config['cond_ae']['device'] = config['device']
|
||||
config['cond_ae']['is_label'] = False
|
||||
config['cond_ae']['latent_dim'] = config['latent_dim']
|
||||
config['cond_ae']['kernel_size'] = 7
|
||||
config['cond_ae']['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']['bottleneck_size'], 9),
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
|
||||
]
|
||||
|
||||
block_setting_reverse = [
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
|
||||
CNBlockConfig(128, 64, 9),
|
||||
CNBlockConfig(64, 32, 9),
|
||||
CNBlockConfig(32, 16, 3),
|
||||
CNBlockConfig(16, 3, 3),
|
||||
# CNBlockConfig(3, 3, 3),
|
||||
]
|
||||
|
||||
config['cond_ae']['encoder_block_settings'] = block_setting
|
||||
config['cond_ae']['decoder_block_settings'] = block_setting_reverse
|
||||
config['cond_ae']['activation'] = 'tanh'
|
||||
|
||||
|
||||
config['lr'] = 0.0001
|
||||
config['cond_ae']['upscale_output'] = False
|
||||
|
||||
config['signature'] = { 'id',
|
||||
'frame_nr',
|
||||
'pose_shape',
|
||||
'uv_body_posed',
|
||||
'mesh_mask',
|
||||
'mesh_body_posed',
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
from utils.convnext_utils import CNBlockConfig
|
||||
|
||||
config = dict()
|
||||
|
||||
config['gpu'] = '0'
|
||||
config['name'] = 'body_torch_cape_m'
|
||||
config['notes'] = ' '
|
||||
config['tags'] = ['torch', 'body']
|
||||
config['script'] = 'scripts/ae_body_train.py'
|
||||
|
||||
config['pretrained_model'] = ''
|
||||
|
||||
config['batch_size'] = 4
|
||||
config['train_dir'] = '/data-local/128/cape_processed/train/male'
|
||||
config['valid_dir'] = '/data-local/128/cape_processed/test/male'
|
||||
config['test_dir'] = '/data-local/128/cape_processed/test/male'
|
||||
config['train_frames'] = None
|
||||
config['valid_frames'] = None
|
||||
|
||||
config['resolution'] = 512
|
||||
config['is_skirt'] = False
|
||||
config['dataset_in_memory'] = False
|
||||
config['is_cape'] = True
|
||||
config['remove_global_orientation'] = True
|
||||
config['device'] = 'cuda'
|
||||
|
||||
config['smpl_res'] = 14475
|
||||
if config['is_cape']:
|
||||
config['smpl_res'] = 3638
|
||||
|
||||
config['latent_dim'] = 128
|
||||
|
||||
config['bottleneck_scale'] = 2
|
||||
|
||||
config['cond_ae'] = dict()
|
||||
config['cond_ae']['device'] = config['device']
|
||||
config['cond_ae']['is_label'] = False
|
||||
config['cond_ae']['latent_dim'] = config['latent_dim']
|
||||
config['cond_ae']['kernel_size'] = 7
|
||||
config['cond_ae']['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']['bottleneck_size'], 9),
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
|
||||
]
|
||||
|
||||
block_setting_reverse = [
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
|
||||
CNBlockConfig(128, 64, 9),
|
||||
CNBlockConfig(64, 32, 9),
|
||||
CNBlockConfig(32, 16, 3),
|
||||
CNBlockConfig(16, 3, 3),
|
||||
# CNBlockConfig(3, 3, 3),
|
||||
]
|
||||
|
||||
config['cond_ae']['encoder_block_settings'] = block_setting
|
||||
config['cond_ae']['decoder_block_settings'] = block_setting_reverse
|
||||
config['cond_ae']['activation'] = 'tanh'
|
||||
|
||||
|
||||
config['lr'] = 0.0001
|
||||
config['cond_ae']['upscale_output'] = False
|
||||
|
||||
config['signature'] = { 'id',
|
||||
'frame_nr',
|
||||
'pose_shape',
|
||||
'uv_body_posed',
|
||||
'mesh_mask',
|
||||
'mesh_body_posed',
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
from utils.convnext_utils import CNBlockConfig
|
||||
|
||||
config = dict()
|
||||
|
||||
config['gpu'] = '1'
|
||||
config['name'] = 'norm_iccvw_1'
|
||||
config['notes'] = ' '
|
||||
config['tags'] = ['torch', 'norm']
|
||||
config['script'] = 'scripts/ae_norm_train.py'
|
||||
|
||||
config['pretrained_model'] = ''
|
||||
|
||||
config['batch_size'] = 4
|
||||
config['train_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['valid_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['test_dir'] = '/data-net/dataset/test_ready'
|
||||
|
||||
config['resolution'] = 512
|
||||
config['is_skirt'] = False
|
||||
config['dataset_in_memory'] = False
|
||||
config['is_cape'] = False
|
||||
config['remove_global_orientation'] = True
|
||||
config['device'] = 'cuda'
|
||||
|
||||
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['smpl_res'] = 14475
|
||||
if config['is_cape']:
|
||||
config['smpl_res'] = 3638
|
||||
|
||||
config['latent_dim'] = 128
|
||||
|
||||
config['bottleneck_scale'] = 2
|
||||
|
||||
config['cond_ae'] = dict()
|
||||
config['cond_ae']['device'] = config['device']
|
||||
config['cond_ae']['is_label'] = False
|
||||
config['cond_ae']['latent_dim'] = config['latent_dim']
|
||||
config['cond_ae']['kernel_size'] = 7
|
||||
config['cond_ae']['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']['bottleneck_size'], 9),
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
|
||||
]
|
||||
|
||||
block_setting_reverse = [
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
|
||||
CNBlockConfig(128, 64, 9),
|
||||
CNBlockConfig(64, 32, 9),
|
||||
CNBlockConfig(32, 16, 3),
|
||||
CNBlockConfig(16, 3, 3),
|
||||
# CNBlockConfig(3, 3, 3),
|
||||
]
|
||||
|
||||
config['cond_ae']['encoder_block_settings'] = block_setting
|
||||
config['cond_ae']['decoder_block_settings'] = block_setting_reverse
|
||||
config['cond_ae']['activation'] = 'tanh'
|
||||
|
||||
|
||||
config['lr'] = 0.0001
|
||||
config['cond_ae']['upscale_output'] = False
|
||||
|
||||
config['signature'] = { 'id',
|
||||
'frame_nr',
|
||||
'uv_mask',
|
||||
'uv_normals',
|
||||
'mesh_mask',
|
||||
'mesh_normals',
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
from utils.convnext_utils import CNBlockConfig
|
||||
|
||||
config = dict()
|
||||
|
||||
config['gpu'] = '4'
|
||||
config['name'] = 'norm_iccvw_1_skirt'
|
||||
config['notes'] = ' '
|
||||
config['tags'] = ['torch', 'norm']
|
||||
config['script'] = 'scripts/ae_norm_train.py'
|
||||
|
||||
config['pretrained_model'] = ''
|
||||
|
||||
config['batch_size'] = 4
|
||||
config['train_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['valid_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['test_dir'] = '/data-net/dataset/test_ready'
|
||||
|
||||
config['resolution'] = 512
|
||||
config['is_skirt'] = True
|
||||
config['dataset_in_memory'] = False
|
||||
config['is_cape'] = False
|
||||
config['remove_global_orientation'] = True
|
||||
config['device'] = 'cuda'
|
||||
|
||||
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['smpl_res'] = 14475
|
||||
if config['is_cape']:
|
||||
config['smpl_res'] = 3638
|
||||
|
||||
config['latent_dim'] = 128
|
||||
|
||||
config['bottleneck_scale'] = 2
|
||||
|
||||
config['cond_ae'] = dict()
|
||||
config['cond_ae']['device'] = config['device']
|
||||
config['cond_ae']['is_label'] = False
|
||||
config['cond_ae']['latent_dim'] = config['latent_dim']
|
||||
config['cond_ae']['kernel_size'] = 7
|
||||
config['cond_ae']['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']['bottleneck_size'], 9),
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
|
||||
]
|
||||
|
||||
block_setting_reverse = [
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
|
||||
CNBlockConfig(128, 64, 9),
|
||||
CNBlockConfig(64, 32, 9),
|
||||
CNBlockConfig(32, 16, 3),
|
||||
CNBlockConfig(16, 3, 3),
|
||||
# CNBlockConfig(3, 3, 3),
|
||||
]
|
||||
|
||||
config['cond_ae']['encoder_block_settings'] = block_setting
|
||||
config['cond_ae']['decoder_block_settings'] = block_setting_reverse
|
||||
config['cond_ae']['activation'] = 'tanh'
|
||||
|
||||
|
||||
config['lr'] = 0.0001
|
||||
config['cond_ae']['upscale_output'] = False
|
||||
|
||||
config['signature'] = { 'id',
|
||||
'frame_nr',
|
||||
'uv_mask',
|
||||
'uv_normals',
|
||||
'mesh_mask',
|
||||
'mesh_normals',
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
from utils.convnext_utils import CNBlockConfig
|
||||
|
||||
config = dict()
|
||||
|
||||
config['gpu'] = '6'
|
||||
config['name'] = 'norm_iccvw_cape_f'
|
||||
config['notes'] = ' '
|
||||
config['tags'] = ['torch', 'norm']
|
||||
config['script'] = 'scripts/ae_norm_train.py'
|
||||
|
||||
config['pretrained_model'] = ''
|
||||
|
||||
config['batch_size'] = 4
|
||||
config['train_dir'] = '/data-local/127_2/cape_processed/train/female'
|
||||
config['valid_dir'] = '/data-local/127_2/cape_processed/test/female'
|
||||
config['test_dir'] = '/data-local/127_2/cape_processed/test/female'
|
||||
config['train_frames'] = None
|
||||
config['valid_frames'] = None
|
||||
|
||||
config['resolution'] = 512
|
||||
config['is_skirt'] = False
|
||||
config['dataset_in_memory'] = False
|
||||
config['is_cape'] = True
|
||||
config['remove_global_orientation'] = True
|
||||
config['device'] = 'cuda'
|
||||
|
||||
config['smpl_res'] = 14475
|
||||
if config['is_cape']:
|
||||
config['smpl_res'] = 3638
|
||||
|
||||
config['latent_dim'] = 128
|
||||
|
||||
config['bottleneck_scale'] = 2
|
||||
|
||||
config['cond_ae'] = dict()
|
||||
config['cond_ae']['device'] = config['device']
|
||||
config['cond_ae']['is_label'] = False
|
||||
config['cond_ae']['latent_dim'] = config['latent_dim']
|
||||
config['cond_ae']['kernel_size'] = 7
|
||||
config['cond_ae']['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']['bottleneck_size'], 9),
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
|
||||
]
|
||||
|
||||
block_setting_reverse = [
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
|
||||
CNBlockConfig(128, 64, 9),
|
||||
CNBlockConfig(64, 32, 9),
|
||||
CNBlockConfig(32, 16, 3),
|
||||
CNBlockConfig(16, 3, 3),
|
||||
# CNBlockConfig(3, 3, 3),
|
||||
]
|
||||
|
||||
config['cond_ae']['encoder_block_settings'] = block_setting
|
||||
config['cond_ae']['decoder_block_settings'] = block_setting_reverse
|
||||
config['cond_ae']['activation'] = 'tanh'
|
||||
|
||||
|
||||
config['lr'] = 0.0001
|
||||
config['cond_ae']['upscale_output'] = False
|
||||
|
||||
config['signature'] = { 'id',
|
||||
'frame_nr',
|
||||
'uv_mask',
|
||||
'uv_normals',
|
||||
'mesh_mask',
|
||||
'mesh_normals',
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
from utils.convnext_utils import CNBlockConfig
|
||||
|
||||
config = dict()
|
||||
|
||||
config['gpu'] = '6'
|
||||
config['name'] = 'norm_iccvw_cape_m'
|
||||
config['notes'] = ' '
|
||||
config['tags'] = ['torch', 'norm']
|
||||
config['script'] = 'scripts/ae_norm_train.py'
|
||||
|
||||
config['pretrained_model'] = ''
|
||||
|
||||
config['batch_size'] = 4
|
||||
config['train_dir'] = '/data-local/127_2/cape_processed/train/male'
|
||||
config['valid_dir'] = '/data-local/127_2/cape_processed/test/male'
|
||||
config['test_dir'] = '/data-local/127_2/cape_processed/test/male'
|
||||
config['train_frames'] = None
|
||||
config['valid_frames'] = None
|
||||
|
||||
config['resolution'] = 512
|
||||
config['is_skirt'] = False
|
||||
config['dataset_in_memory'] = False
|
||||
config['is_cape'] = True
|
||||
config['remove_global_orientation'] = True
|
||||
config['device'] = 'cuda'
|
||||
|
||||
|
||||
config['smpl_res'] = 14475
|
||||
if config['is_cape']:
|
||||
config['smpl_res'] = 3638
|
||||
|
||||
config['latent_dim'] = 128
|
||||
|
||||
config['bottleneck_scale'] = 2
|
||||
|
||||
config['cond_ae'] = dict()
|
||||
config['cond_ae']['device'] = config['device']
|
||||
config['cond_ae']['is_label'] = False
|
||||
config['cond_ae']['latent_dim'] = config['latent_dim']
|
||||
config['cond_ae']['kernel_size'] = 7
|
||||
config['cond_ae']['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']['bottleneck_size'], 9),
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
|
||||
]
|
||||
|
||||
block_setting_reverse = [
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
|
||||
CNBlockConfig(128, 64, 9),
|
||||
CNBlockConfig(64, 32, 9),
|
||||
CNBlockConfig(32, 16, 3),
|
||||
CNBlockConfig(16, 3, 3),
|
||||
# CNBlockConfig(3, 3, 3),
|
||||
]
|
||||
|
||||
config['cond_ae']['encoder_block_settings'] = block_setting
|
||||
config['cond_ae']['decoder_block_settings'] = block_setting_reverse
|
||||
config['cond_ae']['activation'] = 'tanh'
|
||||
|
||||
|
||||
config['lr'] = 0.0001
|
||||
config['cond_ae']['upscale_output'] = False
|
||||
|
||||
config['signature'] = { 'id',
|
||||
'frame_nr',
|
||||
'uv_mask',
|
||||
'uv_normals',
|
||||
'mesh_mask',
|
||||
'mesh_normals',
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
from utils.convnext_utils import CNBlockConfig
|
||||
|
||||
config = dict()
|
||||
|
||||
config['gpu'] = '3'
|
||||
config['name'] = 'static_iccvw_4'
|
||||
config['notes'] = ' '
|
||||
config['tags'] = ['torch', 'static']
|
||||
config['script'] = 'scripts/ae_static_train.py'
|
||||
|
||||
config['pretrained_model'] = 'checkpoints/iccvw/static_iccvw_3'
|
||||
|
||||
config['batch_size'] = 4
|
||||
config['train_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['valid_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['test_dir'] = '/data-net/dataset/test_ready'
|
||||
|
||||
config['resolution'] = 512
|
||||
config['is_skirt'] = False
|
||||
config['dataset_in_memory'] = False
|
||||
config['is_cape'] = False
|
||||
config['remove_global_orientation'] = True
|
||||
config['device'] = 'cuda'
|
||||
|
||||
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['smpl_res'] = 14475
|
||||
if config['is_cape']:
|
||||
config['smpl_res'] = 3638
|
||||
|
||||
config['latent_dim'] = 128
|
||||
|
||||
config['bottleneck_scale'] = 2
|
||||
|
||||
config['cond_ae'] = dict()
|
||||
config['cond_ae']['device'] = config['device']
|
||||
config['cond_ae']['is_label'] = False
|
||||
config['cond_ae']['latent_dim'] = config['latent_dim']
|
||||
config['cond_ae']['kernel_size'] = 7
|
||||
config['cond_ae']['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']['bottleneck_size'], 9),
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
|
||||
]
|
||||
|
||||
block_setting_reverse = [
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3),
|
||||
CNBlockConfig(64, 32, 9),
|
||||
CNBlockConfig(32, 16, 3),
|
||||
CNBlockConfig(16, 8, 3),
|
||||
CNBlockConfig(8, 3, 3),
|
||||
# CNBlockConfig(3, 3, 3),
|
||||
]
|
||||
|
||||
config['cond_ae']['encoder_block_settings'] = block_setting
|
||||
config['cond_ae']['decoder_block_settings'] = block_setting_reverse
|
||||
config['cond_ae']['activation'] = 'tanh'
|
||||
|
||||
|
||||
config['lr'] = 0.0005
|
||||
config['cond_ae']['upscale_output'] = False
|
||||
|
||||
config['signature'] = { 'id',
|
||||
'frame_nr',
|
||||
'uv_static',
|
||||
'mesh_static',
|
||||
'mesh_mask',
|
||||
'uv_mask',
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
from utils.convnext_utils import CNBlockConfig
|
||||
|
||||
config = dict()
|
||||
|
||||
config['gpu'] = '3'
|
||||
config['name'] = 'static_iccvw_1_skirt'
|
||||
config['notes'] = ' '
|
||||
config['tags'] = ['torch', 'static']
|
||||
config['script'] = 'scripts/ae_static_train.py'
|
||||
|
||||
config['pretrained_model'] = ''
|
||||
|
||||
config['batch_size'] = 4
|
||||
config['train_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['valid_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['test_dir'] = '/data-net/dataset/test_ready'
|
||||
|
||||
config['resolution'] = 512
|
||||
config['is_skirt'] = True
|
||||
config['dataset_in_memory'] = False
|
||||
config['is_cape'] = False
|
||||
config['remove_global_orientation'] = True
|
||||
config['device'] = 'cuda'
|
||||
|
||||
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['smpl_res'] = 14475
|
||||
if config['is_cape']:
|
||||
config['smpl_res'] = 3638
|
||||
|
||||
config['latent_dim'] = 128
|
||||
|
||||
config['bottleneck_scale'] = 2
|
||||
|
||||
config['cond_ae'] = dict()
|
||||
config['cond_ae']['device'] = config['device']
|
||||
config['cond_ae']['is_label'] = False
|
||||
config['cond_ae']['latent_dim'] = config['latent_dim']
|
||||
config['cond_ae']['kernel_size'] = 7
|
||||
config['cond_ae']['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']['bottleneck_size'], 9),
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
|
||||
]
|
||||
|
||||
block_setting_reverse = [
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3),
|
||||
CNBlockConfig(64, 32, 9),
|
||||
CNBlockConfig(32, 16, 3),
|
||||
CNBlockConfig(16, 8, 3),
|
||||
CNBlockConfig(8, 3, 3),
|
||||
# CNBlockConfig(3, 3, 3),
|
||||
]
|
||||
|
||||
config['cond_ae']['encoder_block_settings'] = block_setting
|
||||
config['cond_ae']['decoder_block_settings'] = block_setting_reverse
|
||||
config['cond_ae']['activation'] = 'tanh'
|
||||
|
||||
|
||||
config['lr'] = 0.0005
|
||||
config['cond_ae']['upscale_output'] = False
|
||||
|
||||
config['signature'] = { 'id',
|
||||
'frame_nr',
|
||||
'uv_static',
|
||||
'mesh_static',
|
||||
'mesh_mask',
|
||||
'uv_mask',
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
from utils.convnext_utils import CNBlockConfig
|
||||
|
||||
config = dict()
|
||||
|
||||
config['gpu'] = '5'
|
||||
config['name'] = 'static_iccvw_5'
|
||||
config['notes'] = 'lower learning rate'
|
||||
config['tags'] = ['torch', 'static']
|
||||
config['script'] = 'scripts/ae_static_train.py'
|
||||
|
||||
config['pretrained_model'] = 'checkpoints/iccvw/static_iccvw_3'
|
||||
|
||||
config['batch_size'] = 4
|
||||
config['train_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['valid_dir'] = '/data-local/127/cloth3d_processed'
|
||||
config['test_dir'] = '/data-net/dataset/test_ready'
|
||||
|
||||
config['resolution'] = 512
|
||||
config['is_skirt'] = False
|
||||
config['dataset_in_memory'] = False
|
||||
config['is_cape'] = False
|
||||
config['remove_global_orientation'] = True
|
||||
config['device'] = 'cuda'
|
||||
|
||||
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['smpl_res'] = 14475
|
||||
if config['is_cape']:
|
||||
config['smpl_res'] = 3638
|
||||
|
||||
config['latent_dim'] = 128
|
||||
|
||||
config['bottleneck_scale'] = 2
|
||||
|
||||
config['cond_ae'] = dict()
|
||||
config['cond_ae']['device'] = config['device']
|
||||
config['cond_ae']['is_label'] = False
|
||||
config['cond_ae']['latent_dim'] = config['latent_dim']
|
||||
config['cond_ae']['kernel_size'] = 7
|
||||
config['cond_ae']['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']['bottleneck_size'], 9),
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
|
||||
]
|
||||
|
||||
block_setting_reverse = [
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3),
|
||||
CNBlockConfig(64, 32, 9),
|
||||
CNBlockConfig(32, 16, 3),
|
||||
CNBlockConfig(16, 8, 3),
|
||||
CNBlockConfig(8, 3, 3),
|
||||
# CNBlockConfig(3, 3, 3),
|
||||
]
|
||||
|
||||
config['cond_ae']['encoder_block_settings'] = block_setting
|
||||
config['cond_ae']['decoder_block_settings'] = block_setting_reverse
|
||||
config['cond_ae']['activation'] = 'tanh'
|
||||
|
||||
|
||||
config['lr'] = 0.00005
|
||||
config['cond_ae']['upscale_output'] = False
|
||||
|
||||
config['signature'] = { 'id',
|
||||
'frame_nr',
|
||||
'uv_static',
|
||||
'mesh_static',
|
||||
'mesh_mask',
|
||||
'uv_mask',
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
from utils.convnext_utils import CNBlockConfig
|
||||
|
||||
config = dict()
|
||||
|
||||
config['gpu'] = '4'
|
||||
config['name'] = 'static_iccvw_cape_f'
|
||||
config['notes'] = ' '
|
||||
config['tags'] = ['torch', 'static']
|
||||
config['script'] = 'scripts/ae_static_train.py'
|
||||
|
||||
config['pretrained_model'] = ''
|
||||
|
||||
config['batch_size'] = 4
|
||||
config['train_dir'] = '/data-local/127_2/cape_processed/train/female'
|
||||
config['valid_dir'] = '/data-local/127_2/cape_processed/test/female'
|
||||
config['test_dir'] = '/data-local/127_2/cape_processed/test/female'
|
||||
config['train_frames'] = None
|
||||
config['valid_frames'] = None
|
||||
|
||||
config['resolution'] = 512
|
||||
config['is_skirt'] = False
|
||||
config['dataset_in_memory'] = False
|
||||
config['is_cape'] = True
|
||||
config['remove_global_orientation'] = True
|
||||
config['device'] = 'cuda'
|
||||
|
||||
config['smpl_res'] = 14475
|
||||
if config['is_cape']:
|
||||
config['smpl_res'] = 3638
|
||||
|
||||
config['latent_dim'] = 128
|
||||
|
||||
config['bottleneck_scale'] = 2
|
||||
|
||||
config['cond_ae'] = dict()
|
||||
config['cond_ae']['device'] = config['device']
|
||||
config['cond_ae']['is_label'] = False
|
||||
config['cond_ae']['latent_dim'] = config['latent_dim']
|
||||
config['cond_ae']['kernel_size'] = 7
|
||||
config['cond_ae']['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']['bottleneck_size'], 9),
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
|
||||
]
|
||||
|
||||
block_setting_reverse = [
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3),
|
||||
CNBlockConfig(64, 32, 9),
|
||||
CNBlockConfig(32, 16, 3),
|
||||
CNBlockConfig(16, 8, 3),
|
||||
CNBlockConfig(8, 3, 3),
|
||||
# CNBlockConfig(3, 3, 3),
|
||||
]
|
||||
|
||||
config['cond_ae']['encoder_block_settings'] = block_setting
|
||||
config['cond_ae']['decoder_block_settings'] = block_setting_reverse
|
||||
config['cond_ae']['activation'] = 'tanh'
|
||||
|
||||
|
||||
config['lr'] = 0.0005
|
||||
config['cond_ae']['upscale_output'] = False
|
||||
|
||||
config['signature'] = { 'id',
|
||||
'frame_nr',
|
||||
'uv_static',
|
||||
'mesh_static',
|
||||
'mesh_mask',
|
||||
'uv_mask',
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
from utils.convnext_utils import CNBlockConfig
|
||||
|
||||
config = dict()
|
||||
|
||||
config['gpu'] = '5'
|
||||
config['name'] = 'static_iccvw_cape_m'
|
||||
config['notes'] = ' '
|
||||
config['tags'] = ['torch', 'static']
|
||||
config['script'] = 'scripts/ae_static_train.py'
|
||||
|
||||
config['pretrained_model'] = ''
|
||||
|
||||
config['batch_size'] = 4
|
||||
config['train_dir'] = '/data-local/127_2/cape_processed/train/male'
|
||||
config['valid_dir'] = '/data-local/127_2/cape_processed/test/male'
|
||||
config['test_dir'] = '/data-local/127_2/cape_processed/test/male'
|
||||
config['train_frames'] = None
|
||||
config['valid_frames'] = None
|
||||
|
||||
config['resolution'] = 512
|
||||
config['is_skirt'] = False
|
||||
config['dataset_in_memory'] = False
|
||||
config['is_cape'] = True
|
||||
config['remove_global_orientation'] = True
|
||||
config['device'] = 'cuda'
|
||||
|
||||
config['smpl_res'] = 14475
|
||||
if config['is_cape']:
|
||||
config['smpl_res'] = 3638
|
||||
|
||||
config['latent_dim'] = 128
|
||||
|
||||
config['bottleneck_scale'] = 2
|
||||
|
||||
config['cond_ae'] = dict()
|
||||
config['cond_ae']['device'] = config['device']
|
||||
config['cond_ae']['is_label'] = False
|
||||
config['cond_ae']['latent_dim'] = config['latent_dim']
|
||||
config['cond_ae']['kernel_size'] = 7
|
||||
config['cond_ae']['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']['bottleneck_size'], 9),
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
|
||||
]
|
||||
|
||||
block_setting_reverse = [
|
||||
CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3),
|
||||
CNBlockConfig(64, 32, 9),
|
||||
CNBlockConfig(32, 16, 3),
|
||||
CNBlockConfig(16, 8, 3),
|
||||
CNBlockConfig(8, 3, 3),
|
||||
# CNBlockConfig(3, 3, 3),
|
||||
]
|
||||
|
||||
config['cond_ae']['encoder_block_settings'] = block_setting
|
||||
config['cond_ae']['decoder_block_settings'] = block_setting_reverse
|
||||
config['cond_ae']['activation'] = 'tanh'
|
||||
|
||||
|
||||
config['lr'] = 0.0005
|
||||
config['cond_ae']['upscale_output'] = False
|
||||
|
||||
config['signature'] = { 'id',
|
||||
'frame_nr',
|
||||
'uv_static',
|
||||
'mesh_static',
|
||||
'mesh_mask',
|
||||
'uv_mask',
|
||||
}
|
||||
@@ -9,9 +9,6 @@ from plotly.subplots import make_subplots
|
||||
from torchvision.transforms import Resize
|
||||
from torchvision.transforms.functional import InterpolationMode
|
||||
|
||||
# from tensorflow_graphics.geometry.convolution.graph_pooling import pool, unpool
|
||||
# from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
|
||||
|
||||
|
||||
def m2uv_array(m2uv, is_smpl_orig=False):
|
||||
result = np.zeros((14475 if not is_smpl_orig else 3638, 2), dtype=np.int32)
|
||||
@@ -74,35 +71,6 @@ class UV2Mesh(nn.Module):
|
||||
return x
|
||||
|
||||
|
||||
|
||||
class UV2MeshLabel():
|
||||
def __init__(self, res=256, is_skirt=False, is_smpl_orig=False):
|
||||
super(UV2MeshLabel, self).__init__()
|
||||
self.res = res
|
||||
uv_map = np.load('custom_utils/stretched_uv_map_' + ('skirt' if is_skirt else '3') + ('_orig' if is_smpl_orig else '') + '.npy')
|
||||
self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32)
|
||||
m2uv = np.load('custom_utils/m2uv' + ('_skirt' if is_skirt else '') + ('_orig' if is_smpl_orig else '') + '.npy', allow_pickle=True).item()
|
||||
m2uv = m2uv_array(m2uv, is_smpl_orig=is_smpl_orig)
|
||||
self.uv_pixels = self.uv_map_discreet[m2uv]
|
||||
|
||||
# input = channels_first
|
||||
# @tf.function
|
||||
def uv_to_mesh(self, uv):
|
||||
batch_size = uv.shape[0]
|
||||
|
||||
x1_ind = tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1])
|
||||
x2_ind = tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1])
|
||||
|
||||
x1 = tf.gather_nd(uv, x1_ind, batch_dims=1)
|
||||
x2 = tf.gather_nd(uv, x2_ind, batch_dims=1)
|
||||
|
||||
mesh = (x1 + x2) / 2
|
||||
return mesh
|
||||
|
||||
def __call__(self, uv):
|
||||
return self.uv_to_mesh(uv)
|
||||
|
||||
|
||||
def _maskFaces(V, F, mask, return_unfixed=False):
|
||||
V_mask = mask
|
||||
F_mask = np.take(V_mask, F.reshape(-1), axis=0)
|
||||
@@ -119,20 +87,6 @@ def _maskFaces(V, F, mask, return_unfixed=False):
|
||||
return V_new, F_new
|
||||
|
||||
|
||||
# def _maskFaces_tf(V, F, mask):
|
||||
# V_mask = mask
|
||||
# F_mask = tf.gather(V_mask, F.reshape(-1), axis=0)
|
||||
# F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
|
||||
# F_mask = tf.reduce_sum(tf.reshape(F_mask, (-1, 3)), axis=1) == 3
|
||||
|
||||
# V_ind = tf.zeros(V_mask.shape)
|
||||
# V_ind[V_mask] = tf.range(V_mask.sum())
|
||||
# F_new = tf.gather(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
|
||||
# V_new = V[V_mask, :]
|
||||
|
||||
# return V_new, F_new
|
||||
|
||||
|
||||
def maskFacesBatch(V, F, mask):
|
||||
Vs = []
|
||||
Fs = []
|
||||
@@ -272,25 +226,6 @@ def plot_3d_collision(gt, collision_1, collision_2, collision_faces, collision_c
|
||||
return fig
|
||||
|
||||
|
||||
def batch_to_dict(batch, signature):
|
||||
result = dict()
|
||||
for i, key in enumerate(signature.keys()):
|
||||
result[key] = batch[i]
|
||||
return result
|
||||
|
||||
conditions = ['pose_shape', 'uv_static', 'uv_mask', 'uv_shaped_body']
|
||||
def nest_conditions(batch):
|
||||
batch['conditions'] = {d: batch.pop(d) for d in conditions}
|
||||
return batch
|
||||
|
||||
|
||||
def sparse_to_tensor(L):
|
||||
idx = np.mat([L.row, L.col]).transpose().astype(np.float32)
|
||||
data = L.data.astype(dtype=np.float32)
|
||||
shape = L.shape
|
||||
return tf.SparseTensor(idx, data, shape)
|
||||
|
||||
|
||||
|
||||
def get_summary_str(model):
|
||||
# from torchview import draw_graph
|
||||
@@ -304,6 +239,7 @@ def get_summary_str(model):
|
||||
|
||||
return [f'all_params: {total_params}, trainable: {trainable_params}']
|
||||
|
||||
|
||||
def get_summary_html(model):
|
||||
lines = get_summary_str(model)
|
||||
html = ''
|
||||
@@ -312,88 +248,6 @@ def get_summary_html(model):
|
||||
return html
|
||||
|
||||
|
||||
def repeat_laplacian(L, times):
|
||||
idx = np.repeat(np.mat([L.row, L.col]), repeats=times, axis=1)
|
||||
idx_offset = np.linspace(0, L.shape[0] * (times-1), times).astype(np.int32)
|
||||
idx_offset = np.repeat(idx_offset, L.row.shape[0], axis=0)
|
||||
idx += idx_offset[np.newaxis,:]
|
||||
data = np.repeat(L.data, repeats=times, axis=0)
|
||||
shape = (L.shape[0] * times, L.shape[1] * times)
|
||||
return sp.sparse.coo_matrix((data, (np.squeeze(np.asarray(idx[0,:])), np.squeeze(np.asarray(idx[1,:])))), shape)
|
||||
|
||||
|
||||
def merge_verts(Vs, batchDim):
|
||||
V = tf.concat(tf.unstack(Vs, num=batchDim), axis=0)
|
||||
return V
|
||||
|
||||
|
||||
class GraphConvolution(nn.Module):
|
||||
def __init__(self, shape, activation=None, name='gcn', L=None, batchSize=None, normalization=None):
|
||||
super(GraphConvolution, self).__init__()
|
||||
self.shape = shape
|
||||
self.w = self.add_weight(name=name + '/weight', shape=(2, *shape), initializer=tf.initializers.glorot_normal, trainable=True)
|
||||
self.b = self.add_weight(name=name + '/bias', shape=shape[-1], initializer=tf.initializers.zeros, trainable=True)
|
||||
self.activation = activation or (lambda x: x)
|
||||
self.normalization = normalization
|
||||
self.batchSize = batchSize
|
||||
self.indices = tf.cast(tf.linspace(0, self.batchSize * L.shape[1], self.batchSize+1), dtype=tf.int32)
|
||||
self.L = repeat_laplacian(L, self.batchSize)
|
||||
|
||||
def gather(self):
|
||||
return [self.w, self.b]
|
||||
|
||||
def _single(self, X):
|
||||
X0 = tf.einsum('ab,bc->ac', X, self.w[0]) # Node
|
||||
X1 = tf.einsum('ab,bc->ac', X, self.w[1]) # Neigh
|
||||
# Neigh conv
|
||||
X1 = edge_convolution_template(
|
||||
X1, sparse_to_tensor(self.L), sizes=None,
|
||||
edge_function=lambda x, y: y,
|
||||
reduction='weighted',
|
||||
edge_function_kwargs={}
|
||||
)
|
||||
X = X0 + X1 + self.b
|
||||
return X
|
||||
|
||||
def call(self, X, training):
|
||||
X = merge_verts(X, self.batchSize)
|
||||
X = self._single(X)
|
||||
X = tf.stack([X[self.indices[i-1]:self.indices[i]] for i in range(1, len(self.indices))])
|
||||
X = self.activation(X)
|
||||
if self.normalization:
|
||||
X = self.normalization(X, training=training)
|
||||
return X
|
||||
|
||||
|
||||
def replace_arrays(string):
|
||||
out = string[:1]
|
||||
i = 0
|
||||
while i < len(string) and string[i+1:].find('[') != -1:
|
||||
start = string[i+1:].find('[')
|
||||
end = string[i+1+start:].find(']')
|
||||
out += string[i+1:i+1+start]
|
||||
array = string[i+1+start:i+1+start+end+1].replace('\n', '')
|
||||
out += array
|
||||
i = i + 1 + start + end
|
||||
|
||||
out += string[i+1:]
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def group_inputs_for_levels(batch, config):
|
||||
levels_input = []
|
||||
for res in config['pyramid']['resolutions']:
|
||||
res_modifier = '_' + str(res) if res != 512 else ''
|
||||
if config['use_skinning']:
|
||||
uv_orig = batch['uv_unposed' + res_modifier]
|
||||
else:
|
||||
uv_orig = batch['uv_posed' + res_modifier]
|
||||
levels_input.append([uv_orig, batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier], batch['uv_mask' + res_modifier], batch['uv_mask']])
|
||||
conds = [batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature']]
|
||||
return levels_input, conds
|
||||
|
||||
|
||||
class Accumulator():
|
||||
def __init__(self):
|
||||
self.dict = dict()
|
||||
|
||||
@@ -0,0 +1,530 @@
|
||||
import os
|
||||
import sys
|
||||
import glob
|
||||
import tqdm
|
||||
import math
|
||||
import time
|
||||
import random
|
||||
import shutil
|
||||
import collections
|
||||
import argparse
|
||||
import pickle
|
||||
import numpy as np
|
||||
import scipy as sp
|
||||
import trimesh as tm
|
||||
# import cmapy as cm
|
||||
import tables as tb
|
||||
import scipy.io as sio
|
||||
import scipy.interpolate as spi
|
||||
import matplotlib.pyplot as plt
|
||||
import scipy.stats as stat
|
||||
import cv2
|
||||
import plotly.express as ex
|
||||
import plotly.graph_objects as go
|
||||
|
||||
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
||||
from data.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo, read_posed_body_frame, read_posed_body, read_posed_garment, read_curvatures
|
||||
from data.cloth3d_io import readPC2
|
||||
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "6"
|
||||
|
||||
def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rbf', garment_uv_mask=None):
|
||||
verts_uv = np.zeros((res * 2, res, 3))
|
||||
uv_values = np.zeros((uv_map.shape[0], 3))
|
||||
uv_values[m2uv[:, 0]] = verts
|
||||
uv_values[m2uv[:, 1]] += verts
|
||||
uv_values = uv_values[uv_map_mask[type], :] # add garment mask
|
||||
uv_values /= uv_counter[type][:, np.newaxis]
|
||||
return_garment_uv_mask = False
|
||||
|
||||
if mask is not None:
|
||||
garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
|
||||
garment_mask[m2uv[:, 0]] = mask
|
||||
garment_mask[m2uv[:, 1]] += mask
|
||||
garment_mask = garment_mask[uv_map_mask[type]] # add garment mask
|
||||
else:
|
||||
garment_mask = np.sum(uv_values, axis=-1) != 0
|
||||
|
||||
# test
|
||||
verts_uv[uv_map_discreet_unique[type][:, 0], uv_map_discreet_unique[type][:, 1], :] = uv_values
|
||||
# verts_uv_save = verts_uv.copy()
|
||||
|
||||
# garment_mask = np.sum(uv_values, axis=-1) != 0
|
||||
unique_uv_values = uv_values[garment_mask, :]
|
||||
uv_map_discreet_unique_unique = uv_map_discreet_unique[type][garment_mask, :]
|
||||
rbf_dest_current = rbf_dest[type]
|
||||
|
||||
if garment_uv_mask is None and mask is not None:
|
||||
garment_uv_mask = np.zeros((res * 2, res))
|
||||
garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255
|
||||
garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (4 * 2 + 1, 4 * 2 + 1), (4, 4)), iterations=4)
|
||||
garment_uv_mask_border = garment_uv_mask.astype(bool)
|
||||
garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)
|
||||
garment_uv_mask_border = ~garment_uv_mask_border & garment_uv_mask
|
||||
garment_uv_mask = garment_uv_mask * uv_mask[type][:,:,0]
|
||||
rbf_dest_current = rbf_dest[type] * garment_uv_mask
|
||||
return_garment_uv_mask = True
|
||||
|
||||
|
||||
rbf_dest_current = np.nonzero(rbf_dest_current)
|
||||
rbf_dest_current = np.array(rbf_dest_current).T
|
||||
|
||||
# uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).T
|
||||
if method == 'rbf':
|
||||
interpolator = spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values, neighbors=neighbours, kernel='linear')
|
||||
elif method == 'nearest':
|
||||
interpolator = spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
final = interpolator(rbf_dest_current)
|
||||
|
||||
result = np.zeros((res * 2, res, 3))
|
||||
|
||||
# not sure, shouldn't be needed
|
||||
# rbf_dest[:, [0,1]] = rbf_dest[:, [1, 0]]
|
||||
|
||||
result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = final
|
||||
result[uv_map_discreet_unique[type][:, 0], uv_map_discreet_unique[type][:, 1], :] = uv_values
|
||||
if garment_uv_mask is not None:
|
||||
result = result * garment_uv_mask[:,:,np.newaxis]
|
||||
|
||||
if return_garment_uv_mask:
|
||||
return result, garment_uv_mask, garment_uv_mask_border
|
||||
else:
|
||||
return result
|
||||
|
||||
|
||||
def convert_to_uv_batch_body(res, m2uv, uv_map, verts, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None):
|
||||
batch = verts.shape[0]
|
||||
verts_uv = np.zeros((batch, res * 2, res, 3))
|
||||
uv_values = np.zeros((batch, uv_map.shape[0], 3))
|
||||
uv_values[:, m2uv[:, 0]] = verts
|
||||
uv_values[:, m2uv[:, 1]] += verts
|
||||
uv_values = uv_values[:, uv_map_mask['non-skirt'], :]
|
||||
uv_values /= uv_counter['non-skirt'][np.newaxis, :, np.newaxis]
|
||||
return_garment_uv_mask = False
|
||||
# test
|
||||
verts_uv[:, uv_map_discreet_unique['non-skirt'][:, 0], uv_map_discreet_unique['non-skirt'][:, 1], :] = uv_values
|
||||
|
||||
if mask is not None:
|
||||
garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
|
||||
garment_mask[m2uv[:, 0]] = mask
|
||||
garment_mask[m2uv[:, 1]] += mask
|
||||
garment_mask = garment_mask[uv_map_mask['non-skirt']] # add garment mask
|
||||
else:
|
||||
garment_mask = np.sum(uv_values[0,:,:], axis=-1) != 0
|
||||
|
||||
unique_uv_values = uv_values[:, garment_mask, :]
|
||||
uv_map_discreet_unique_unique = uv_map_discreet_unique['non-skirt'][garment_mask, :]
|
||||
rbf_dest_current = rbf_dest['non-skirt']
|
||||
|
||||
# add back mesh vertices that are not part of the garment
|
||||
uv_map_discreet_unique_unique_inv = uv_map_discreet_unique['non-skirt'][np.invert(garment_mask), :]
|
||||
rbf_dest_fill = np.zeros((res * 2, res), dtype=bool)
|
||||
rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True
|
||||
|
||||
if garment_uv_mask is None and mask is not None:
|
||||
garment_uv_mask = np.zeros((res * 2, res))
|
||||
garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255
|
||||
garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=4)
|
||||
garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)
|
||||
return_garment_uv_mask = True
|
||||
|
||||
if garment_uv_mask is not None:
|
||||
rbf_dest_current = (rbf_dest['non-skirt'] + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border
|
||||
|
||||
rbf_dest_current = np.nonzero(rbf_dest_current)
|
||||
rbf_dest_current = np.array(rbf_dest_current).T
|
||||
|
||||
final = []
|
||||
for i in range(batch):
|
||||
print('processing frame: ' + str(i))
|
||||
if method == 'rbf':
|
||||
interpolator = spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values[i,:,:], neighbors=neighbours, kernel='linear')
|
||||
elif method == 'nearest':
|
||||
interpolator = spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values[i,:,:])
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
result = np.zeros((res * 2, res, 3), np.float32)
|
||||
result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = interpolator(rbf_dest_current)
|
||||
result[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1], :] = unique_uv_values[i,:,:]
|
||||
|
||||
if garment_uv_mask is not None:
|
||||
result *= garment_uv_mask[:,:,np.newaxis]
|
||||
|
||||
final.append(result)
|
||||
|
||||
if return_garment_uv_mask:
|
||||
return final, garment_uv_mask
|
||||
else:
|
||||
return final
|
||||
|
||||
|
||||
def convert_to_uv_batch(res, m2uv, uv_map, verts, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None, id=None):
|
||||
batch = len(verts)
|
||||
|
||||
assert method == 'rbf'
|
||||
|
||||
assert mask is not None
|
||||
garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
|
||||
garment_mask[m2uv[:, 0]] = mask
|
||||
garment_mask[m2uv[:, 1]] += mask
|
||||
garment_mask = garment_mask[uv_map_mask[type]] # add garment mask
|
||||
|
||||
rbf_dest_current = rbf_dest[type]
|
||||
|
||||
# add back mesh vertices that are not part of the garment
|
||||
uv_map_discreet_unique_unique_inv = uv_map_discreet_unique[type][np.invert(garment_mask), :]
|
||||
rbf_dest_fill = np.zeros((res * 2, res), dtype=bool)
|
||||
rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True
|
||||
|
||||
assert not (garment_uv_mask is None and mask is not None)
|
||||
assert garment_uv_mask is not None
|
||||
rbf_dest_current = (rbf_dest[type] + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border
|
||||
|
||||
rbf_dest_current = np.nonzero(rbf_dest_current)
|
||||
rbf_dest_current = np.array(rbf_dest_current).T
|
||||
|
||||
out = []
|
||||
|
||||
for i in range(batch):
|
||||
print('processing frame: ' + str(i))
|
||||
|
||||
uv_values = np.zeros((uv_map.shape[0], 3))
|
||||
uv_values[m2uv[:, 0]] = verts[i]
|
||||
uv_values[m2uv[:, 1]] += verts[i]
|
||||
uv_values = uv_values[uv_map_mask[type], :]
|
||||
uv_values /= uv_counter[type][:, np.newaxis]
|
||||
|
||||
unique_uv_values = uv_values[garment_mask, :]
|
||||
uv_map_discreet_unique_unique = uv_map_discreet_unique[type][garment_mask, :]
|
||||
|
||||
interpolator = spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values[:,:], neighbors=neighbours, kernel='linear')
|
||||
result = np.zeros((res * 2, res, 3), np.float32)
|
||||
result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = interpolator(rbf_dest_current)
|
||||
result[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1], :] = unique_uv_values[:,:]
|
||||
|
||||
if garment_uv_mask is not None:
|
||||
result *= garment_uv_mask[:,:,np.newaxis]
|
||||
|
||||
out.append(result)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def convert_to_uv_batch_label(res, m2uv, uv_map, values, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None):
|
||||
batch = values.shape[0]
|
||||
values = (values + 1) * mask
|
||||
values_uv = np.zeros((batch, res * 2, res))
|
||||
missing = np.zeros((res * 2, res))
|
||||
uv_values = np.zeros((batch, uv_map.shape[0]))
|
||||
uv_values[:, m2uv[:, 0]] = values
|
||||
uv_values[:, m2uv[:, 1]] = values
|
||||
uv_values = uv_values[:, uv_map_mask]
|
||||
return_garment_uv_mask = False
|
||||
|
||||
if mask is not None:
|
||||
garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
|
||||
garment_mask[m2uv[:, 0]] = mask
|
||||
garment_mask[m2uv[:, 1]] += mask
|
||||
garment_mask = garment_mask[uv_map_mask] # add garment mask
|
||||
else:
|
||||
raise Exception
|
||||
|
||||
|
||||
# test
|
||||
values_uv[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = uv_values
|
||||
missing[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = (uv_values[0,:] == 0)
|
||||
|
||||
unique_uv_values = uv_values[:, garment_mask]
|
||||
uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
|
||||
# rbf_dest_current = rbf_dest.astype(bool) + garment_uv_mask_border * uv_mask[:,:,0]
|
||||
|
||||
if garment_uv_mask is None and mask is not None:
|
||||
garment_uv_mask = np.zeros((res * 2, res))
|
||||
garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255
|
||||
garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=4)
|
||||
garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)
|
||||
return_garment_uv_mask = True
|
||||
|
||||
rbf_dest_current = rbf_dest * garment_uv_mask + missing * garment_uv_mask
|
||||
rbf_dest_current = np.nonzero(rbf_dest_current)
|
||||
rbf_dest_current = np.array(rbf_dest_current).T
|
||||
|
||||
# uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).T
|
||||
final = []
|
||||
for i in range(batch):
|
||||
print('processing frame: ' + str(i))
|
||||
interpolator = spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values[i,:])
|
||||
final.append(interpolator(rbf_dest_current))
|
||||
final = np.asarray(final)
|
||||
|
||||
result = np.zeros((batch, res * 2, res))
|
||||
result[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = uv_values
|
||||
result[:, rbf_dest_current[:, 0], rbf_dest_current[:, 1]] = final
|
||||
if garment_uv_mask is not None:
|
||||
result = result * garment_uv_mask[np.newaxis,:,:]
|
||||
|
||||
if return_garment_uv_mask:
|
||||
return result, garment_uv_mask
|
||||
else:
|
||||
return result
|
||||
|
||||
|
||||
def m2uv_array(m2uv):
|
||||
if m2uv is None:
|
||||
return None
|
||||
result = np.zeros((3638, 2), dtype=np.int32)
|
||||
for k, v in m2uv.items():
|
||||
if len(v) == 2:
|
||||
result[k] = list(v)
|
||||
elif len(v) == 1:
|
||||
x = v.pop()
|
||||
result[k] = [x, x]
|
||||
else:
|
||||
assert False
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def uv_to_mesh(uv, m2uv):
|
||||
res = uv.shape[0]
|
||||
mesh = np.zeros((14475, 3))
|
||||
uv_pixels = uv_map_discreet[m2uv]
|
||||
# counter = np.ones((14475))
|
||||
# counter[m2uv[:, 0] == m2uv[:, 1]] = 2
|
||||
mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
|
||||
return mesh
|
||||
pass
|
||||
|
||||
|
||||
def verify_reconstruction(orig, reconstructed):
|
||||
error = np.sqrt(np.sum((orig - reconstructed) ** 2, axis=-1))
|
||||
colors = np.array([cm.color('jet', 1 - i) for i in error])
|
||||
error_sum = np.sum(error)
|
||||
|
||||
|
||||
verts_uv = np.zeros((uv_res * 2, uv_res, 3))
|
||||
uv_values = np.zeros((uv_map.shape[0], 3))
|
||||
uv_values[m2uv[:, 0]] = colors
|
||||
uv_values[m2uv[:, 1]] += colors
|
||||
uv_counter = np.zeros((uv_map.shape[0]))
|
||||
uv_counter[m2uv[:, 0]] = 1
|
||||
uv_counter[m2uv[:, 1]] += 1
|
||||
uv_values /= uv_counter[:, np.newaxis]
|
||||
verts_uv[uv_map_discreet[:, 0], uv_map_discreet[:, 1], :] = uv_values
|
||||
fig = plt.figure(figsize=(15, 15))
|
||||
ax = fig.add_subplot(111)
|
||||
ax.scatter(uv_map[:, 0], uv_map[:, 1], s=5, c=uv_values/255)
|
||||
plt.show()
|
||||
|
||||
|
||||
|
||||
fig = plt.figure(figsize=(15, 15))
|
||||
ax = fig.add_subplot(121, projection='3d')
|
||||
x = orig
|
||||
ax.set_xlim3d(-1, 1)
|
||||
ax.set_ylim3d(-1, 1)
|
||||
ax.set_zlim3d(-1, 1)
|
||||
ax.set_xlabel("X")
|
||||
ax.set_ylabel("Y")
|
||||
ax.set_zlabel("Z")
|
||||
ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5)
|
||||
|
||||
ax = fig.add_subplot(122, projection='3d')
|
||||
x = reconstructed#[error > 0.1, :]
|
||||
ax.set_xlim3d(-1, 1)
|
||||
ax.set_ylim3d(-1, 1)
|
||||
ax.set_zlim3d(-1, 1)
|
||||
ax.set_xlabel("X")
|
||||
ax.set_ylabel("Y")
|
||||
ax.set_zlabel("Z")
|
||||
ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5, c=colors/255) #[error > 0.1, :]
|
||||
|
||||
plt.show()
|
||||
|
||||
pass
|
||||
|
||||
# https://stackoverflow.com/questions/38187286/find-unique-pairs-in-list-of-pairs
|
||||
def unique_pairs(a):
|
||||
a.sort(axis=1)
|
||||
b = np.ascontiguousarray(a).view(
|
||||
np.dtype((np.void, a.dtype.itemsize * a.shape[1]))
|
||||
)
|
||||
_, inv, ct = np.unique(b, return_inverse=True, return_counts=True)
|
||||
return ct[inv] == 1
|
||||
|
||||
|
||||
def unique_uv_map_mask(uv_map_discreet):
|
||||
mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool)
|
||||
temp = np.zeros((uv_res * 2, uv_res))
|
||||
for p in uv_map_discreet:
|
||||
temp[p[0], p[1]] += 1
|
||||
# temp = temp == 1
|
||||
for i, p in enumerate(uv_map_discreet):
|
||||
if temp[p[0], p[1]] == 1:
|
||||
mask[i] = True
|
||||
elif temp[p[0], p[1]] == 2:
|
||||
mask[i] = True
|
||||
temp[p[0], p[1]] = 0
|
||||
return mask
|
||||
|
||||
|
||||
def test(verts, faces, values):
|
||||
uv_values = np.zeros((uv_map.shape[0], 3))
|
||||
uv_values[m2uv[:, 0]] = values
|
||||
uv_values[m2uv[:, 1]] += values
|
||||
uv_counter = np.zeros((uv_map.shape[0]))
|
||||
uv_counter[m2uv[:, 0]] = 1
|
||||
uv_counter[m2uv[:, 1]] += 1
|
||||
uv_values /= uv_counter[:, np.newaxis]
|
||||
x = uv_values
|
||||
verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1)
|
||||
im = ((x - x.min()) / (x.max() - x.min())) * 255
|
||||
tm.Trimesh(verts, faces, vertex_colors=im.astype(np.int32)).show()
|
||||
|
||||
|
||||
def get_shapedirs():
|
||||
to_keep = sio.loadmat(os.path.join('smpl/keep_remove.mat'))['keep'].reshape(-1)
|
||||
shapedirs = []
|
||||
with open(os.path.join('smpl/super_model_f.pkl'), 'rb') as f:
|
||||
female = pickle.load(f, encoding="latin1")
|
||||
shapedirs.append(female['shapedirs'][to_keep, :])
|
||||
with open(os.path.join('smpl/super_model_m.pkl'), 'rb') as f:
|
||||
male = pickle.load(f, encoding="latin1")
|
||||
shapedirs.append(male['shapedirs'][to_keep, :])
|
||||
return shapedirs
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/test/female')
|
||||
parser.add_argument("-s", "--start", help = "start of range", default=0, type=int)
|
||||
parser.add_argument("-e", "--end", help = "end of range", default=100, type=int)
|
||||
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
|
||||
parser.add_argument("-d", "--debug", help = "end of range", default=False, action='store_true')
|
||||
args = parser.parse_args()
|
||||
|
||||
F = {'non-skirt': np.load("smpl/smpl_faces_orig.npy"), 'skirt': None}
|
||||
uv_map = {'non-skirt': np.load('custom_utils/stretched_uv_map_3_orig.npy'), 'skirt': None}
|
||||
m2uv = {'non-skirt': np.load('custom_utils/m2uv_orig.npy', allow_pickle=True).item(), 'skirt': None}
|
||||
m2uv = {k: m2uv_array(v) for k, v in m2uv.items()}
|
||||
uv_mask = {'non-skirt': np.load('custom_utils/mask_256x512.npy'), 'skirt': np.load('custom_utils/mask_256x512_skirt.npy')}
|
||||
|
||||
|
||||
uv_res = 256
|
||||
uv_map_discreet = dict()
|
||||
uv_map_discreet_unique = dict()
|
||||
uv_counter = dict()
|
||||
rbf_dest = dict()
|
||||
uv_map_mask = dict()
|
||||
for k,v in uv_mask.items():
|
||||
if k == 'skirt':
|
||||
continue
|
||||
uv_mask[k] = (cv2.resize(v * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
|
||||
# uv_mask_dilate = cv2.erode(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_CROSS, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=1).astype(bool)[:, :, np.newaxis]
|
||||
# uv_mask_border = uv_mask * np.invert(uv_mask_dilate)
|
||||
uv_mask[k] = cv2.dilate(v.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
|
||||
|
||||
uv_map_discreet[k] = np.floor(uv_map[k] * uv_res).astype(np.int32)
|
||||
uv_map_mask[k] = unique_uv_map_mask(uv_map_discreet[k])
|
||||
|
||||
uv_counter[k] = np.zeros((uv_map[k].shape[0]))
|
||||
uv_counter[k][m2uv[k][:, 0]] = 1
|
||||
uv_counter[k][m2uv[k][:, 1]] += 1
|
||||
uv_counter[k] = uv_counter[k][uv_map_mask[k]]
|
||||
|
||||
uv_map_discreet_unique[k] = uv_map_discreet[k][uv_map_mask[k], :]
|
||||
|
||||
rbf_dest[k] = np.ones((uv_res * 2, uv_res))
|
||||
rbf_dest[k][uv_map_discreet_unique[k][:, 0], uv_map_discreet_unique[k][:, 1]] = 0
|
||||
rbf_dest[k][np.invert(uv_mask[k][:,:,0])] = 0
|
||||
|
||||
|
||||
#########################################################################################
|
||||
|
||||
folder = args.folder
|
||||
# data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_normals', 'body_posed', 'mesh_unposed'] # 'mesh_unposed',
|
||||
data_to_convert = ['mesh_static', 'mesh_normals'] # 'mesh_unposed',
|
||||
|
||||
frames = sorted(os.listdir(folder + '/mesh_posed_body_no_orient'))#[args.start:args.end]
|
||||
frames = [[frame.split('.')[0], frame.split('.')[1]] for frame in frames]
|
||||
ids = list(set([frame[0] for frame in frames]))
|
||||
frames_by_id = {id: [frame[1] for frame in frames if frame[0] == id] for id in ids}
|
||||
|
||||
shapedirs = get_shapedirs()
|
||||
|
||||
root = '/'.join(folder.split('/')[:-1])
|
||||
type = 'non-skirt'
|
||||
|
||||
os.makedirs(folder, exist_ok=True)
|
||||
os.makedirs(os.path.join(folder, 'uv_posed_no_orient'), exist_ok=True)
|
||||
os.makedirs(os.path.join(folder, 'uv_static'), exist_ok=True)
|
||||
os.makedirs(os.path.join(folder, 'uv_normals'), exist_ok=True)
|
||||
os.makedirs(os.path.join(folder, 'uv_unposed_no_shape'), exist_ok=True)
|
||||
os.makedirs(os.path.join(folder, 'uv_body_posed_no_orient'), exist_ok=True)
|
||||
os.makedirs(os.path.join(folder, 'uv_mask'), exist_ok=True)
|
||||
|
||||
mask = np.ones((3638), dtype=bool)
|
||||
|
||||
for id in tqdm.tqdm(ids):
|
||||
|
||||
|
||||
# x = readPC2(os.path.join(folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V']
|
||||
# mesh_posed = np.zeros((x.shape[0], 14475, 3), dtype=np.float32)
|
||||
# mesh_posed[:,mask, :] = x
|
||||
if 'mesh_unposed' in data_to_convert:
|
||||
mesh_unposed = np.zeros((x.shape[0], 14475, 3), dtype=np.float32)
|
||||
mesh_unposed[:,mask, :] = readPC2(os.path.join(folder, 'mesh_unposed_no_shape' ,f'{id}_{garment}.pc16'), True)['V']
|
||||
if 'mesh_static' in data_to_convert:
|
||||
mesh_static = np.load(os.path.join(folder, 'mesh_static', f'{id}.npz'))['mesh_static']
|
||||
if 'mesh_normals' in data_to_convert:
|
||||
mesh_normals = np.stack([np.load(os.path.join(folder, 'mesh_normals', f'{id}.{int(i):06d}.npz'))['mesh_normals'] for i in frames_by_id[id]], axis=0)
|
||||
if 'body_posed' in data_to_convert and k != 1:
|
||||
body_posed = readPC2(os.path.join(folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V']
|
||||
|
||||
start = time.perf_counter()
|
||||
|
||||
uv_static, uv_mask_local, garment_uv_mask_border = convert_to_uv(uv_res, m2uv[type], uv_map[type], mesh_static, neighbours=3, mask=mask, method='rbf', garment_uv_mask=None)
|
||||
uv_static = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_static[None, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
|
||||
uv_static = uv_static[0].astype(np.float32)
|
||||
|
||||
if 'mesh_posed' in data_to_convert:
|
||||
uv_posed = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
|
||||
if 'mesh_unposed' in data_to_convert:
|
||||
uv_unposed = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
|
||||
if 'mesh_normals' in data_to_convert:
|
||||
uv_normals = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_normals[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
|
||||
if 'body_posed' in data_to_convert and k != 1:
|
||||
uv_body_posed = convert_to_uv_batch_body(uv_res, m2uv['non-skirt'], uv_map['non-skirt'], body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border)
|
||||
|
||||
print(time.perf_counter() - start)
|
||||
|
||||
# ex.imshow(uv_mask_local).show()
|
||||
# ex.imshow((uv_static+1)/2).show()
|
||||
# ex.imshow((uv_posed[0]+1)/2).show()
|
||||
# ex.imshow((uv_unposed[0]+1)/2).show()
|
||||
# ex.imshow((uv_normals[0]+1.5)/3).show()
|
||||
# ex.imshow((uv_body_posed[0]+1)/2).show()
|
||||
|
||||
# pass
|
||||
# np.savez_compressed(os.path.join(folder, 'uv_static', f'{id}_{garment}.npz'), uv_static=uv_static)
|
||||
# np.savez_compressed(os.path.join(folder, 'uv_mask', f'{id}_{garment}.npz'), uv_mask=uv_mask_local)
|
||||
|
||||
if 'mesh_posed' in data_to_convert:
|
||||
for i in range(len(uv_posed)):
|
||||
np.savez_compressed(os.path.join(folder, 'uv_posed_no_orient', f'{id}_{i}_{garment}.npz'), uv_posed=uv_posed[i])
|
||||
if 'mesh_unposed' in data_to_convert:
|
||||
for i in range(len(uv_unposed)):
|
||||
np.savez_compressed(os.path.join(folder, 'uv_unposed_no_shape', f'{id}_{i}_{garment}.npz'), uv_unposed=uv_unposed[i])
|
||||
if 'mesh_normals' in data_to_convert:
|
||||
for i, frame_nr in enumerate(frames_by_id[id]):
|
||||
np.savez_compressed(os.path.join(folder, 'uv_normals', f'{id}.{frame_nr}.npz'), uv_normals=uv_normals[i])
|
||||
if 'body_posed' in data_to_convert and k != 1:
|
||||
for i in range(len(uv_body_posed)):
|
||||
np.savez_compressed(os.path.join(folder, 'uv_body_posed_no_orient', f'{id}_{i}.npz'), uv_body_posed=uv_body_posed[i])
|
||||
|
||||
# pass
|
||||
+143
-75
@@ -1,7 +1,7 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
# os.environ["CUDA_VISIBLE_DEVICES"] = '6'
|
||||
# os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'
|
||||
|
||||
import math
|
||||
import time
|
||||
@@ -14,13 +14,14 @@ import scipy.io as sio
|
||||
import plotly.express as ex
|
||||
import torchvision
|
||||
|
||||
from torch.utils.data import Dataset
|
||||
from torchvision.transforms import Resize
|
||||
from torchvision.transforms.functional import InterpolationMode
|
||||
|
||||
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
||||
|
||||
from data.utils import Frame
|
||||
from data.cloth3d_io import readPC2, readPC2Frame
|
||||
from data.dataset_cape import DatasetCape
|
||||
|
||||
|
||||
def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=None):
|
||||
@@ -60,7 +61,6 @@ class StaticCache():
|
||||
self.folder = folder
|
||||
assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
|
||||
self.resolutions = resolutions
|
||||
self.filled = ''
|
||||
self.is_cape = is_cape
|
||||
self.device = device
|
||||
|
||||
@@ -74,50 +74,65 @@ class StaticCache():
|
||||
self.uv_masks[res] = down.cpu().numpy()
|
||||
|
||||
|
||||
def read(self, id):
|
||||
def read(self, id, garment):
|
||||
data = dict()
|
||||
|
||||
mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', f'{id}_{garment}.npy')); mesh_mask = mesh_mask.astype(bool)
|
||||
data['mesh_mask'] = mesh_mask
|
||||
|
||||
for res in self.resolutions:
|
||||
res_string = ''
|
||||
if res != 512:
|
||||
res_string = '_' + str(res)
|
||||
|
||||
if 'uv_static' + res_string in self.signature:
|
||||
uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; data['uv_static' + res_string] = uv_static
|
||||
uv_static = np.load(os.path.join(self.folder, f'uv_static{res_string}', f'{id}_{garment}.npz')); uv_static = uv_static['uv_static']; data['uv_static' + res_string] = uv_static
|
||||
if 'uv_mask' + res_string in self.signature:
|
||||
if self.is_cape:
|
||||
data['uv_mask' + res_string] = self.uv_masks[res]
|
||||
else:
|
||||
uv_mask = np.load(os.path.join(self.folder, 'uv_mask' + res_string, id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask
|
||||
uv_mask = np.load(os.path.join(self.folder, f'uv_mask{res_string}', f'{id}_{garment}.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask
|
||||
|
||||
if 'mesh_static' in self.signature:
|
||||
mesh_static = np.load(os.path.join(self.folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static.astype(np.float32)
|
||||
if 'mesh_mask' in self.signature:
|
||||
if self.is_cape:
|
||||
data['mesh_mask'] = np.ones((3638), np.float32)
|
||||
else:
|
||||
mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask
|
||||
mesh_static = np.load(os.path.join(self.folder, 'mesh_static', f'{id}_{garment}.npy'))
|
||||
data['mesh_static'] = np.zeros((14475, 3), dtype=np.float32)
|
||||
data['mesh_static'][mesh_mask] = mesh_static.astype(np.float32)
|
||||
|
||||
# if 'uv_shaped_body' in self.signature:
|
||||
# uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; data['uv_shaped_body'] = uv_shaped_body
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def __getitem__(self, id) -> dict:
|
||||
if id not in self.store:
|
||||
item = self.read(id)
|
||||
def __getitem__(self, frame: Frame) -> dict:
|
||||
if (frame.id, frame.garment) not in self.store:
|
||||
item = self.read(frame.id, frame.garment)
|
||||
stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()}
|
||||
for key, value in stored_data.items():
|
||||
if key.startswith('uv_') and value.shape[-1] == 3:
|
||||
stored_data[key] = torch.permute(value, [2,0,1])
|
||||
self.store[id] = stored_data
|
||||
self.store[(frame.id, frame.garment)] = stored_data
|
||||
|
||||
data = { key: value.to(self.device) for key, value in self.store[id].items()}
|
||||
data = { key: value.to(self.device) for key, value in self.store[(frame.id, frame.garment)].items()}
|
||||
return data
|
||||
|
||||
|
||||
class Dataset(Dataset):
|
||||
def __init__(self, folder, signature, resolutions=[512], frames_file='frames_all.txt', device='cuda', use_cache=False, is_cape=False):
|
||||
|
||||
def read_frames_file(frames_file, shuffle=False):
|
||||
frames = []
|
||||
with open(frames_file, 'r') as f:
|
||||
lines = f.readlines()
|
||||
lines = [line.strip() for line in lines]
|
||||
frames = [Frame(line.split()[0], int(line.split()[1]), line.split()[2]) for line in lines]
|
||||
|
||||
if shuffle:
|
||||
random.shuffle(frames)
|
||||
|
||||
return frames
|
||||
|
||||
|
||||
class DatasetCloth3D(torch.utils.data.Dataset):
|
||||
def __init__(self, folder, signature, resolutions=[512], frames_file='', device='cuda', use_cache=False, is_cape=False):
|
||||
self.folder = folder
|
||||
self.device = device
|
||||
self.is_cape = is_cape
|
||||
@@ -125,13 +140,20 @@ class Dataset(Dataset):
|
||||
self.resolutions = resolutions
|
||||
|
||||
self.transforms = tv.transforms.ToTensor()
|
||||
self.frames = read_frames(folder, frames_file=frames_file, is_cape=is_cape, shuffle=True)
|
||||
self.frames = read_frames_file(os.path.join(folder, frames_file), shuffle=True)
|
||||
|
||||
self.body_faces = np.load('smpl/smpl_faces.npy')
|
||||
self.static_cache = StaticCache(folder, signature, is_cape=is_cape, resolutions=resolutions, device=device)
|
||||
|
||||
self.cache = dict() if use_cache else None
|
||||
|
||||
ids = set([frame.id for frame in self.frames])
|
||||
self.pose_cache = dict()
|
||||
for id in ids:
|
||||
self.pose_cache[id] = np.load(os.path.join(self.folder, 'pose', f'{id}.npy'))
|
||||
self.shape_cache = dict()
|
||||
for id in ids:
|
||||
self.shape_cache[id] = np.load(os.path.join(self.folder, 'shape', f'{id}.npy'))
|
||||
|
||||
|
||||
def __len__(self):
|
||||
return len(self.frames)
|
||||
@@ -142,76 +164,59 @@ class Dataset(Dataset):
|
||||
if self.cache is not None:
|
||||
if frame in self.cache.keys():
|
||||
data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in self.cache[frame].items()}
|
||||
data = { **data, **self.static_cache[frame.id] }
|
||||
data = { **data, **self.static_cache[frame] }
|
||||
return data
|
||||
|
||||
separator = '_'
|
||||
if self.is_cape:
|
||||
separator = '.'
|
||||
|
||||
uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient_2'
|
||||
if self.is_cape:
|
||||
uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient'
|
||||
|
||||
item = dict()
|
||||
mesh_mask = self.static_cache[frame]['mesh_mask'].cpu().numpy().astype(bool)
|
||||
|
||||
item['id'] = frame.id
|
||||
item['frame_nr'] = frame.frame_nr
|
||||
item['garment'] = frame.garment
|
||||
|
||||
if 'id' in self.signature:
|
||||
item['id'] = frame.id
|
||||
if 'frame_nr' in self.signature:
|
||||
item['frame_nr'] = frame.frame_nr
|
||||
if 'pose_shape' in self.signature:
|
||||
pose_shape = np.load(os.path.join(self.folder, 'pose_shape', frame.id + separator + frame.frame_nr + '.npy')).astype(np.float32)
|
||||
pose_shape[:3] = 0
|
||||
item['pose_shape'] = pose_shape
|
||||
pose = self.pose_cache[frame.id][frame.frame_nr]
|
||||
pose[:3] = 0
|
||||
item['pose_shape'] = np.concatenate((pose, self.shape_cache[frame.id]), axis=0)
|
||||
if 'mesh_body_posed' in self.signature:
|
||||
mesh_posed_body_no_orient = np.load(os.path.join(self.folder, 'mesh_posed_body_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
|
||||
item['mesh_body_posed'] = mesh_posed_body_no_orient
|
||||
mesh_body_posed_no_orient = readPC2Frame(os.path.join(self.folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), frame.frame_nr, True)
|
||||
item['mesh_body_posed'] = mesh_body_posed_no_orient
|
||||
if 'mesh_posed' in self.signature:
|
||||
mesh_posed = np.load(os.path.join(self.folder, 'mesh_posed_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed_no_orient'].astype(np.float32)
|
||||
item['mesh_posed'] = mesh_posed
|
||||
mesh_posed = readPC2Frame(os.path.join(self.folder, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
|
||||
item['mesh_posed'] = np.zeros((14475, 3), dtype=np.float32)
|
||||
item['mesh_posed'][mesh_mask] = mesh_posed
|
||||
if 'mesh_unposed' in self.signature:
|
||||
mesh_unposed_no_shape = np.load(os.path.join(self.folder, 'mesh_unposed_no_shape', frame.id + separator + frame.frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
|
||||
item['mesh_unposed'] = mesh_unposed_no_shape
|
||||
if 'mesh_curvature' in self.signature:
|
||||
mesh_curvature = np.load(os.path.join(self.folder, 'mesh_curvature', frame.id + separator + frame.frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
|
||||
item['mesh_curvature'] = mesh_curvature.astype(np.int32)
|
||||
mesh_unposed_no_shape = readPC2Frame(os.path.join(self.folder, 'mesh_unposed_no_shape', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
|
||||
item['mesh_unposed'] = np.zeros((14475, 3), dtype=np.float32)
|
||||
item['mesh_unposed'][mesh_mask] = mesh_unposed_no_shape
|
||||
if 'mesh_normals' in self.signature:
|
||||
mesh_normals = np.load(os.path.join(self.folder, 'mesh_normals', frame.id + separator + frame.frame_nr + '.npz')); mesh_normals = mesh_normals['mesh_normals']
|
||||
item['mesh_normals'] = mesh_normals
|
||||
mesh_normals = readPC2Frame(os.path.join(self.folder, 'mesh_normals', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
|
||||
item['mesh_normals'] = np.zeros((14475, 3), dtype=np.float32)
|
||||
item['mesh_normals'][mesh_mask] = mesh_normals
|
||||
|
||||
for res in self.resolutions:
|
||||
res_string = ''
|
||||
if res != 512:
|
||||
res_string = '_' + str(res)
|
||||
|
||||
if 'uv_curvature' + res_string in self.signature:
|
||||
uv_curvature = np.load(os.path.join(self.folder, 'uv_curvature' + res_string, frame.id + separator + frame.frame_nr + '.npz'));
|
||||
if res_string != '':
|
||||
uv_curvature = uv_curvature['uv_curvature']
|
||||
else:
|
||||
uv_curvature = uv_curvature['curvature_uv']
|
||||
uv_curvature = torch.nn.functional.one_hot(torch.tensor(uv_curvature, device=self.device, dtype=torch.long), 7).float().to(self.device)
|
||||
item['uv_curvature' + res_string] = uv_curvature
|
||||
if 'uv_normals' + res_string in self.signature:
|
||||
uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_normals = uv_normals['uv_normals']
|
||||
item['uv_normals' + res_string] = uv_normals
|
||||
|
||||
if 'uv_unposed' + res_string in self.signature:
|
||||
uv_unposed_no_shape = np.load(os.path.join(self.folder, 'uv_unposed_no_shape' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']
|
||||
uv_unposed_no_shape = np.load(os.path.join(self.folder, 'uv_unposed_no_shape' + res_string, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed']
|
||||
item['uv_unposed' + res_string] = uv_unposed_no_shape
|
||||
|
||||
if 'uv_body_posed' + res_string in self.signature:
|
||||
uv_posed_body_no_orient = np.load(os.path.join(self.folder, uv_posed_body_no_orient_folder_name + res_string, frame.id + separator + frame.frame_nr + '.npz'))
|
||||
uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']
|
||||
item['uv_body_posed' + res_string] = uv_posed_body_no_orient
|
||||
uv_body_posed_no_orient = np.load(os.path.join(self.folder, 'uv_body_posed_no_orient' + res_string, f'{frame.id}_{frame.frame_nr}.npz')); uv_body_posed_no_orient = uv_body_posed_no_orient['uv_body_posed']
|
||||
item['uv_body_posed' + res_string] = uv_body_posed_no_orient
|
||||
|
||||
if 'uv_posed' + res_string in self.signature:
|
||||
uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient']
|
||||
uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_posed = uv_posed['uv_posed']
|
||||
item['uv_posed' + res_string] = uv_posed
|
||||
|
||||
if 'uv_normals' + res_string in self.signature:
|
||||
uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_normals = uv_normals['uv_normals']
|
||||
item['uv_normals' + res_string] = uv_normals
|
||||
|
||||
|
||||
data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in item.items()}
|
||||
data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor or type(value) == int) else value) for key, value in item.items()}
|
||||
|
||||
for key, value in data.items():
|
||||
if key.startswith('uv_') and (value.shape[-1] == 3 or value.shape[-1] == 7):
|
||||
@@ -221,38 +226,101 @@ class Dataset(Dataset):
|
||||
self.cache[frame] = data
|
||||
|
||||
data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in data.items()}
|
||||
data = { **data, **self.static_cache[frame.id] }
|
||||
data = { **data, **self.static_cache[frame] }
|
||||
return data
|
||||
|
||||
def Dataset(folder, signature, resolutions=[512], frames_file='', device='cuda', use_cache=False, is_cape=False):
|
||||
if not is_cape:
|
||||
return DatasetCloth3D(folder, signature, resolutions=resolutions, frames_file=frames_file, device=device, use_cache=use_cache, is_cape=False)
|
||||
else:
|
||||
return DatasetCape(folder, signature, resolutions=resolutions, frames_file=frames_file, device=device, use_cache=use_cache, is_cape=True)
|
||||
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
from torch.utils.data import DataLoader
|
||||
import plotly.graph_objects as go
|
||||
import plotly.express as ex
|
||||
from plotly.subplots import make_subplots
|
||||
from custom_utils.custom_utils import _maskFaces
|
||||
|
||||
F = np.load('smpl/smpl_faces.npy')
|
||||
F_skirt = np.load('smpl/smpl_faces_skirt.npy')
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
signature = { 'id',
|
||||
'frame_nr',
|
||||
'pose_shape',
|
||||
signature = { 'pose_shape',
|
||||
'uv_posed',
|
||||
'uv_body_posed',
|
||||
'mesh_body_posed',
|
||||
'mesh_posed',
|
||||
'mesh_unposed',
|
||||
'uv_unposed',
|
||||
'uv_curvature',
|
||||
'mesh_curvature',
|
||||
'uv_normals',
|
||||
'mesh_normals',
|
||||
'uv_static',
|
||||
'mesh_static',
|
||||
'mesh_mask',
|
||||
'uv_mask',
|
||||
'uv_shaped_body',
|
||||
}
|
||||
|
||||
dataset = Dataset('/data-local/127/train', signature, resolutions=[512], device=device, frames_file='frames_jumpsuit.txt')
|
||||
dataset = Dataset('/data-local/127/cloth3d_processed', signature, resolutions=[512], device=device, frames_file='frames_train.txt')
|
||||
train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
|
||||
|
||||
for batch in train_dataloader:
|
||||
batch
|
||||
b = 0
|
||||
mesh_mask = batch['mesh_mask'][b].cpu().numpy()
|
||||
faces = F_skirt if batch['garment'][b] in ['Skirt', 'Dress'] else F
|
||||
_, faces = _maskFaces(np.zeros((14475, 3)), faces, mesh_mask)
|
||||
|
||||
mesh_static = batch['mesh_static'][b].cpu().numpy()[mesh_mask]
|
||||
mesh_posed = batch['mesh_posed'][b].cpu().numpy()[mesh_mask]
|
||||
mesh_unposed = batch['mesh_unposed'][b].cpu().numpy()[mesh_mask]
|
||||
mesh_body_posed = batch['mesh_body_posed'][b].cpu().numpy()
|
||||
|
||||
fig = make_subplots(rows=2, cols=2,
|
||||
specs=[[{'type': 'mesh3d'}, {'type': 'mesh3d'}],
|
||||
[{'type': 'mesh3d'}, {'type': 'mesh3d'}]])
|
||||
|
||||
fig.add_trace(go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,2], z=mesh_static[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=1)
|
||||
fig.add_trace(go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=2)
|
||||
fig.add_trace(go.Mesh3d(x=mesh_unposed[:,0], y=mesh_unposed[:,2], z=mesh_unposed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=2, col=1)
|
||||
fig.add_trace(go.Mesh3d(x=mesh_body_posed[:,0], y=mesh_body_posed[:,2], z=mesh_body_posed[:,1], i=F[:,0], j=F[:,1], k=F[:,2], opacity=1), row=2, col=2)
|
||||
|
||||
fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], autorange=False),
|
||||
yaxis=dict(range=[-1, 1], autorange=False),
|
||||
zaxis=dict(range=[-1, 1], autorange=False),
|
||||
aspectmode='cube'),
|
||||
scene2=dict(xaxis=dict(range=[-1, 1], autorange=False),
|
||||
yaxis=dict(range=[-1, 1], autorange=False),
|
||||
zaxis=dict(range=[-1, 1], autorange=False),
|
||||
aspectmode='cube'),
|
||||
scene3=dict(xaxis=dict(range=[-1, 1], autorange=False),
|
||||
yaxis=dict(range=[-1, 1], autorange=False),
|
||||
zaxis=dict(range=[-1, 1], autorange=False),
|
||||
aspectmode='cube'),
|
||||
scene4=dict(xaxis=dict(range=[-1, 1], autorange=False),
|
||||
yaxis=dict(range=[-1, 1], autorange=False),
|
||||
zaxis=dict(range=[-1, 1], autorange=False),
|
||||
aspectmode='cube'))
|
||||
|
||||
fig.show()
|
||||
|
||||
|
||||
uv_mask = batch['uv_mask'][b].cpu().numpy()
|
||||
uv_mask = np.repeat(uv_mask[np.newaxis, :, :], 3, axis=0)
|
||||
uv_static = (batch['uv_static'][b].cpu().numpy() + 1) / 2
|
||||
uv_posed = (batch['uv_posed'][b].cpu().numpy() + 1) / 2
|
||||
uv_unposed = (batch['uv_unposed'][b].cpu().numpy() + 1) / 2
|
||||
uv_body_posed = (batch['uv_body_posed'][b].cpu().numpy() + 1) / 2
|
||||
uv_normals = (batch['uv_normals'][b].cpu().numpy() + 1.5) / 3
|
||||
|
||||
img = np.concatenate([uv_mask, uv_static, uv_posed, uv_unposed, uv_body_posed, uv_normals], axis=2)
|
||||
img = np.transpose(img, (1, 2, 0))
|
||||
ex.imshow(img).show()
|
||||
|
||||
pass
|
||||
|
||||
pass
|
||||
@@ -281,7 +349,7 @@ if __name__ == '__main__':
|
||||
# for i, batch in enumerate(dataset):
|
||||
# print(i)
|
||||
# batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
|
||||
# ex.imshow(batch['uv_curvature'][0].numpy().argmax(-1) / 6.0).show()
|
||||
# ex.imshow(batch['uv_normals'][0].numpy().argmax(-1) / 6.0).show()
|
||||
# sum_uv += batch['uv_unposed'].numpy().sum(axis=0)
|
||||
# sum_mask += batch['uv_mask'].numpy().sum(axis=0)
|
||||
|
||||
|
||||
@@ -0,0 +1,306 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
# os.environ["CUDA_VISIBLE_DEVICES"] = '6'
|
||||
|
||||
import math
|
||||
import time
|
||||
import torch
|
||||
import random
|
||||
import pytorch3d
|
||||
import torchvision as tv
|
||||
import numpy as np
|
||||
import scipy.io as sio
|
||||
import plotly.express as ex
|
||||
import torchvision
|
||||
|
||||
from torchvision.transforms import Resize
|
||||
from torchvision.transforms.functional import InterpolationMode
|
||||
|
||||
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
||||
|
||||
from data.utils import Frame
|
||||
|
||||
|
||||
def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=None):
|
||||
if frames_file is None:
|
||||
uv_static_folder = os.path.join(folder, 'pose_shape')
|
||||
frames = sorted(os.listdir(uv_static_folder))
|
||||
|
||||
if is_cape:
|
||||
frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0]
|
||||
frames = [Frame(frame.split('.')[0], frame.split('.')[1]) for frame in frames]
|
||||
else:
|
||||
frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
|
||||
frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames]
|
||||
|
||||
# if testing:
|
||||
# frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') ))
|
||||
# frames = [frame.split('.')[0] + '_4' for frame in frames]
|
||||
|
||||
|
||||
if frames_file is not None:
|
||||
frames = []
|
||||
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 shuffle:
|
||||
random.shuffle(frames)
|
||||
|
||||
return frames
|
||||
|
||||
|
||||
class StaticCache():
|
||||
def __init__(self, folder, signature, resolutions=[512], is_cape=False, device='cuda') -> None:
|
||||
self.store = dict()
|
||||
self.signature = signature
|
||||
self.folder = folder
|
||||
assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
|
||||
self.resolutions = resolutions
|
||||
self.filled = ''
|
||||
self.is_cape = is_cape
|
||||
self.device = device
|
||||
|
||||
if is_cape:
|
||||
self.uv_masks = dict()
|
||||
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_EXACT)(uv_mask[None,:,:])
|
||||
down = torch.squeeze(down).to(torch.float32)
|
||||
self.uv_masks[res] = down.cpu().numpy()
|
||||
|
||||
|
||||
def read(self, id):
|
||||
data = dict()
|
||||
|
||||
for res in self.resolutions:
|
||||
res_string = ''
|
||||
if res != 512:
|
||||
res_string = '_' + str(res)
|
||||
|
||||
if 'uv_static' + res_string in self.signature:
|
||||
uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; data['uv_static' + res_string] = uv_static
|
||||
if 'uv_mask' + res_string in self.signature:
|
||||
if self.is_cape:
|
||||
data['uv_mask' + res_string] = self.uv_masks[res]
|
||||
else:
|
||||
uv_mask = np.load(os.path.join(self.folder, 'uv_mask' + res_string, id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask
|
||||
|
||||
if 'mesh_static' in self.signature:
|
||||
mesh_static = np.load(os.path.join(self.folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static.astype(np.float32)
|
||||
if 'mesh_mask' in self.signature:
|
||||
if self.is_cape:
|
||||
data['mesh_mask'] = np.ones((3638), np.float32)
|
||||
else:
|
||||
mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask
|
||||
# if 'uv_shaped_body' in self.signature:
|
||||
# uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; data['uv_shaped_body'] = uv_shaped_body
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def __getitem__(self, id) -> dict:
|
||||
if id not in self.store:
|
||||
item = self.read(id)
|
||||
stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()}
|
||||
for key, value in stored_data.items():
|
||||
if key.startswith('uv_') and value.shape[-1] == 3:
|
||||
stored_data[key] = torch.permute(value, [2,0,1])
|
||||
self.store[id] = stored_data
|
||||
|
||||
data = { key: value.to(self.device) for key, value in self.store[id].items()}
|
||||
return data
|
||||
|
||||
|
||||
class DatasetCape(torch.utils.data.Dataset):
|
||||
def __init__(self, folder, signature, resolutions=[512], frames_file='frames_all.txt', device='cuda', use_cache=False, is_cape=False):
|
||||
self.folder = folder
|
||||
self.device = device
|
||||
self.is_cape = is_cape
|
||||
self.signature = signature
|
||||
self.resolutions = resolutions
|
||||
|
||||
self.transforms = tv.transforms.ToTensor()
|
||||
self.frames = read_frames(folder, frames_file=frames_file, is_cape=is_cape, shuffle=True)
|
||||
|
||||
self.body_faces = np.load('smpl/smpl_faces.npy')
|
||||
self.static_cache = StaticCache(folder, signature, is_cape=is_cape, resolutions=resolutions, device=device)
|
||||
|
||||
self.cache = dict() if use_cache else None
|
||||
|
||||
|
||||
def __len__(self):
|
||||
return len(self.frames)
|
||||
|
||||
|
||||
def __getitem__(self, idx):
|
||||
frame : Frame = self.frames[idx]
|
||||
if self.cache is not None:
|
||||
if frame in self.cache.keys():
|
||||
data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in self.cache[frame].items()}
|
||||
data = { **data, **self.static_cache[frame.id] }
|
||||
return data
|
||||
|
||||
separator = '_'
|
||||
if self.is_cape:
|
||||
separator = '.'
|
||||
|
||||
uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient_2'
|
||||
if self.is_cape:
|
||||
uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient'
|
||||
|
||||
item = dict()
|
||||
|
||||
if 'id' in self.signature:
|
||||
item['id'] = frame.id
|
||||
if 'frame_nr' in self.signature:
|
||||
item['frame_nr'] = frame.frame_nr
|
||||
if 'pose_shape' in self.signature:
|
||||
pose_shape = np.load(os.path.join(self.folder, 'pose_shape', frame.id + separator + frame.frame_nr + '.npy')).astype(np.float32)
|
||||
pose_shape[:3] = 0
|
||||
item['pose_shape'] = pose_shape
|
||||
if 'mesh_body_posed' in self.signature:
|
||||
mesh_posed_body_no_orient = np.load(os.path.join(self.folder, 'mesh_posed_body_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
|
||||
item['mesh_body_posed'] = mesh_posed_body_no_orient
|
||||
if 'mesh_posed' in self.signature:
|
||||
mesh_posed = np.load(os.path.join(self.folder, 'mesh_posed_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed_no_orient'].astype(np.float32)
|
||||
item['mesh_posed'] = mesh_posed
|
||||
if 'mesh_unposed' in self.signature:
|
||||
mesh_unposed_no_shape = np.load(os.path.join(self.folder, 'mesh_unposed_no_shape', frame.id + separator + frame.frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
|
||||
item['mesh_unposed'] = mesh_unposed_no_shape
|
||||
if 'mesh_curvature' in self.signature:
|
||||
mesh_curvature = np.load(os.path.join(self.folder, 'mesh_curvature', frame.id + separator + frame.frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
|
||||
item['mesh_curvature'] = mesh_curvature.astype(np.int32)
|
||||
if 'mesh_normals' in self.signature:
|
||||
mesh_normals = np.load(os.path.join(self.folder, 'mesh_normals', frame.id + separator + frame.frame_nr + '.npz')); mesh_normals = mesh_normals['mesh_normals']
|
||||
item['mesh_normals'] = mesh_normals
|
||||
|
||||
for res in self.resolutions:
|
||||
res_string = ''
|
||||
if res != 512:
|
||||
res_string = '_' + str(res)
|
||||
|
||||
if 'uv_curvature' + res_string in self.signature:
|
||||
uv_curvature = np.load(os.path.join(self.folder, 'uv_curvature' + res_string, frame.id + separator + frame.frame_nr + '.npz'));
|
||||
if res_string != '':
|
||||
uv_curvature = uv_curvature['uv_curvature']
|
||||
else:
|
||||
uv_curvature = uv_curvature['curvature_uv']
|
||||
uv_curvature = torch.nn.functional.one_hot(torch.tensor(uv_curvature, device=self.device, dtype=torch.long), 7).float().to(self.device)
|
||||
item['uv_curvature' + res_string] = uv_curvature
|
||||
|
||||
if 'uv_unposed' + res_string in self.signature:
|
||||
uv_unposed_no_shape = np.load(os.path.join(self.folder, 'uv_unposed_no_shape' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']
|
||||
item['uv_unposed' + res_string] = uv_unposed_no_shape
|
||||
|
||||
if 'uv_body_posed' + res_string in self.signature:
|
||||
uv_posed_body_no_orient = np.load(os.path.join(self.folder, uv_posed_body_no_orient_folder_name + res_string, frame.id + separator + frame.frame_nr + '.npz'))
|
||||
uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']
|
||||
item['uv_body_posed' + res_string] = uv_posed_body_no_orient
|
||||
|
||||
if 'uv_posed' + res_string in self.signature:
|
||||
uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient']
|
||||
item['uv_posed' + res_string] = uv_posed
|
||||
|
||||
if 'uv_normals' + res_string in self.signature:
|
||||
uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_normals = uv_normals['uv_normals']
|
||||
item['uv_normals' + res_string] = uv_normals
|
||||
|
||||
|
||||
data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in item.items()}
|
||||
|
||||
for key, value in data.items():
|
||||
if key.startswith('uv_') and (value.shape[-1] == 3 or value.shape[-1] == 7):
|
||||
data[key] = torch.permute(value, [2,0,1])
|
||||
|
||||
if self.cache is not None:
|
||||
self.cache[frame] = data
|
||||
|
||||
data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in data.items()}
|
||||
data = { **data, **self.static_cache[frame.id] }
|
||||
return data
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
signature = { 'id',
|
||||
'frame_nr',
|
||||
'pose_shape',
|
||||
'uv_posed',
|
||||
'uv_body_posed',
|
||||
'mesh_body_posed',
|
||||
'mesh_posed',
|
||||
'mesh_unposed',
|
||||
'uv_unposed',
|
||||
'uv_curvature',
|
||||
'mesh_curvature',
|
||||
'uv_static',
|
||||
'mesh_static',
|
||||
'mesh_mask',
|
||||
'uv_mask',
|
||||
'uv_shaped_body',
|
||||
}
|
||||
|
||||
dataset = Dataset('/data-local/127/train', signature, resolutions=[512], device=device, frames_file='frames_jumpsuit.txt')
|
||||
train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
|
||||
|
||||
for batch in train_dataloader:
|
||||
batch
|
||||
pass
|
||||
|
||||
pass
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
# if __name__ == '__main__':
|
||||
# os.environ["CUDA_VISIBLE_DEVICES"] = "6"
|
||||
# from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
|
||||
# import plotly.graph_objects as go
|
||||
# import plotly.express as ex
|
||||
# # import cmapy as cm
|
||||
# uv2mesh = UV2Mesh(is_skirt=True)
|
||||
|
||||
# F = np.load("custom_utils/smpl_faces.npy")
|
||||
|
||||
# sum_uv = np.zeros((512, 256, 3), dtype=np.float32)
|
||||
# sum_mask = np.zeros((512, 256), dtype=np.float32)
|
||||
# for i, batch in enumerate(dataset):
|
||||
# print(i)
|
||||
# batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
|
||||
# ex.imshow(batch['uv_curvature'][0].numpy().argmax(-1) / 6.0).show()
|
||||
# sum_uv += batch['uv_unposed'].numpy().sum(axis=0)
|
||||
# sum_mask += batch['uv_mask'].numpy().sum(axis=0)
|
||||
|
||||
# sum_mask[sum_mask == 0] = 1
|
||||
# avg_garment = sum_uv / sum_mask[:, :, np.newaxis]
|
||||
|
||||
# # x = np.transpose(avg_garment, [1,2,0])
|
||||
# x = uv2mesh(avg_garment[tf.newaxis,:,:,:])[0]
|
||||
# go.Figure(data=[
|
||||
# go.Mesh3d(
|
||||
# # 8 vertices of a cube
|
||||
# x=x[:,0],
|
||||
# y=x[:,1],
|
||||
# z=x[:,2],
|
||||
# # i, j and k give the vertices of triangles
|
||||
# i = F[:,0],
|
||||
# j = F[:,1],
|
||||
# k = F[:,2]
|
||||
# )]).show()
|
||||
|
||||
# np.save('custom_data/train_dress_30k_no_shape_avg_uv.npy', avg_garment)
|
||||
|
||||
# print()
|
||||
@@ -0,0 +1,223 @@
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import tqdm
|
||||
import cv2
|
||||
import torch
|
||||
import scipy.interpolate as spi
|
||||
import numpy as np
|
||||
import plotly.graph_objects as go
|
||||
import plotly.express as ex
|
||||
from torchvision.transforms import Resize
|
||||
from torchvision.transforms.functional import InterpolationMode
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "4"
|
||||
|
||||
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
||||
from data.utils import Frame
|
||||
|
||||
|
||||
def fill(img, mask):
|
||||
coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0))
|
||||
img_idx = np.transpose(np.stack(np.nonzero((mask))), (1,0))
|
||||
img_data = img[img_idx[:,0], img_idx[:,1], :]
|
||||
|
||||
interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data)
|
||||
# interpolator = spi.RBFInterpolator(y=img_idx, d=img_data, neighbors=5, kernel='linear')
|
||||
|
||||
img_2 = interpolator(coords)
|
||||
res = np.zeros((512,256,3), dtype=np.float32)
|
||||
res[coords[:,0], coords[:,1], :] = img_2
|
||||
res[img_idx[:,0], img_idx[:,1], :] = img_data
|
||||
|
||||
return res
|
||||
|
||||
|
||||
def fill_label(img, mask):
|
||||
coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0))
|
||||
img_idx = np.transpose(np.stack(np.nonzero((mask))), (1,0))
|
||||
img_data = img[img_idx[:,0], img_idx[:,1]]
|
||||
|
||||
interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data)
|
||||
img_2 = interpolator(coords)
|
||||
res = np.zeros((512,256), dtype=np.float32)
|
||||
res[coords[:,0], coords[:,1]] = img_2
|
||||
res[img_idx[:,0], img_idx[:,1]] = img_data
|
||||
|
||||
return res
|
||||
|
||||
|
||||
def generate_downscaled_uv_map(folder, frame: Frame, size):
|
||||
img = np.load(os.path.join(folder, 'uv_mask', f'{frame.id}_{frame.garment}.npz')); img = img['uv_mask'].astype(np.float32)
|
||||
down = resizers_nearest[size](torch.from_numpy(img[None,:,:]).to(device))
|
||||
down = np.squeeze(down.cpu().numpy().astype(bool))
|
||||
if not os.path.exists(os.path.join(folder, 'uv_mask' + '_' + str(size))):
|
||||
os.makedirs(os.path.join(folder, 'uv_mask' + '_' + str(size)), exist_ok=True)
|
||||
np.savez_compressed(os.path.join(folder, 'uv_mask' + '_' + str(size), f'{frame.id}_{frame.garment}.npz'), uv_mask = down)
|
||||
|
||||
|
||||
class UVMaskStore():
|
||||
def __init__(self, folder, size):
|
||||
self.store = dict()
|
||||
self.folder = folder
|
||||
self.size = size
|
||||
|
||||
def read(self, frame: Frame):
|
||||
path = os.path.join(self.folder, 'uv_mask' + ('_' + str(self.size) if self.size != 512 else ''), f'{frame.id}_{frame.garment}.npz')
|
||||
if not os.path.exists(path):
|
||||
generate_downscaled_uv_map(self.folder, frame, self.size)
|
||||
|
||||
mask = np.load(path)
|
||||
mask = mask['uv_mask']
|
||||
return mask
|
||||
|
||||
def __getitem__(self, frame: Frame) -> dict:
|
||||
if (frame.id, frame.garment) not in self.store:
|
||||
self.store[(frame.id, frame.garment)] = self.read(frame)
|
||||
|
||||
return self.store[(frame.id, frame.garment)]
|
||||
|
||||
|
||||
def downscale_uv_maps_posed_body(folder, type, sizes):
|
||||
for size in sizes:
|
||||
os.makedirs(os.path.join(folder, type + str(size)), exist_ok=True)
|
||||
|
||||
frames = sorted(os.listdir(os.path.join(folder, type)))[int(args.start):int(args.end)]
|
||||
frames = [Frame(frame.split('_')[0], frame.split('_')[1].split('.')[0], '') for frame in frames]
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
img = np.load(os.path.join(folder, type, f'{frame.id}_{frame.frame_nr}.npz'))['uv_body_posed']
|
||||
img = fill(img, uv_masks_body[512])
|
||||
for size in sizes:
|
||||
mask = uv_masks_body[size]
|
||||
down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
|
||||
down = down * mask[None,:,:]
|
||||
np.savez_compressed(os.path.join(folder, type + str(size), f'{frame.id}_{frame.frame_nr}.npz'), uv_body_posed = down)
|
||||
|
||||
|
||||
def downscale_uv_maps_posed_no_orient(folder, type, sizes):
|
||||
for size in sizes:
|
||||
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
|
||||
|
||||
frames = sorted(os.listdir(os.path.join(folder, type)))[int(args.start):int(args.end)]
|
||||
frames = [Frame(frame.split('_')[0], frame.split('_')[1], frame.split('_')[2].split('.')[0]) for frame in frames]
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
img = np.load(os.path.join(folder, type, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))['uv_posed']
|
||||
img = fill(img, uv_mask_store[512][frame])
|
||||
for size in sizes:
|
||||
mask = uv_mask_store[size][frame]
|
||||
down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
|
||||
down = down * mask[None,:,:]
|
||||
np.savez_compressed(os.path.join(folder, type + '_' + str(size), f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'), uv_posed = down)
|
||||
|
||||
|
||||
def downscale_uv_maps_normals(folder, type, sizes):
|
||||
for size in sizes:
|
||||
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
|
||||
|
||||
frames = sorted(os.listdir(os.path.join(folder, type)))[int(args.start):int(args.end)]
|
||||
frames = [Frame(frame.split('_')[0], frame.split('_')[1], frame.split('_')[2].split('.')[0]) for frame in frames]
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
img = np.load(os.path.join(folder, type, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))['uv_normals']
|
||||
img = fill(img, uv_mask_store[512][frame])
|
||||
for size in sizes:
|
||||
mask = uv_mask_store[size][frame]
|
||||
down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
|
||||
down = down * mask[None,:,:]
|
||||
np.savez_compressed(os.path.join(folder, type + '_' + str(size), f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'), uv_normals = down)
|
||||
|
||||
|
||||
def downscale_uv_maps_unposed(folder, type, sizes):
|
||||
for size in sizes:
|
||||
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
|
||||
|
||||
frames = sorted(os.listdir(os.path.join(folder, type)))[int(args.start):int(args.end)]
|
||||
frames = [Frame(frame.split('_')[0], frame.split('_')[1], frame.split('_')[2].split('.')[0]) for frame in frames]
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
img = np.load(os.path.join(folder, type, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))['uv_unposed']
|
||||
img = fill(img, uv_mask_store[512][frame])
|
||||
for size in sizes:
|
||||
mask = uv_mask_store[size][frame]
|
||||
down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
|
||||
down = down * mask[None,:,:]
|
||||
np.savez_compressed(os.path.join(folder, type + '_' + str(size), f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'), uv_unposed = down)
|
||||
|
||||
|
||||
def downscale_uv_maps_static(folder, type, sizes):
|
||||
for size in sizes:
|
||||
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
|
||||
|
||||
frames = sorted(os.listdir(os.path.join(folder, type)))
|
||||
import random
|
||||
random.shuffle(frames)
|
||||
frames = [(frame.split('_')[0], frame.split('_')[1].split('.')[0]) for frame in frames]
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
img = np.load(os.path.join(folder, type, f'{frame[0]}_{frame[1]}.npz'))['uv_static']
|
||||
img = fill(img, uv_mask_store[512][Frame(frame[0], 0, frame[1])])
|
||||
for size in sizes:
|
||||
mask = uv_mask_store[size][Frame(frame[0], 0, frame[1])]
|
||||
down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
|
||||
down = down * mask[None,:,:]
|
||||
np.savez_compressed(os.path.join(folder, type + '_' + str(size), f'{frame[0]}_{frame[1]}.npz'), uv_static = down)
|
||||
|
||||
|
||||
def downscale_uv_maps_mask(folder, type, size):
|
||||
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
|
||||
|
||||
frames = sorted(os.listdir(os.path.join(folder, type)))[int(args.start):int(args.end)]
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
id, _ = frame.split('.')
|
||||
img = np.load(os.path.join(folder, type, frame)); img = img['uv_mask'].astype(np.float32)
|
||||
down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
|
||||
down = np.squeeze(down.numpy().astype(bool))
|
||||
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_mask = down)
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127/cloth3d_processed')
|
||||
parser.add_argument("-s", "--start", help = "start of range", default=0)
|
||||
parser.add_argument("-e", "--end", help = "end of range", default=25000)
|
||||
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
|
||||
parser.add_argument("-d", "--debug", help = "end of range", default=False, action='store_true')
|
||||
args = parser.parse_args()
|
||||
|
||||
# sizes = [32, 64, 128, 256]
|
||||
device = 'cuda'
|
||||
sizes = [32, 64, 128, 256]
|
||||
folder = args.folder
|
||||
resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BICUBIC).to(device) for res in sizes }
|
||||
resizers_nearest = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST_EXACT).to(device) for res in sizes }
|
||||
|
||||
uv_masks_body = dict()
|
||||
uv_mask_body = np.load('custom_utils/mask_256x512.npy')
|
||||
uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
|
||||
uv_masks_body[512] = uv_mask_body.astype(bool)
|
||||
for res in [32, 64, 128, 256]:
|
||||
down = resizers_nearest[res](torch.from_numpy(uv_mask_body[None,:,:].astype(np.int32)).to(device))
|
||||
down = np.squeeze(down.cpu().numpy().astype(bool))
|
||||
uv_masks_body[res] = down
|
||||
|
||||
|
||||
uv_mask_store = { size: UVMaskStore(folder, size=size) for size in sizes }
|
||||
uv_mask_store[512] = UVMaskStore(folder, size=512)
|
||||
|
||||
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=32)
|
||||
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=64)
|
||||
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=128)
|
||||
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=256)
|
||||
# downscale_uv_maps_static(folder=folder, type='uv_static', sizes=sizes)
|
||||
|
||||
downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes)
|
||||
downscale_uv_maps_posed_no_orient(folder=folder, type='uv_posed_no_orient', sizes=sizes)
|
||||
downscale_uv_maps_normals(folder=folder, type='uv_normals', sizes=sizes)
|
||||
downscale_uv_maps_posed_body(folder=folder, type='uv_body_posed_no_orient', sizes=sizes)
|
||||
|
||||
|
||||
print()
|
||||
+208
-12
@@ -6,8 +6,8 @@ parser = argparse.ArgumentParser()
|
||||
# parser.add_argument("-f", "--folder", help = "end of range")
|
||||
parser.add_argument("-s", "--start", help = "start of range", type=int, default=0)
|
||||
parser.add_argument("-e", "--end", help = "end of range", type=int, default=1)
|
||||
parser.add_argument("-g", "--gpu", help = "gpu id")
|
||||
# parser.add_argument("-n", "--name", help = "end of range")
|
||||
parser.add_argument("-g", "--gpu", help = "gpu id", default='0')
|
||||
parser.add_argument("-d", "--debug", help = "end of range", default=False, action='store_true')
|
||||
args = parser.parse_args()
|
||||
|
||||
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
|
||||
@@ -30,7 +30,7 @@ import torch
|
||||
from scipy.spatial.transform import Rotation as R
|
||||
|
||||
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
||||
from data.cloth3d_io import readPC2, loadInfo, writePC2, readOBJ, quads2tris
|
||||
from data.cloth3d_io import readPC2, loadInfo, writePC2, readOBJ, quads2tris, readPC2Frame
|
||||
from data.utils import Frame
|
||||
from smpl.smpl_torch import SMPL
|
||||
from custom_utils.custom_utils import _maskFaces
|
||||
@@ -160,7 +160,7 @@ def deorient(verts, pose):
|
||||
|
||||
|
||||
def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder):
|
||||
ids = os.listdir(folder)
|
||||
ids = sorted(os.listdir(folder))
|
||||
os.makedirs(os.path.join(out_folder, 'mesh_posed_no_orient'), exist_ok=True)
|
||||
os.makedirs(os.path.join(out_folder, 'mesh_mask'), exist_ok=True)
|
||||
|
||||
@@ -179,12 +179,18 @@ def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder):
|
||||
continue
|
||||
|
||||
verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
|
||||
# if info['poses'].shape[1] == verts.shape[0]:
|
||||
# continue
|
||||
reg_data = sio.loadmat(os.path.join(folder_reg, id, garment + '_weights_fixed.mat'))
|
||||
verts = (verts[:, reg_data['F'], :] * reg_data['W'][None, :, :, None]).sum(axis=2)
|
||||
mask_idx = sio.loadmat(os.path.join(folder_reg, id, garment + '_mask.mat'))['mask_idx'].reshape((-1)) - 1
|
||||
mask = np.zeros((14475), dtype=bool)
|
||||
mask[mask_idx] = True
|
||||
|
||||
if info['poses'].shape[1] != verts.shape[0]:
|
||||
assert False, 'Somethign wrong here'
|
||||
info['poses'] = info['poses'][:, :verts.shape[0]]
|
||||
|
||||
verts = deorient(verts, info['poses'])
|
||||
|
||||
# # plot as scatter plot with plotly
|
||||
@@ -193,13 +199,12 @@ def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder):
|
||||
# fig = go.Figure(data=[go.Scatter3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], mode='markers', marker=dict(size=1))])
|
||||
# fig.show()
|
||||
|
||||
|
||||
writePC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), verts, True)
|
||||
np.save(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy'), mask)
|
||||
|
||||
|
||||
def create_pose_shape(folder, out_folder):
|
||||
ids = os.listdir(folder)
|
||||
ids = sorted(os.listdir(folder))
|
||||
os.makedirs(os.path.join(out_folder, 'pose'), exist_ok=True)
|
||||
os.makedirs(os.path.join(out_folder, 'shape'), exist_ok=True)
|
||||
|
||||
@@ -216,7 +221,7 @@ def create_pose_shape(folder, out_folder):
|
||||
|
||||
|
||||
def create_mesh_static(folder, out_folder):
|
||||
ids = os.listdir(folder)
|
||||
ids = sorted(os.listdir(folder))
|
||||
os.makedirs(os.path.join(out_folder, 'mesh_static'), exist_ok=True)
|
||||
|
||||
for id in tqdm(ids):
|
||||
@@ -247,7 +252,7 @@ def create_mesh_static(folder, out_folder):
|
||||
|
||||
|
||||
def create_mesh_body_posed_no_orient(folder, out_folder):
|
||||
ids = os.listdir(folder)
|
||||
ids = sorted(os.listdir(folder))
|
||||
os.makedirs(os.path.join(out_folder, 'mesh_body_posed_no_orient'), exist_ok=True)
|
||||
|
||||
for id in tqdm(ids):
|
||||
@@ -288,10 +293,11 @@ def calculate_normals(vertices, faces):
|
||||
|
||||
|
||||
def create_mesh_normals(folder, out_folder):
|
||||
ids = os.listdir(folder)
|
||||
ids = sorted(os.listdir(folder))
|
||||
os.makedirs(os.path.join(out_folder, 'mesh_normals'), exist_ok=True)
|
||||
# F = np.load('smpl/smpl_faces.npy')
|
||||
# F_skirt = np.load('smpl/smpl_faces_skirt.npy')
|
||||
# ids = ['00438']
|
||||
|
||||
for id in tqdm(ids):
|
||||
garment_dir = os.path.join(folder, id)
|
||||
@@ -306,7 +312,7 @@ def create_mesh_normals(folder, out_folder):
|
||||
continue
|
||||
|
||||
verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
|
||||
verts = deorient(verts, info['poses'])
|
||||
verts = deorient(verts, info['poses'][:, :verts.shape[0]])
|
||||
_, faces, _, _ = readOBJ(os.path.join(garment_dir, garment + '.obj'))
|
||||
faces = quads2tris(faces)
|
||||
normals = calculate_normals(verts, faces.astype(np.int32))
|
||||
@@ -356,9 +362,64 @@ def create_mesh_normals(folder, out_folder):
|
||||
|
||||
pass
|
||||
|
||||
def create_mesh_normals_cape(folder, out_folder):
|
||||
frames = os.listdir(os.path.join(folder, 'mesh_posed'))
|
||||
os.makedirs(os.path.join(folder, 'mesh_normals'), exist_ok=True)
|
||||
F = np.load('smpl/smpl_faces_orig.npy')
|
||||
# F_skirt = np.load('smpl/smpl_faces_skirt.npy')
|
||||
|
||||
for frame in tqdm(frames):
|
||||
id, frame_nr = frame.split('.')[:2]
|
||||
|
||||
|
||||
verts = np.load(os.path.join(folder, 'mesh_posed', f'{id}.{frame_nr}.npz'))['mesh_posed']
|
||||
normals = calculate_normals(verts[None,:,:], F)[0]
|
||||
|
||||
np.savez_compressed(os.path.join(folder, 'mesh_normals', f'{id}.{frame_nr}.npz'), mesh_normals=normals)
|
||||
|
||||
# normals = readPC2(os.path.join(out_folder, 'mesh_normals', f'{id}_{garment}.pc16'), True)['V']
|
||||
# colors = np.clip(np.abs(normals), 0, 1)
|
||||
# verts = (verts[:, reg_data['F'], :] * reg_data['W'][None, :, :, None]).sum(axis=2)
|
||||
# plot as scatter plot with plotly
|
||||
# import plotly.graph_objects as go
|
||||
# frame = 5
|
||||
# fig = go.Figure(data=[go.Scatter3d(
|
||||
# x=verts[:,0],
|
||||
# y=verts[:,2],
|
||||
# z=verts[:,1],
|
||||
# mode='markers',
|
||||
# marker=dict(
|
||||
# size=1,
|
||||
# color=normals[:,:],
|
||||
# colorscale='Viridis',
|
||||
# opacity=0.8
|
||||
# )
|
||||
# )])
|
||||
# moved_verts = verts + 0.05 * normals
|
||||
|
||||
# # Create the scatter plot for the moved vertices
|
||||
# fig.add_trace(go.Scatter3d(
|
||||
# x=moved_verts[:,0],
|
||||
# y=moved_verts[:,2],
|
||||
# z=moved_verts[:,1],
|
||||
# mode='markers',
|
||||
# marker=dict(
|
||||
# size=1,
|
||||
# color='red',
|
||||
# opacity=0.8
|
||||
# )
|
||||
# ))
|
||||
# fig.show()
|
||||
|
||||
# # create a plotly mesh plot for verts and faces
|
||||
# fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)])
|
||||
# fig.show()
|
||||
|
||||
pass
|
||||
|
||||
|
||||
def verify_registration_ans_mask(folder, out_folder):
|
||||
ids = os.listdir(folder)
|
||||
ids = sorted(os.listdir(folder))
|
||||
|
||||
for id in ids:
|
||||
garment_dir = os.path.join(folder, id)
|
||||
@@ -510,9 +571,136 @@ def create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=0, end=1)
|
||||
# fig.show()
|
||||
pass
|
||||
|
||||
|
||||
def create_frames_file(folder, out_folder, split=0.9):
|
||||
ids = os.listdir(folder)
|
||||
|
||||
frames_non_skirt = []
|
||||
frames_skirt = []
|
||||
ids_non_skirt = set()
|
||||
ids_skirt = set()
|
||||
|
||||
for id in tqdm(ids):
|
||||
garment_dir = os.path.join(folder, id)
|
||||
if not os.path.isdir(garment_dir):
|
||||
continue
|
||||
|
||||
info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
|
||||
garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
|
||||
frames = info['poses'].shape[1]
|
||||
|
||||
assert info['poses'].shape[0] == 72
|
||||
|
||||
for garment in garments:
|
||||
for i in range(frames):
|
||||
if garment == 'Skirt' or garment == 'Dress':
|
||||
frames_skirt.append(Frame(id, i, garment))
|
||||
ids_skirt.add(id)
|
||||
else:
|
||||
frames_non_skirt.append(Frame(id, i, garment))
|
||||
ids_non_skirt.add(id)
|
||||
|
||||
nr_train_non_skirt = int(len(ids_non_skirt) * split)
|
||||
nr_train_skirt = int(len(ids_skirt) * split)
|
||||
|
||||
ids_skirt = list(ids_skirt)
|
||||
ids_non_skirt = list(ids_non_skirt)
|
||||
|
||||
with open(os.path.join(out_folder, 'frames_train.txt'), 'w') as f:
|
||||
for frame in frames_non_skirt:
|
||||
if frame.id in ids_non_skirt[:nr_train_non_skirt]:
|
||||
f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
|
||||
|
||||
with open(os.path.join(out_folder, 'frames_train_skirt.txt'), 'w') as f:
|
||||
for frame in frames_skirt:
|
||||
if frame.id in ids_skirt[:nr_train_skirt]:
|
||||
f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
|
||||
|
||||
with open(os.path.join(out_folder, 'frames_valid.txt'), 'w') as f:
|
||||
for frame in frames_non_skirt:
|
||||
if frame.id in ids_non_skirt[nr_train_non_skirt:]:
|
||||
f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
|
||||
|
||||
with open(os.path.join(out_folder, 'frames_valid_skirt.txt'), 'w') as f:
|
||||
for frame in frames_skirt:
|
||||
if frame.id in ids_skirt[nr_train_skirt:]:
|
||||
f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
|
||||
|
||||
|
||||
def create_frames_file_partial(folder, out_folder, split=0.9, part=0.1):
|
||||
ids = os.listdir(folder)
|
||||
|
||||
frames_non_skirt = []
|
||||
frames_skirt = []
|
||||
ids_non_skirt = set()
|
||||
ids_skirt = set()
|
||||
|
||||
for id in tqdm(ids):
|
||||
garment_dir = os.path.join(folder, id)
|
||||
if not os.path.isdir(garment_dir):
|
||||
continue
|
||||
|
||||
info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
|
||||
garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
|
||||
frames = info['poses'].shape[1]
|
||||
|
||||
assert info['poses'].shape[0] == 72
|
||||
|
||||
for garment in garments:
|
||||
for i in range(frames):
|
||||
if garment == 'Skirt' or garment == 'Dress':
|
||||
frames_skirt.append(Frame(id, i, garment))
|
||||
ids_skirt.add(id)
|
||||
else:
|
||||
frames_non_skirt.append(Frame(id, i, garment))
|
||||
ids_non_skirt.add(id)
|
||||
|
||||
nr_train_non_skirt = int(len(ids_non_skirt) * split * part)
|
||||
nr_train_skirt = int(len(ids_skirt) * split * part)
|
||||
|
||||
ids_skirt = list(ids_skirt)
|
||||
ids_non_skirt = list(ids_non_skirt)
|
||||
|
||||
with open(os.path.join(out_folder, f'frames_train_{int(part*100)}p.txt'), 'w') as f:
|
||||
for frame in frames_non_skirt:
|
||||
if frame.id in ids_non_skirt[:nr_train_non_skirt]:
|
||||
f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
|
||||
|
||||
with open(os.path.join(out_folder, f'frames_train_skirt_{int(part*100)}p.txt'), 'w') as f:
|
||||
for frame in frames_skirt:
|
||||
if frame.id in ids_skirt[:nr_train_skirt]:
|
||||
f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
|
||||
|
||||
|
||||
|
||||
def check_irregular_poses(folder, folder_reg, out_folder):
|
||||
frame: Frame = Frame('01512', 50, 'Dress')
|
||||
mesh_posed = readPC2Frame(os.path.join(out_folder, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
|
||||
# mesh_static = np.load(f'{out_folder}/mesh_static/{frame.id}_{frame.garment}.npy')
|
||||
|
||||
mesh_mask = np.load(f'{out_folder}/mesh_mask/{frame.id}_{frame.garment}.npy')
|
||||
mesh_body_posed = readPC2Frame(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), frame.frame_nr, True)
|
||||
|
||||
F = np.load('smpl/smpl_faces_skirt.npy')
|
||||
_, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask)
|
||||
|
||||
fig = go.Figure(data=[go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)])
|
||||
fig.add_trace(go.Mesh3d(x=mesh_body_posed[:,0], y=mesh_body_posed[:,2], z=mesh_body_posed[:,1], i=F[:,0], j=F[:,1], k=F[:,2], color='lightblue', opacity=1))
|
||||
fig.update_layout(scene = dict(
|
||||
xaxis = dict(range=[-1,1],),
|
||||
yaxis = dict(range=[-1,1],),
|
||||
zaxis = dict(range=[-1,1],),),)
|
||||
fig.update_layout(scene_aspectmode='cube')
|
||||
fig.show()
|
||||
pass
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
|
||||
# folder = '/data-net/dataset/cloth3d/test'
|
||||
# folder_reg = '/data-local/127/cloth3d_registration/test'
|
||||
# out_folder = '/data-local/127/cloth3d_processed_test'
|
||||
|
||||
folder = '/data-net/dataset/cloth3d/train'
|
||||
folder_reg = '/data-local/127/cloth3d_registration/train'
|
||||
out_folder = '/data-local/127/cloth3d_processed'
|
||||
@@ -524,10 +712,18 @@ if __name__ == '__main__':
|
||||
# split_data(folder, 0.9)
|
||||
# create_extras_file(folder)
|
||||
|
||||
# check_irregular_poses(folder, folder_reg, out_folder)
|
||||
|
||||
# create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder)
|
||||
# verify_registration_ans_mask(folder, out_folder)
|
||||
# create_pose_shape(folder, out_folder)
|
||||
# create_mesh_static(folder, out_folder)
|
||||
# create_mesh_body_posed_no_orient(folder, out_folder)
|
||||
# create_mesh_normals(folder, out_folder)
|
||||
create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end)
|
||||
# create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end)
|
||||
|
||||
# create_frames_file(folder, out_folder)
|
||||
create_frames_file_partial(folder, out_folder, part=0.5)
|
||||
|
||||
# folder = '/data-local/127_2/cape_processed/test/female'
|
||||
# create_mesh_normals_cape(folder, None)
|
||||
+4
-4
@@ -1,10 +1,10 @@
|
||||
import tables as tb
|
||||
|
||||
class Frame:
|
||||
def __init__(self, id, frame_nr, garment=''):
|
||||
self.id = id
|
||||
self.frame_nr = frame_nr
|
||||
self.garment = garment
|
||||
def __init__(self, id: str, frame_nr: int, garment: str=''):
|
||||
self.id: str = id
|
||||
self.frame_nr: int = frame_nr
|
||||
self.garment: str = garment
|
||||
|
||||
def __hash__(self):
|
||||
return hash((self.id, self.frame_nr, self.garment))
|
||||
|
||||
+10
-4
@@ -1,7 +1,7 @@
|
||||
# Use an existing image as the base
|
||||
FROM nvcr.io/nvidia/pytorch:23.04-py3
|
||||
# FROM nvcr.io/nvidia/pytorch:23.04-py3
|
||||
# FROM pytorch/pytorch:1.13.1-cuda11.6-cudnn8-devel
|
||||
# FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-devel
|
||||
FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel
|
||||
#
|
||||
# Add the current user to the image
|
||||
ARG USER_NAME=lhunor
|
||||
@@ -10,7 +10,7 @@ ARG GROUP_NAME=guest
|
||||
ARG GROUP_ID=11000
|
||||
|
||||
# Update the system and install required packages
|
||||
RUN apt-get update && apt-get upgrade -y && apt-get install -y sudo git apt-utils htop
|
||||
RUN apt-get update && apt-get install -y sudo git apt-utils htop
|
||||
|
||||
# Create a new group with the same GID as the host
|
||||
RUN groupadd --gid $GROUP_ID $GROUP_NAME
|
||||
@@ -25,6 +25,12 @@ RUN echo '%sudo ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers
|
||||
# Add the new user to the sudo group
|
||||
RUN usermod -aG sudo $USER_NAME
|
||||
|
||||
# install conda stuff still as root
|
||||
RUN conda update -n base conda
|
||||
RUN conda install -c fvcore -c iopath -c conda-forge fvcore iopath -y
|
||||
RUN conda install -c bottler nvidiacub -y
|
||||
RUN conda install pytorch3d -c pytorch3d -y
|
||||
|
||||
# Set the default user to the new user
|
||||
USER $USER_NAME
|
||||
|
||||
@@ -33,8 +39,8 @@ WORKDIR /workspace
|
||||
ENV PATH="$PATH:/home/$USER_NAME/.local/bin"
|
||||
RUN pip install --upgrade pip
|
||||
COPY requirements.txt /tmp/
|
||||
RUN pip install --requirement /tmp/requirements.txt
|
||||
RUN pip install torch-cluster -f https://data.pyg.org/whl/torch-2.0.0+cu117.html
|
||||
RUN pip install --requirement /tmp/requirements.txt
|
||||
|
||||
|
||||
# defining the volumes that will be mounted
|
||||
|
||||
@@ -41,3 +41,25 @@ services:
|
||||
reservations:
|
||||
devices:
|
||||
- capabilities: [gpu]
|
||||
|
||||
|
||||
lhunor_iccvw:
|
||||
image: lhunor_iccvw:pt2.0.1
|
||||
container_name: lhunor_iccvw_pt2.0.1
|
||||
ipc: host
|
||||
tty: true
|
||||
ulimits:
|
||||
memlock: -1
|
||||
stack: 67108864
|
||||
volumes:
|
||||
- '/home-net/lhunor:/home-net'
|
||||
- '/data-net/guest/lhunor/data:/data-net'
|
||||
- '/data-fast/126-data1/lhunor/data:/data-local/126'
|
||||
- '/data-fast/127-data1/lhunor/data:/data-local/127'
|
||||
- '/data-fast/127-data2/lhunor/data:/data-local/127_2'
|
||||
- '/data-fast/128-data1/lhunor/data:/data-local/128'
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- capabilities: [gpu]
|
||||
@@ -1,13 +1,10 @@
|
||||
torch==2.0.0
|
||||
numpy
|
||||
matplotlib
|
||||
opencv-python
|
||||
wandb
|
||||
scipy
|
||||
plotly
|
||||
pandas
|
||||
cmapy
|
||||
pytorch3d
|
||||
tables
|
||||
numpy-stl
|
||||
torchmetrics
|
||||
|
||||
+13
-15
@@ -11,10 +11,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.ae_body_cape_m import config
|
||||
from configs.iccvw.ae_body_cape_m 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']:
|
||||
@@ -32,10 +34,7 @@ import plotly.graph_objects as go
|
||||
import plotly.express as ex
|
||||
import torch.optim as optim
|
||||
|
||||
if config['is_skirt']:
|
||||
from losses.losses import edge_loss_skirt as edge_loss
|
||||
else:
|
||||
from losses.losses import edge_loss
|
||||
from losses.losses import edge_loss
|
||||
|
||||
from data.dataset import Dataset
|
||||
# from torch.utils.data import DataLoader
|
||||
@@ -49,6 +48,7 @@ from pytorch_msssim import MS_SSIM
|
||||
|
||||
|
||||
def init_wandb():
|
||||
wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
|
||||
run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
|
||||
artifact = wandb.Artifact('code', type='code')
|
||||
artifact.add_file('scripts/ae_body_train.py')
|
||||
@@ -64,7 +64,7 @@ def init_wandb():
|
||||
artifact.add_file(file_name)
|
||||
run.log_artifact(artifact)
|
||||
|
||||
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
|
||||
uv2mesh = UV2Mesh(res=256, is_skirt=False, is_smpl_orig=config['is_cape'])
|
||||
|
||||
|
||||
|
||||
@@ -121,7 +121,7 @@ def plot_everything(batch: dict, type: str):
|
||||
img = torch.concat((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")
|
||||
F = np.load("smpl/smpl_faces.npy")
|
||||
garment_unposed, faces = mesh_pred.cpu().numpy(), F
|
||||
gt = [mesh_orig[i].cpu().numpy() for i in range(pred.shape[0])]
|
||||
|
||||
@@ -175,9 +175,7 @@ if __name__ == '__main__':
|
||||
if not args['debug']:
|
||||
init_wandb()
|
||||
|
||||
skirt_modifier = '_skirt' if config['is_skirt'] else ''
|
||||
|
||||
uv_mask_body = np.load('custom_utils/mask_256x512' + ('_skirt' if config['is_skirt'] else '') + '.npy')
|
||||
uv_mask_body = np.load('custom_utils/mask_256x512.npy')
|
||||
uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
|
||||
uv_mask_body = torch.from_numpy(uv_mask_body).float().to(config['device'])
|
||||
|
||||
@@ -198,15 +196,15 @@ if __name__ == '__main__':
|
||||
load_model(model.encoder, config['pretrained_model'] + '_enc')
|
||||
load_model(model.decoder, config['pretrained_model'] + '_dec')
|
||||
|
||||
ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_enc', top_n=5)
|
||||
ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_dec', top_n=5)
|
||||
ckpt_saver_enc = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_enc', top_n=5)
|
||||
ckpt_saver_dec = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_dec', top_n=5)
|
||||
|
||||
run_name = config['name']
|
||||
print(f'started training: {run_name}')
|
||||
|
||||
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()
|
||||
|
||||
|
||||
+14
-11
@@ -11,10 +11,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.ae_norm_1 import config
|
||||
from configs.iccvw.ae_norm_cape_m 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']:
|
||||
@@ -49,7 +51,8 @@ from pytorch_msssim import MS_SSIM
|
||||
|
||||
|
||||
def init_wandb():
|
||||
run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127_2/temp', tags=config['tags'])
|
||||
wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
|
||||
run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
|
||||
artifact = wandb.Artifact('code', type='code')
|
||||
artifact.add_file('scripts/ae_norm_train.py')
|
||||
artifact.add_file('losses/losses.py')
|
||||
@@ -70,9 +73,9 @@ uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_
|
||||
|
||||
def train_step(batch):
|
||||
model.train()
|
||||
pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.1
|
||||
pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.4
|
||||
|
||||
ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_normals'] + 1.1))
|
||||
ms_ssim = torch.mean(_MSSSIM(pred + 1.4, batch['uv_normals'] + 1.4))
|
||||
loss = 1e3 * (1 - ms_ssim)
|
||||
|
||||
optimizer.zero_grad()
|
||||
@@ -93,13 +96,13 @@ def valid_step(batch):
|
||||
model.eval()
|
||||
|
||||
with torch.inference_mode():
|
||||
pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.1
|
||||
pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.4
|
||||
|
||||
mesh_pred = uv2mesh(pred)
|
||||
mesh_orig = batch['mesh_normals']
|
||||
|
||||
l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_normals'] - pred)), axis=1))
|
||||
ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_normals'] + 1.1))
|
||||
ms_ssim = torch.mean(_MSSSIM(pred + 1.4, batch['uv_normals'] + 1.4))
|
||||
loss = 1e3 * (1 - ms_ssim)
|
||||
|
||||
metric = torch.mean(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1)))
|
||||
@@ -109,7 +112,7 @@ def valid_step(batch):
|
||||
|
||||
def plot_everything(batch: dict, type: str):
|
||||
model.eval()
|
||||
pred = model([batch['uv_normals'], batch['uv_mask']]).detach() * 1.1
|
||||
pred = model([batch['uv_normals'], batch['uv_mask']]).detach() * 1.4
|
||||
|
||||
# mesh_pred = uv2mesh(pred)
|
||||
# mesh_orig = batch['mesh_normals']
|
||||
@@ -117,7 +120,7 @@ def plot_everything(batch: dict, type: str):
|
||||
orig_img = torch.permute(batch['uv_normals'], [0,2,3,1])
|
||||
pred_img = torch.permute(pred, [0,2,3,1])
|
||||
img = torch.concat((orig_img, pred_img), axis=2)
|
||||
imgs = [wandb.Image((img[i].cpu().numpy() + 1.1) / 2.2) for i in range(4)]
|
||||
imgs = [wandb.Image((img[i].cpu().numpy() + 1.4) / 2.8) for i in range(4)]
|
||||
|
||||
# F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
|
||||
# garment_unposed, faces = mesh_pred.cpu().numpy(), F
|
||||
@@ -196,7 +199,7 @@ if __name__ == '__main__':
|
||||
optimizer = optim.Adam(model.parameters(), lr=config['lr'])
|
||||
|
||||
# _MSSSIM = MSSSIM(data_range=2).to(config['device'])
|
||||
_MSSSIM = MS_SSIM(data_range=2.2, size_average=True, channel=3)
|
||||
_MSSSIM = MS_SSIM(data_range=2.8, size_average=True, channel=3)
|
||||
|
||||
if config['pretrained_model'] != '':
|
||||
load_model(model.encoder, config['pretrained_model'] + '_enc')
|
||||
|
||||
+15
-10
@@ -10,11 +10,13 @@ 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.ae_static_cape_m import config
|
||||
# from configs.ae_static_orig_2 import config
|
||||
# from configs.ae_static_cape_m import config
|
||||
from configs.iccvw.ae_static_1 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']:
|
||||
@@ -49,6 +51,7 @@ from pytorch_msssim import MS_SSIM
|
||||
|
||||
|
||||
def init_wandb():
|
||||
wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
|
||||
run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
|
||||
artifact = wandb.Artifact('code', type='code')
|
||||
artifact.add_file('scripts/ae_static_train.py')
|
||||
@@ -91,7 +94,7 @@ def train_step(batch):
|
||||
mesh_orig = batch['mesh_static']
|
||||
l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_static'] - pred)), axis=1))
|
||||
e_loss = 1e-2 * torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
|
||||
metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
|
||||
metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
|
||||
|
||||
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, }
|
||||
|
||||
@@ -205,18 +208,20 @@ if __name__ == '__main__':
|
||||
_MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
|
||||
|
||||
if config['pretrained_model'] != '':
|
||||
load_model(model.encoder, config['pretrained_model'] + '_enc')
|
||||
load_model(model.decoder, config['pretrained_model'] + '_dec')
|
||||
it = load_model(model.encoder, config['pretrained_model'] + '_enc', optimizer=optimizer)
|
||||
load_model(model.decoder, config['pretrained_model'] + '_dec', optimizer=optimizer)
|
||||
pre_name = config['pretrained_model']
|
||||
print(f'continued training from: {pre_name} at iteration: {it}')
|
||||
|
||||
ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_enc', top_n=5)
|
||||
ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_dec', top_n=5)
|
||||
ckpt_saver_enc = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_enc', top_n=5)
|
||||
ckpt_saver_dec = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_dec', top_n=5)
|
||||
|
||||
run_name = config['name']
|
||||
print(f'started training: {run_name}')
|
||||
|
||||
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()
|
||||
|
||||
|
||||
+25
-15
@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
|
||||
args = vars(parser.parse_args())
|
||||
|
||||
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = '6'
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = '2'
|
||||
|
||||
|
||||
import time
|
||||
@@ -49,26 +49,36 @@ gender_modifier = '' if not is_cape else ('_m' if is_male else '_f')
|
||||
smpl_res = 3638 if is_cape else 14475
|
||||
use_cache = False
|
||||
name=f'iccv_{low_res}' + skirt_modifier + cape_modifier + gender_modifier
|
||||
train_dir = '/data-local/127/cloth3d_processed'
|
||||
valid_dir = '/data-local/127/cloth3d_processed'
|
||||
|
||||
if not is_skirt:
|
||||
train_dir = '/data-local/127/train'
|
||||
valid_dir = '/data-local/127/valid'
|
||||
if is_skirt:
|
||||
train_frames = 'frames_train_skirt.txt'
|
||||
valid_frames = 'frames_valid_skirt.txt'
|
||||
else:
|
||||
train_dir = '/data-local/128/train_skirt'
|
||||
valid_dir = '/data-local/127/valid_skirt'
|
||||
train_frames = 'frames_train.txt'
|
||||
valid_frames = 'frames_valid.txt'
|
||||
|
||||
if is_cape:
|
||||
train_frames = None
|
||||
valid_frames = None
|
||||
|
||||
if is_cape:
|
||||
if is_male:
|
||||
train_dir = '/data-local/127/cape_processed/train/male'
|
||||
valid_dir = '/data-local/127/cape_processed/test/male'
|
||||
train_dir = '/data-local/127_2/cape_processed/train/male'
|
||||
valid_dir = '/data-local/127_2/cape_processed/test/male'
|
||||
else:
|
||||
train_dir = '/data-local/127/cape_processed/train/female'
|
||||
valid_dir = '/data-local/127/cape_processed/test/female'
|
||||
train_dir = '/data-local/127_2/cape_processed/train/female'
|
||||
valid_dir = '/data-local/127_2/cape_processed/test/female'
|
||||
|
||||
print(f'nohup python scripts/upscaler_train.py --name {name} > nohup/up/up_{name}.out 2>&1 &')
|
||||
if args['debug']:
|
||||
exit()
|
||||
|
||||
|
||||
def init_wandb():
|
||||
run = wandb.init(project='upscale', name=name, notes='', dir='/data-local/127/temp', tags=['upscaler', 'iccv'])
|
||||
wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
|
||||
run = wandb.init(project='upscale', name=name, notes='', dir='/data-local/127/temp', tags=['upscaler', 'iccvw'])
|
||||
artifact = wandb.Artifact('code', type='code')
|
||||
artifact.add_file('scripts/upscaler_train.py')
|
||||
artifact.add_file('losses/losses.py')
|
||||
@@ -255,8 +265,8 @@ if __name__ == '__main__':
|
||||
}
|
||||
|
||||
# dataset = Dataset('/data-local/127/train', config['signature'], resolutions=[512], device=device, frames_file='frames_jumpsuit.txt')
|
||||
dataset_train = Dataset(train_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=None, resolutions=[512, low_res], is_cape=is_cape)
|
||||
dataset_valid = Dataset(valid_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=None, resolutions=[512, low_res], is_cape=is_cape)
|
||||
dataset_train = Dataset(train_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=train_frames, resolutions=[512, low_res], is_cape=is_cape)
|
||||
dataset_valid = Dataset(valid_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=valid_frames, resolutions=[512, low_res], is_cape=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()
|
||||
@@ -279,13 +289,13 @@ if __name__ == '__main__':
|
||||
_MSSSIM_small = MS_SSIM(data_range=2, size_average=True, channel=3, weights=[1.0])
|
||||
|
||||
step = 0
|
||||
ckpt_saver = CheckpointSaver('checkpoints/iccv/torch/' + name, top_n=5)
|
||||
ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + name, top_n=5)
|
||||
|
||||
print(f'started training: {name}')
|
||||
|
||||
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()
|
||||
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
import os
|
||||
import sys
|
||||
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
||||
|
||||
import numpy as np
|
||||
|
||||
from custom_utils.custom_utils import save_stl, write_obj
|
||||
from data.cloth3d_io import readPC2Frame
|
||||
from data.utils import Frame
|
||||
from custom_utils.custom_utils import _maskFaces
|
||||
import plotly.graph_objects as go
|
||||
from plotly.subplots import make_subplots
|
||||
|
||||
|
||||
def remove_long_edges(vertices, faces):
|
||||
# Compute edge lengths
|
||||
edges = np.vstack((faces[:,[0,1]], faces[:,[1,2]], faces[:,[2,0]]))
|
||||
edge_lengths = np.linalg.norm(vertices[edges[:,0]] - vertices[edges[:,1]], axis=1)
|
||||
|
||||
# Find edges longer than 0.1
|
||||
long_edges = edges[edge_lengths > 0.1]
|
||||
|
||||
# Find vertices that are part of long edges
|
||||
vertices_to_remove = np.unique(long_edges)
|
||||
|
||||
# Remove vertices and update faces
|
||||
mask = np.ones(vertices.shape[0], dtype=bool)
|
||||
mask[vertices_to_remove] = False
|
||||
new_vertices = vertices[mask]
|
||||
|
||||
|
||||
return new_vertices, mask
|
||||
|
||||
|
||||
def plot_unposing_example():
|
||||
frame = Frame('00438', 146, 'Jumpsuit')
|
||||
mesh_static = np.load(f'{folder_test}/mesh_static/{frame.id}_{frame.garment}.npy')
|
||||
|
||||
mesh_mask = np.load(f'{folder_test}/mesh_mask/{frame.id}_{frame.garment}.npy')
|
||||
mesh_posed = readPC2Frame(os.path.join(folder_test, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
|
||||
mesh_unposed = readPC2Frame(os.path.join(folder_test, 'mesh_unposed_no_shape', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
|
||||
mesh_unposed_shaped = readPC2Frame(os.path.join('temp', f'{frame.id}_{frame.garment}_unposed_shaped.pc16'), frame.frame_nr, True)[mesh_mask]
|
||||
|
||||
|
||||
F = np.load('smpl/smpl_faces.npy')
|
||||
_, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask)
|
||||
|
||||
mesh_static, v_mask = remove_long_edges(mesh_static, faces.astype(np.int32))
|
||||
|
||||
mesh_posed = mesh_posed[v_mask]
|
||||
mesh_unposed = mesh_unposed[v_mask]
|
||||
mesh_unposed_shaped = mesh_unposed_shaped[v_mask]
|
||||
new_mask = np.zeros_like(mesh_mask)
|
||||
new_mask[mesh_mask] = v_mask
|
||||
_, faces = _maskFaces(np.zeros((14475, 3)), F, new_mask)
|
||||
|
||||
fig = make_subplots(rows=1, cols=4, specs=[[{'type': 'mesh3d'}, {'type': 'mesh3d'}, {'type': 'mesh3d'}, {'type': 'mesh3d'}]])
|
||||
|
||||
fig.add_trace(go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,2], z=mesh_static[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1), row=1, col=1)
|
||||
fig.add_trace(go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightblue', opacity=1), row=1, col=2)
|
||||
fig.add_trace(go.Mesh3d(x=mesh_unposed_shaped[:,0], y=mesh_unposed_shaped[:,2], z=mesh_unposed_shaped[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightgreen', opacity=1), row=1, col=3)
|
||||
fig.add_trace(go.Mesh3d(x=mesh_unposed[:,0], y=mesh_unposed[:,2], z=mesh_unposed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightyellow', opacity=1), row=1, col=4)
|
||||
|
||||
fig.update_layout(scene=dict(
|
||||
xaxis=dict(range=[-1, 1]),
|
||||
yaxis=dict(range=[-1, 1]),
|
||||
zaxis=dict(range=[-1, 1]),
|
||||
), scene_aspectmode='cube')
|
||||
|
||||
# fig.show()
|
||||
faces = faces.astype(np.uint32)
|
||||
|
||||
save_stl(f'temp/unposing_plots/{frame.id}_{frame.frame_nr}_{frame.garment}_mesh_static.stl', mesh_static, faces)
|
||||
save_stl(f'temp/unposing_plots/{frame.id}_{frame.frame_nr}_{frame.garment}_mesh_posed.stl', mesh_posed, faces)
|
||||
save_stl(f'temp/unposing_plots/{frame.id}_{frame.frame_nr}_{frame.garment}_mesh_unposed.stl', mesh_unposed, faces)
|
||||
save_stl(f'temp/unposing_plots/{frame.id}_{frame.frame_nr}_{frame.garment}_mesh_unposed_shaped.stl', mesh_unposed_shaped, faces)
|
||||
|
||||
pass
|
||||
|
||||
|
||||
def plot_normals_example():
|
||||
frame = Frame('00438', 146, 'Jumpsuit')
|
||||
mesh_static = np.load(f'{folder_test}/mesh_static/{frame.id}_{frame.garment}.npy')
|
||||
|
||||
mesh_mask = np.load(f'{folder_test}/mesh_mask/{frame.id}_{frame.garment}.npy')
|
||||
mesh_posed = readPC2Frame(os.path.join(folder_test, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
|
||||
# mesh_unposed = readPC2Frame(os.path.join(folder_test, 'mesh_unposed_no_shape', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
|
||||
# mesh_unposed_shaped = readPC2Frame(os.path.join('temp', f'{frame.id}_{frame.garment}_unposed_shaped.pc16'), frame.frame_nr, True)[mesh_mask]
|
||||
mesh_normals = readPC2Frame(os.path.join(folder_test, 'mesh_normals', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
|
||||
uv_normals = np.load(os.path.join(folder_test, 'uv_normals', f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))
|
||||
|
||||
F = np.load('smpl/smpl_faces.npy')
|
||||
_, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask)
|
||||
|
||||
mesh_static, v_mask = remove_long_edges(mesh_static, faces.astype(np.int32))
|
||||
|
||||
mesh_posed = mesh_posed[v_mask]
|
||||
# mesh_unposed = mesh_unposed[v_mask]
|
||||
# mesh_unposed_shaped = mesh_unposed_shaped[v_mask]
|
||||
new_mask = np.zeros_like(mesh_mask)
|
||||
new_mask[mesh_mask] = v_mask
|
||||
_, faces = _maskFaces(np.zeros((14475, 3)), F, new_mask)
|
||||
|
||||
fig = make_subplots(rows=1, cols=4, specs=[[{'type': 'mesh3d'}, {'type': 'mesh3d'}, {'type': 'mesh3d'}, {'type': 'mesh3d'}]])
|
||||
|
||||
fig.add_trace(go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,2], z=mesh_static[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1), row=1, col=1)
|
||||
fig.add_trace(go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightblue', opacity=1), row=1, col=2)
|
||||
fig.add_trace(go.Mesh3d(x=mesh_unposed_shaped[:,0], y=mesh_unposed_shaped[:,2], z=mesh_unposed_shaped[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightgreen', opacity=1), row=1, col=3)
|
||||
fig.add_trace(go.Mesh3d(x=mesh_unposed[:,0], y=mesh_unposed[:,2], z=mesh_unposed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightyellow', opacity=1), row=1, col=4)
|
||||
|
||||
fig.update_layout(scene=dict(
|
||||
xaxis=dict(range=[-1, 1]),
|
||||
yaxis=dict(range=[-1, 1]),
|
||||
zaxis=dict(range=[-1, 1]),
|
||||
), scene_aspectmode='cube')
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
folder_train = '/data-local/127/cloth3d_processed'
|
||||
folder_test = '/data-local/127/cloth3d_processed_test'
|
||||
|
||||
# plot_unposing_example()
|
||||
plot_normals_example()
|
||||
|
||||
pass
|
||||
Reference in New Issue
Block a user