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

Compare commits

base: luni/cvpr_2023:455b5d02e76ff2f2dfbeacd34a28f00bc9aaf22a
Branches Tags
luni/cvpr_2023:main luni/cvpr_2023:iccv_review luni/cvpr_2023:vto luni/cvpr_2023:extrapolation luni/cvpr_2023:skip_conn
..
compare: luni/cvpr_2023:vto
Branches Tags
luni/cvpr_2023:main luni/cvpr_2023:iccv_review luni/cvpr_2023:vto luni/cvpr_2023:extrapolation luni/cvpr_2023:skip_conn
3 Commits
455b5d02e7 .. vto
Author SHA1 Message Date
luni 3abe7554bc vto conversion, some progress, not working properl 2023-05-29 10:28:31 +02:00
luni e34e2143e7 initial work 2023-05-25 10:57:55 +02:00
luni 53217c6e5d direct training 2023-05-23 22:57:08 +02:00
81 changed files with 1265 additions and 7038 deletions
Expand all files Collapse all files
configs/archive/abl_all.py → configs/abl_all.py
View File
configs/archive/abl_all_2.py → configs/abl_all_2.py
View File
configs/archive/ae_body_1.py → configs/ae_body_1.py
View File
configs/archive/ae_body_cape_f.py → configs/ae_body_cape_f.py
View File
configs/archive/ae_body_cape_m.py → configs/ae_body_cape_m.py
+4 -4
View File
@@ -2,7 +2,7 @@ from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '0'
config['gpu'] = '4'
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/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_dir'] = '/data-local/127/cape_processed/train/male'
config['valid_dir'] = '/data-local/127/cape_processed/test/male'
config['test_dir'] = '/data-local/127/cape_processed/test/male'
config['train_frames'] = None
config['valid_frames'] = None
configs/archive/ae_curv_1.py → configs/ae_curv_1.py
View File
configs/archive/ae_curv_cape_f.py → configs/ae_curv_cape_f.py
View File
configs/archive/ae_curv_cape_m.py → configs/ae_curv_cape_m.py
View File
configs/archive/ae_static.py → configs/ae_static.py
View File
configs/archive/ae_static_cape_f.py → configs/ae_static_cape_f.py
View File
configs/archive/ae_static_cape_m.py → configs/ae_static_cape_m.py
View File
configs/archive/ae_static_orig.py → configs/ae_static_orig.py
View File
configs/archive/ae_static_orig_2.py → configs/ae_static_orig_2.py
View File
configs/archive/abl_norm.py
-192
View File
@@ -1,192 +0,0 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '0'
config['name'] = 'abl_torch_norm'
config['notes'] = ''
config['tags'] = ['iccv', 'abl_norm', 'torch']
config['script'] = 'scripts/cvae.py'
config['pretrained_model'] = ''
config['use_skinning'] = True # and deshaping
config['remove_global_orientation'] = True
config['use_normals'] = True
config['use_extra_losses'] = True # edge, collision, 3d mesh
config['use_offsets'] = True
config['use_norm_loss'] = True
ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_normals'], config['use_extra_losses'], config['use_offsets']].count(False)
assert ablation_count == 1 or ablation_count == 0
config['avg_garment_path'] = 'custom_data/avg/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv train_no_shape_one_hot_avg_uv
if config['use_skinning']:
assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
config['batch_size'] = 4
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = False
config['is_cape'] = False
config['device'] = 'cuda'
config['train_dir'] = '/data-local/127/train'
config['valid_dir'] = '/data-local/127/valid'
config['test_dir'] = '/data-net/dataset/test_ready'
config['train_frames'] = None
config['valid_frames'] = None
config['static_model'] = 'checkpoints/iccv/static_torch_orig_2'
config['body_model'] = 'checkpoints/iccv/body_torch_1'
config['normals_model'] = 'checkpoints/iccv/torch/norm_torch_1'
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['latent_dim'] = 128
###############
config['lr'] = 0.0001# 0.0001
###############
config['starting_balancing_coeff'] = 0.0001
config['signature']: set = { 'id',
'frame_nr',
'pose_shape',
'uv_body_posed',
'mesh_body_posed',
'mesh_posed',
'uv_normals',
'uv_static',
'mesh_mask',
'uv_mask',
}
if config['use_skinning']:
config['signature'].add('uv_unposed')
config['signature'].add('mesh_unposed')
else:
config['signature'].add('uv_posed')
config['upscale_enhancer'] = ''
config['freeze_encoders'] = True
config['use_one_hot'] = False
config['bottleneck_scale'] = 2
config['cvae'] = dict()
config['cvae']['device'] = config['device']
config['cvae']['is_label'] = False
config['cvae']['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
config['cvae']['kernel_size'] = 7
config['cvae']['bottleneck_size'] = 64 * config['bottleneck_scale']
config['cvae']['latent_dim'] = config['cvae']['bottleneck_size']
block_setting = [
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 3),
CNBlockConfig(128, 256, 3),
CNBlockConfig(256, config['cvae']['bottleneck_size'], 9),
CNBlockConfig(config['cvae']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cvae']['bottleneck_size'] * 2, 256, 3),
CNBlockConfig(256, 128, 9),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 3),
CNBlockConfig(32, 3, 3),
]
config['cvae']['encoder_block_settings'] = block_setting
config['cvae']['decoder_block_settings'] = block_setting_reverse
####################### cond_ae_norm ############################
config['cond_ae_norm'] = dict()
config['cond_ae_norm']['device'] = config['device']
config['cond_ae_norm']['is_label'] = True
config['cond_ae_norm']['latent_dim'] = config['latent_dim']
config['cond_ae_norm']['kernel_size'] = 7
config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 9),
CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
]
config['cond_ae_norm']['encoder_block_settings'] = block_setting
config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_norm']['activation'] = 'tanh'
####################### cond_ae_body ############################
config['cond_ae_body'] = dict()
config['cond_ae_body']['device'] = config['device']
config['cond_ae_body']['is_label'] = False
config['cond_ae_body']['latent_dim'] = config['latent_dim']
config['cond_ae_body']['kernel_size'] = 7
config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 9),
CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
]
config['cond_ae_body']['encoder_block_settings'] = block_setting
config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_body']['activation'] = 'tanh'
####################### cond_ae_static ############################
config['cond_ae_static'] = dict()
config['cond_ae_static']['device'] = config['device']
config['cond_ae_static']['is_label'] = False
config['cond_ae_static']['latent_dim'] = config['latent_dim']
config['cond_ae_static']['kernel_size'] = 7
config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(8, 16, 3),
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 8, 3),
CNBlockConfig(8, 3, 3),
]
config['cond_ae_static']['encoder_block_settings'] = block_setting
config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_static']['activation'] = 'tanh'
configs/archive/ae_norm_1.py
-73
View File
@@ -1,73 +0,0 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '0'
config['name'] = 'norm_torch_2'
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/train'
config['valid_dir'] = '/data-local/127/valid'
config['test_dir'] = '/data-net/dataset/test_ready'
config['train_frames'] = None
config['valid_frames'] = None
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = False
config['is_cape'] = False
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',
}
configs/archive/cvae_norm_1.py
-172
View File
@@ -1,172 +0,0 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '0'
config['name'] = 'cvae_norm_1'
config['notes'] = ''
config['tags'] = ['iccv', 'cvae_norm', 'torch']
config['script'] = 'scripts/cvae_norm.py'
config['pretrained_model'] = ''
config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
config['batch_size'] = 4
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = False
config['is_cape'] = False
config['device'] = 'cuda'
config['train_dir'] = '/data-local/127/train'
config['valid_dir'] = '/data-local/127/valid'
config['test_dir'] = '/data-net/dataset/test_ready'
config['train_frames'] = None
config['valid_frames'] = None
config['static_model'] = 'checkpoints/iccv/static_torch_orig_2'
config['body_model'] = 'checkpoints/iccv/body_torch_1'
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['latent_dim'] = 128
###############
config['lr'] = 0.0001# 0.0001
###############
config['starting_balancing_coeff'] = 0.0001
config['signature']: set = { 'id',
'frame_nr',
'pose_shape',
'uv_body_posed',
'uv_normals',
'uv_static',
'mesh_mask',
'uv_mask',
'mesh_normals',
}
config['upscale_enhancer'] = ''
config['freeze_encoders'] = True
config['use_one_hot'] = False
config['bottleneck_scale'] = 2
config['cvae'] = dict()
config['cvae']['device'] = config['device']
config['cvae']['is_label'] = False
config['cvae']['input_features'] = 3+3+3
config['cvae']['kernel_size'] = 7
config['cvae']['bottleneck_size'] = 64 * config['bottleneck_scale']
config['cvae']['latent_dim'] = config['cvae']['bottleneck_size']
block_setting = [
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 3),
CNBlockConfig(128, 256, 3),
CNBlockConfig(256, config['cvae']['bottleneck_size'], 9),
CNBlockConfig(config['cvae']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(176, 256, 3),
CNBlockConfig(256, 128, 9),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 3),
CNBlockConfig(32, 3, 3),
]
config['cvae']['encoder_block_settings'] = block_setting
config['cvae']['decoder_block_settings'] = block_setting_reverse
####################### cond_ae_curv ############################
config['cond_ae_curv'] = dict()
config['cond_ae_curv']['device'] = config['device']
config['cond_ae_curv']['is_label'] = True
config['cond_ae_curv']['latent_dim'] = config['latent_dim']
config['cond_ae_curv']['kernel_size'] = 7
config['cond_ae_curv']['bottleneck_size'] = 40 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 9),
CNBlockConfig(128, config['cond_ae_curv']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], 128, 3),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
]
config['cond_ae_curv']['encoder_block_settings'] = block_setting
config['cond_ae_curv']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_curv']['activation'] = 'tanh'
####################### cond_ae_body ############################
config['cond_ae_body'] = dict()
config['cond_ae_body']['device'] = config['device']
config['cond_ae_body']['is_label'] = False
config['cond_ae_body']['latent_dim'] = config['latent_dim']
config['cond_ae_body']['kernel_size'] = 7
config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 9),
CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
]
config['cond_ae_body']['encoder_block_settings'] = block_setting
config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_body']['activation'] = 'tanh'
####################### cond_ae_static ############################
config['cond_ae_static'] = dict()
config['cond_ae_static']['device'] = config['device']
config['cond_ae_static']['is_label'] = False
config['cond_ae_static']['latent_dim'] = config['latent_dim']
config['cond_ae_static']['kernel_size'] = 7
config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(8, 16, 3),
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 8, 3),
CNBlockConfig(8, 3, 3),
]
config['cond_ae_static']['encoder_block_settings'] = block_setting
config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_static']['activation'] = 'tanh'
configs/direct_1.py
+77
View File
@@ -0,0 +1,77 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '5'
config['name'] = 'direct_1'
config['notes'] = ' '
config['tags'] = ['torch', 'direct']
config['script'] = 'scripts/direct_train.py'
config['pretrained_model'] = ''
config['use_curvature'] = False
config['batch_size'] = 4
config['train_dir'] = '/data-local/127/train'
config['valid_dir'] = '/data-local/127/valid'
config['test_dir'] = '/data-net/dataset/test_ready'
config['train_frames'] = None
config['valid_frames'] = None
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = False
config['is_cape'] = False
config['remove_global_orientation'] = True
config['device'] = 'cuda'
config['input_features'] = 6
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['latent_dim'] = 128
config['bottleneck_scale'] = 2
config['is_label'] = False
config['kernel_size'] = 7
config['bottleneck_size'] = 64 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 128, 3),
CNBlockConfig(128, 256, 9),
CNBlockConfig(256, config['bottleneck_size'], 9),
CNBlockConfig(config['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['bottleneck_size'], 256, 3),
CNBlockConfig(256, 128, 9),
CNBlockConfig(128, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
# CNBlockConfig(3, 3, 3),
]
config['encoder_block_settings'] = block_setting
config['decoder_block_settings'] = block_setting_reverse
config['activation'] = 'tanh'
config['lr'] = 0.0001
config['reg_loss_coeff'] = 0.1
config['signature'] = { 'id',
'frame_nr',
'pose_shape',
'uv_body_posed',
'mesh_mask',
'mesh_body_posed',
'mesh_posed',
'uv_unposed',
'mesh_unposed',
'uv_static',
'uv_mask',
}
configs/direct_2.py
+78
View File
@@ -0,0 +1,78 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '0'
config['name'] = 'direct_2'
config['notes'] = 'with curvature input'
config['tags'] = ['torch', 'direct']
config['script'] = 'scripts/direct_train.py'
config['pretrained_model'] = ''
config['use_curvature'] = True
config['batch_size'] = 4
config['train_dir'] = '/data-local/127/train'
config['valid_dir'] = '/data-local/127/valid'
config['test_dir'] = '/data-net/dataset/test_ready'
config['train_frames'] = None
config['valid_frames'] = None
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = False
config['is_cape'] = False
config['remove_global_orientation'] = True
config['device'] = 'cuda'
config['input_features'] = 10
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['latent_dim'] = 128
config['bottleneck_scale'] = 2
config['is_label'] = False
config['kernel_size'] = 7
config['bottleneck_size'] = 64 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 128, 3),
CNBlockConfig(128, 256, 9),
CNBlockConfig(256, config['bottleneck_size'], 9),
CNBlockConfig(config['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['bottleneck_size'], 256, 3),
CNBlockConfig(256, 128, 9),
CNBlockConfig(128, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
# CNBlockConfig(3, 3, 3),
]
config['encoder_block_settings'] = block_setting
config['decoder_block_settings'] = block_setting_reverse
config['activation'] = 'tanh'
config['lr'] = 0.0001
config['reg_loss_coeff'] = 0.1
config['signature'] = { 'id',
'frame_nr',
'pose_shape',
'uv_curvature',
'mesh_mask',
'mesh_body_posed',
'mesh_posed',
'uv_unposed',
'mesh_unposed',
'uv_static',
'uv_mask',
'uv_body_posed',
}
configs/iccvw/abl_all.py
-195
View File
@@ -1,195 +0,0 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '2'
config['name'] = 'abl_iccvw_all_1'
config['notes'] = ''
config['tags'] = ['iccv', 'abl_all', 'torch']
config['script'] = 'scripts/cvae.py'
config['pretrained_model'] = ''
config['use_skinning'] = True # and deshaping
config['use_normals'] = True
config['use_extra_losses'] = True # edge, collision, 3d mesh
config['use_norm_loss'] = True
ablation_count = [config['use_skinning'], config['use_normals'], config['use_extra_losses']].count(False)
assert ablation_count == 1 or ablation_count == 0
config['avg_garment_path'] = 'custom_data/avg/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv train_no_shape_one_hot_avg_uv
if config['use_skinning']:
assert (config['avg_garment_path'] != '')
config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
config['batch_size'] = 4
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = False
config['is_cape'] = False
config['device'] = 'cuda'
config['train_dir'] = '/data-local/127/cloth3d_processed'
config['valid_dir'] = '/data-local/127/cloth3d_processed'
config['test_dir'] = '/data-net/dataset/test_ready'
if config['is_skirt']:
config['train_frames'] = 'frames_train_skirt.txt'
config['valid_frames'] = 'frames_valid_skirt.txt'
else:
config['train_frames'] = 'frames_train.txt'
config['valid_frames'] = 'frames_valid.txt'
config['static_model'] = 'checkpoints/iccvw/static_iccvw_4'
config['body_model'] = 'checkpoints/iccvw/body_iccvw_1'
config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1'
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['latent_dim'] = 128
###############
config['lr'] = 0.0001# 0.0001
###############
config['starting_balancing_coeff'] = 0.0001
config['balancing_activation_limit'] = 0.01
config['signature']: set = { 'id',
'frame_nr',
'pose_shape',
'uv_body_posed',
'mesh_body_posed',
'mesh_posed',
'uv_normals',
'uv_static',
'mesh_mask',
'uv_mask',
}
if config['use_skinning']:
config['signature'].add('uv_unposed')
config['signature'].add('mesh_unposed')
else:
config['signature'].add('uv_posed')
config['upscale_enhancer'] = ''
config['freeze_encoders'] = True
config['use_one_hot'] = False
config['bottleneck_scale'] = 2
config['cvae'] = dict()
config['cvae']['device'] = config['device']
config['cvae']['is_label'] = False
config['cvae']['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
config['cvae']['kernel_size'] = 7
config['cvae']['bottleneck_size'] = 64 * config['bottleneck_scale']
config['cvae']['latent_dim'] = config['cvae']['bottleneck_size']
block_setting = [
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 3),
CNBlockConfig(128, 256, 3),
CNBlockConfig(256, config['cvae']['bottleneck_size'], 9),
CNBlockConfig(config['cvae']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cvae']['bottleneck_size'] * 2, 256, 3),
CNBlockConfig(256, 128, 9),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 3),
CNBlockConfig(32, 3, 3),
]
config['cvae']['encoder_block_settings'] = block_setting
config['cvae']['decoder_block_settings'] = block_setting_reverse
####################### cond_ae_norm ############################
config['cond_ae_norm'] = dict()
config['cond_ae_norm']['device'] = config['device']
config['cond_ae_norm']['latent_dim'] = config['latent_dim']
config['cond_ae_norm']['kernel_size'] = 7
config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 9),
CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
]
config['cond_ae_norm']['encoder_block_settings'] = block_setting
config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_norm']['activation'] = 'tanh'
####################### cond_ae_body ############################
config['cond_ae_body'] = dict()
config['cond_ae_body']['device'] = config['device']
config['cond_ae_body']['is_label'] = False
config['cond_ae_body']['latent_dim'] = config['latent_dim']
config['cond_ae_body']['kernel_size'] = 7
config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 9),
CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
]
config['cond_ae_body']['encoder_block_settings'] = block_setting
config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_body']['activation'] = 'tanh'
####################### cond_ae_static ############################
config['cond_ae_static'] = dict()
config['cond_ae_static']['device'] = config['device']
config['cond_ae_static']['is_label'] = False
config['cond_ae_static']['latent_dim'] = config['latent_dim']
config['cond_ae_static']['kernel_size'] = 7
config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(8, 16, 3),
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 8, 3),
CNBlockConfig(8, 3, 3),
]
config['cond_ae_static']['encoder_block_settings'] = block_setting
config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_static']['activation'] = 'tanh'
configs/iccvw/ae_body_1.py
-78
View File
@@ -1,78 +0,0 @@
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',
}
configs/iccvw/ae_body_1_10p.py
-78
View File
@@ -1,78 +0,0 @@
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',
}
configs/iccvw/ae_body_1_25p.py
-78
View File
@@ -1,78 +0,0 @@
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',
}
configs/iccvw/ae_body_1_50p.py
-78
View File
@@ -1,78 +0,0 @@
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',
}
configs/iccvw/ae_body_1_skirt.py
-78
View File
@@ -1,78 +0,0 @@
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',
}
configs/iccvw/ae_body_cape_f.py
-73
View File
@@ -1,73 +0,0 @@
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',
}
configs/iccvw/ae_body_cape_m.py
-73
View File
@@ -1,73 +0,0 @@
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',
}
configs/iccvw/ae_norm_1.py
-77
View File
@@ -1,77 +0,0 @@
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),
]
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',
}
configs/iccvw/ae_norm_1_skirt.py
-78
View File
@@ -1,78 +0,0 @@
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',
}
configs/iccvw/ae_norm_cape_f.py
-73
View File
@@ -1,73 +0,0 @@
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',
}
configs/iccvw/ae_norm_cape_m.py
-74
View File
@@ -1,74 +0,0 @@
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',
}
configs/iccvw/ae_static_1.py
-78
View File
@@ -1,78 +0,0 @@
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',
}
configs/iccvw/ae_static_1_skirt.py
-78
View File
@@ -1,78 +0,0 @@
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',
}
configs/iccvw/ae_static_5.py
-78
View File
@@ -1,78 +0,0 @@
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',
}
configs/iccvw/ae_static_cape_f.py
-73
View File
@@ -1,73 +0,0 @@
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',
}
configs/iccvw/ae_static_cape_m.py
-73
View File
@@ -1,73 +0,0 @@
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',
}
configs/archive/pyr_32.py → configs/pyr_32.py
View File
configs/archive/pyr_64_cape_m.py → configs/pyr_64_cape_m.py
View File
configs/archive/pyr_64_cape_m_3.py → configs/pyr_64_cape_m_3.py
View File
custom_data/avg/train_no_shape_avg_uv.npy
BIN
View File
Binary file not shown.
custom_data/avg/train_no_shape_skirt_avg_uv.npy
BIN
View File
Binary file not shown.
custom_data/avg/train_skirt_no_shape_avg_uv.npy
BIN
View File
Binary file not shown.
custom_data/zero_body.npy
BIN
View File
Binary file not shown.
custom_utils/custom_utils.py
+148 -1
View File
@@ -9,6 +9,9 @@ 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)
@@ -40,6 +43,7 @@ class UV2Mesh(nn.Module):
if uv.shape[1] == 3 or uv.shape[1] == 7:
uv = torch.permute(uv, [0,2,3,1])
if uv.shape[1] != 512:
raise NotImplementedError
uv = Resize(uv[:,:,None], size=(512, 256), interpolation=InterpolationMode.BILINEAR)
# batch_size = uv.shape[0]
@@ -71,6 +75,35 @@ 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)
@@ -87,6 +120,20 @@ 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 = []
@@ -226,6 +273,25 @@ 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
@@ -239,7 +305,6 @@ 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 = ''
@@ -248,6 +313,88 @@ 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()
data/calculate_normals.py
-112
View File
@@ -1,112 +0,0 @@
import os
import sys
import numpy as np
import trimesh as tm
import scipy.io as sio
import plotly.graph_objects as go
import plotly.express as px
from tqdm import tqdm
import torch
from scipy.spatial.transform import Rotation as R
import argparse
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from data.cloth3d_io import readOBJ, readPC2
from data.utils import Frame
from smpl.smpl_torch import SMPL
def calculate_normals(vertices, mesh_file, poses):
_, faces, _, _ = readOBJ(mesh_file)
rodr = smpl.rodrigues(torch.tensor(poses.T.reshape(-1, 1, 3), dtype=torch.float32)).numpy().reshape(vertices.shape[0], -1, 3, 3)
normals = []
for i in range(vertices.shape[0]):
r = R.from_matrix(rodr[i, 0, :, :])
# assert r.as_matrix() == rodr[frame, 0, :, :]
out = r.apply(vertices[i], inverse=True)
# create a plotly mesh plot of out and vertices [i]
# fig = go.Figure(data=[
# go.Mesh3d(
# x=out[:, 0], y=out[:, 2], z=out[:, 1],
# # i, j and k give the vertices of triangles
# i=faces[:, 0], j=faces[:, 1], k=faces[:, 2],
# name='y',
# showscale=False
# ),
# go.Mesh3d(
# x=vertices[i][:, 0], y=vertices[i][:, 1], z=vertices[i][:, 2],
# # i, j and k give the vertices of triangles
# i=faces[:, 0], j=faces[:, 1], k=faces[:, 2],
# name='x',
# showscale=False
# )
# ])
# fig.show()
mesh = tm.Trimesh(vertices=out, faces=faces, process=False, validate=False)
normals.append(mesh.vertex_normals)
return np.stack(normals)
def get_frames(folder):
uv_static_folder = os.path.join(folder, 'pose_shape')
frames = sorted(os.listdir(uv_static_folder))
frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames]
return frames
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("-f", "--folder", help = "end of range", default='/data-net/dataset/00761')
parser.add_argument("-f_obj", "--folder_object", help = "end of range", default='/data-net/dataset/cloth3d/train')
parser.add_argument("-s", "--start", help = "start of range", default=0)
parser.add_argument("-e", "--end", help = "end of range", default=10)
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
args = parser.parse_args()
obj_folder = args.folder_object # '/data-net/dataset/cloth3d/train'
folder = args.folder #'/data-local/127/train'
os.makedirs(os.path.join(folder, 'mesh_normals'), exist_ok=True)
frames = get_frames(folder)
smpl = SMPL('cpu', 'smpl')
ids = {frame.id for frame in frames}
ids = sorted(list(ids))[int(args.start):int(args.end)]
for id in tqdm(ids):
# find obj files in the folder
garment_obj_folder = os.path.join(obj_folder, id[:5])
garments = [file.split('.')[0] for file in os.listdir(garment_obj_folder) if file.endswith('.obj')]
if len(id) == 5:
garment = garments[0]
elif len(id) == 6:
garment = garments[1]
else:
raise ValueError('Invalid frame id')
mesh_file = os.path.join(garment_obj_folder, f'{garment}.obj')
vertices =readPC2(os.path.join(garment_obj_folder, f'{garment}.pc16'), True)['V']
poses = sio.loadmat(os.path.join(garment_obj_folder, f'info.mat'))['poses']
normals = calculate_normals(vertices, mesh_file, poses)
map = sio.loadmat(os.path.join(garment_obj_folder, f'{garment}_map.mat'))['map'][0]
mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy'))
mesh_normals = normals[:,map,:]
mesh_normals = mesh_normals * mask[None,:, None]
for i in range(vertices.shape[0]):
# np.savez(os.path.join(folder, 'mesh_normals', f'{id}_{i}.npz'), mesh_normals=mesh_normals)
np.savez_compressed(os.path.join(folder, 'mesh_normals', f'{id}_{i}.npz'), mesh_normals=mesh_normals[i])
# pass
# print(frame.id, frame.frame_nr)
data/cloth3d_io.py
+1 -12
View File
@@ -342,6 +342,7 @@ import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
import plotly.graph_objects as go
from custom_utils.custom_utils import _maskFaces
def plot_3d(mesh, faces):
fig = go.Figure()
@@ -369,19 +370,7 @@ def plot_3d(mesh, faces):
return fig
def quads2tris(F):
out = []
for f in F:
if len(f) == 3: out += [f]
elif len(f) == 4: out += [[f[0],f[1],f[2]],
[f[0],f[2],f[3]]]
else: print("This should not happen...")
return np.array(out, np.int32)
if __name__ == '__main__':
from custom_utils.custom_utils import _maskFaces
folder = '/home/workspace/data/dataset/cloth3d/train'
F = np.load("custom_utils/smpl_faces.npy")#.astype(np.int32)
data/convert_to_uv.py
+147 -124
View File
@@ -12,7 +12,7 @@ import pickle
import numpy as np
import scipy as sp
import trimesh as tm
# import cmapy as cm
import cmapy as cm
import tables as tb
import scipy.io as sio
import scipy.interpolate as spi
@@ -24,7 +24,7 @@ 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
from data.utils import save_to_h5file
os.environ["CUDA_VISIBLE_DEVICES"] = "6"
@@ -34,26 +34,26 @@ def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rb
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]
uv_values = uv_values[uv_map_mask, :] # add garment mask
uv_values /= uv_counter[:, 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
garment_mask = garment_mask[uv_map_mask] # 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[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 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]
uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
rbf_dest_current = rbf_dest
if garment_uv_mask is None and mask is not None:
garment_uv_mask = np.zeros((res * 2, res))
@@ -62,8 +62,8 @@ def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rb
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
garment_uv_mask = garment_uv_mask * uv_mask[:,:,0]
rbf_dest_current = rbf_dest * garment_uv_mask
return_garment_uv_mask = True
@@ -86,7 +86,7 @@ def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rb
# 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
result[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
if garment_uv_mask is not None:
result = result * garment_uv_mask[:,:,np.newaxis]
@@ -102,26 +102,26 @@ def convert_to_uv_batch_body(res, m2uv, uv_map, verts, neighbours=15, mask=None,
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]
uv_values = uv_values[:, uv_map_mask, :]
uv_values /= uv_counter[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
verts_uv[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 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
garment_mask = garment_mask[uv_map_mask] # 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']
uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
rbf_dest_current = rbf_dest
# 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), :]
uv_map_discreet_unique_unique_inv = uv_map_discreet_unique[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
@@ -133,7 +133,7 @@ def convert_to_uv_batch_body(res, m2uv, uv_map, verts, neighbours=15, mask=None,
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 = (rbf_dest + 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
@@ -172,18 +172,18 @@ def convert_to_uv_batch(res, m2uv, uv_map, verts, neighbours=15, mask=None, meth
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
garment_mask = garment_mask[uv_map_mask] # add garment mask
rbf_dest_current = rbf_dest[type]
rbf_dest_current = rbf_dest
# 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), :]
uv_map_discreet_unique_unique_inv = uv_map_discreet_unique[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 = (rbf_dest + 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
@@ -196,11 +196,11 @@ def convert_to_uv_batch(res, m2uv, uv_map, verts, neighbours=15, mask=None, meth
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]
uv_values = uv_values[uv_map_mask, :]
uv_values /= uv_counter[:, np.newaxis]
unique_uv_values = uv_values[garment_mask, :]
uv_map_discreet_unique_unique = uv_map_discreet_unique[type][garment_mask, :]
uv_map_discreet_unique_unique = uv_map_discreet_unique[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)
@@ -386,12 +386,12 @@ def test(verts, faces, values):
def get_shapedirs():
to_keep = sio.loadmat(os.path.join('smpl/keep_remove.mat'))['keep'].reshape(-1)
to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1)
shapedirs = []
with open(os.path.join('smpl/super_model_f.pkl'), 'rb') as f:
with open(os.path.join('custom_utils/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:
with open(os.path.join('custom_utils/super_model_m.pkl'), 'rb') as f:
male = pickle.load(f, encoding="latin1")
shapedirs.append(male['shapedirs'][to_keep, :])
return shapedirs
@@ -399,139 +399,162 @@ def get_shapedirs():
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, type=int)
parser.add_argument("-e", "--end", help = "end of range", default=100, type=int)
parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/data/dataset/cloth3d/train')
parser.add_argument("-s", "--start", help = "start of range", default=70500)
parser.add_argument("-e", "--end", help = "end of range", default=90000)
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
parser.add_argument("-sk", "--skirt", help = "end of range", default=0)
args = parser.parse_args()
F = {'non-skirt': np.load("smpl/smpl_faces.npy"), 'skirt': np.load("smpl/smpl_faces_skirt.npy")}
uv_map = {'non-skirt': np.load('custom_utils/stretched_uv_map_3.npy'), 'skirt': np.load('custom_utils/stretched_uv_map_skirt.npy')}
m2uv = {'non-skirt': np.load('custom_utils/m2uv.npy', allow_pickle=True).item(), 'skirt': np.load('custom_utils/m2uv_skirt.npy', allow_pickle=True).item()}
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')}
F = np.load("custom_utils/smpl_faces" + ("_skirt" if args.skirt == 1 else "") + ".npy")
uv_map = np.load('custom_utils/stretched_uv_map_' + ("skirt" if args.skirt == 1 else "3") + '.npy')
m2uv = np.load('custom_utils/m2uv' + ("_skirt" if args.skirt == 1 else "") + '.npy', allow_pickle=True).item()
m2uv = m2uv_array(m2uv)
uv_mask = np.load('custom_utils/mask_256x512' + ("_skirt" if args.skirt == 1 else "") + '.npy')
curvature_label_type = '6_25 '
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():
uv_mask[k] = (cv2.resize(v * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
uv_mask = (cv2.resize(uv_mask * 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_mask = cv2.dilate(uv_mask.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_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
uv_map_mask = unique_uv_map_mask(uv_map_discreet)
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_counter = np.zeros((uv_map.shape[0]))
uv_counter[m2uv[:, 0]] = 1
uv_counter[m2uv[:, 1]] += 1
uv_counter = uv_counter[uv_map_mask]
uv_map_discreet_unique[k] = uv_map_discreet[k][uv_map_mask[k], :]
uv_map_discreet_unique = uv_map_discreet[uv_map_mask, :]
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
rbf_dest = np.ones((uv_res * 2, uv_res))
rbf_dest[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = 0
rbf_dest[np.invert(uv_mask[:,:,0])] = 0
#########################################################################################
use_no_orient = False
no_shape = False
is_skirt = True if args.skirt == 1 else False
folder = args.folder
data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_normals', 'body_posed', 'mesh_unposed'] # 'mesh_unposed',
neighbourhood = 25
frames_file = 'frames_all.txt'
data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_curvature', 'body_posed'] # 'mesh_unposed',
ids = sorted(os.listdir(folder + '/mesh_body_posed_no_orient'))[args.start:args.end]
ids = [id.split('.')[0].split('_')[0] for id in ids]
print("skirt" if args.skirt == 1 else "not skirt")
print(f'params: -f {args.folder} -s {args.start} -e {args.end} -n {args.name} -skirt {args.skirt}')
skirt_modifier = '_skirt' if is_skirt else ''
orient_mod = '_no_orient' if use_no_orient else ''
shape_mod = '_no_shape' if no_shape else ''
frames = []
with open(os.path.join(folder, frames_file), 'r') as f:
lines = f.readlines()
lines = [line.strip() for line in lines]
frames = lines[int(args.start):int(args.end)]
ids = [frame.split()[0] for frame in frames]
ids = sorted(list(set(ids)))
shapedirs = get_shapedirs()
root = '/'.join(folder.split('/')[:-1])
type = folder.split('/')[-1]
# h5file_curv = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='r')
# h5file = tb.open_file(os.path.join(root, f'{type}_extras_uv.h5'), mode='a')
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)
extras_folder = os.path.join(root, f'{type}_extras')
os.makedirs(extras_folder, exist_ok=True)
os.makedirs(os.path.join(extras_folder, 'uv_posed' + orient_mod), exist_ok=True)
os.makedirs(os.path.join(extras_folder, 'uv_static'), exist_ok=True)
os.makedirs(os.path.join(extras_folder, 'uv_curvature'), exist_ok=True)
os.makedirs(os.path.join(extras_folder, 'uv_unposed'), exist_ok=True)
os.makedirs(os.path.join(extras_folder, 'uv_body_posed' + orient_mod), exist_ok=True)
os.makedirs(os.path.join(extras_folder, 'uv_mask'), exist_ok=True)
for id in tqdm.tqdm(ids):
garments = glob.glob(os.path.join(folder, 'mesh_mask', f'{id}_*'))
garments = [garment.split('/')[-1].split('.')[0].split('_')[-1] for garment in garments]
# if os.path.exists(os.path.join(folder, 'uv_static', f'{id}_{garment}.npz')):
info = loadInfo(os.path.join(folder, id, 'info.mat'))
garments = list(info['outfit'].keys())
assert info['poses'].shape[1] <= 300
# put all this in a function and call for all garments
# also masking is not the same as just filling the empty values with neighbouring
# values, at least the plots are different
for garment in garments:
# garment = garments[0]
# if all the uv maps have been generated, skip this id
# if all([os.path.exists(os.path.join(extras_folder, 'uv_body_posed' + orient_mod, f'{id}_{i}_{garment}.npz')) for i in range(info['poses'].shape[1])]):
# continue
for k, garment in enumerate(garments):
# if not ('Dress' in garments or 'Skirt' in garments):
# if k != 1:
# files = glob.glob(os.path.join(folder, 'uv_body_posed_no_orient', f'{id}_*_{garment}.npz'))
# for file in files:
# shutil.move(file, '_'.join(file.split('_')[:-1]) + '.npz')
# continue
mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}_{garment}.npy'))
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.zeros((14475, 3), dtype=np.float32)
mesh_static[mask, :] = np.load(os.path.join(folder, 'mesh_static', f'{id}_{garment}.npy'))
if 'mesh_normals' in data_to_convert:
mesh_normals = np.zeros((x.shape[0], 14475, 3), dtype=np.float32)
mesh_normals[:,mask, :] = readPC2(os.path.join(folder, 'mesh_normals', f'{id}_{garment}.pc16'), True)['V']
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()
type = 'skirt' if garment == 'Dress' or garment == 'Skirt' else 'non-skirt'
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)
mask = read_garment_mask(folder, id, garment)
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)
# mesh_posed_frame = read_posed_garment_frame(folder, id, garment, 0, mask=mask, shape=info['shape'], pose=info['poses'][:,0], deorient=use_no_orient, unshape=no_shape)
mesh_posed = read_posed_garment(folder, id, garment, mask=mask, shape=info['shape'], poses=info['poses'], deorient=use_no_orient, unshape=no_shape)
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)
raise NotImplementedError()
mesh_unposed = read_unposed_garment(folder, id, garment, mask=mask, shape=info['shape'], unshape=no_shape)
if 'mesh_static' in data_to_convert:
mesh_static = read_garment_template(folder, id, garment, mask)
if 'mesh_curvature' in data_to_convert:
mesh_curvature = read_curvatures(folder + '_extras', id, garment, count=info['poses'].shape[1])
if 'body_posed' in data_to_convert:
body_posed = read_posed_body(folder, id, info['poses'], use_no_orient)
start = time.perf_counter()
uv_static, uv_mask_local, garment_uv_mask_border = convert_to_uv(uv_res, m2uv, uv_map, mesh_static, neighbours=3, mask=mask, method='rbf', garment_uv_mask=None)
uv_static = convert_to_uv_batch(uv_res, m2uv, uv_map, 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.astype(np.float32)
if 'mesh_posed' in data_to_convert:
uv_posed = convert_to_uv_batch(uv_res, m2uv, uv_map, 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:
raise NotImplementedError()
uv_unposed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
if 'mesh_curvature' in data_to_convert:
uv_curvature = convert_to_uv_batch_label(uv_res, m2uv, uv_map, mesh_curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
if 'body_posed' in data_to_convert:
uv_body_posed = convert_to_uv_batch_body(uv_res, m2uv, uv_map, 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)
# mesh_static
# save_to_h5file(h5file, id, 0, garment, uv_static, 'uv_static', 'float32')
# save_to_h5file(h5file, id, 0, garment, uv_mask, 'uv_mask', 'float32')
np.savez_compressed(os.path.join(extras_folder, 'uv_static', f'{id}_{garment}.npz'), uv_static=uv_static)
np.savez_compressed(os.path.join(extras_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])
np.savez_compressed(os.path.join(extras_folder, 'uv_posed' + orient_mod, f'{id}_{i}_{garment}.npz'), uv_posed=uv_posed[i])
# save_to_h5file(h5file, id, i, garment, uv_posed[i, :, :, :], 'uv_posed' + orient_mod + shape_mod, 'float32')
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 in range(len(uv_normals)):
np.savez_compressed(os.path.join(folder, 'uv_normals', f'{id}_{i}_{garment}.npz'), uv_normals=uv_normals[i])
if 'body_posed' in data_to_convert and k != 1:
raise NotImplementedError()
for i in range(uv_unposed.shape[0]):
save_to_h5file(h5file, id, i, garment, uv_unposed[i, :, :, :], 'uv_unposed' + orient_mod + shape_mod, 'float32')
if 'mesh_curvature' in data_to_convert:
for i in range(len(uv_curvature)):
np.savez_compressed(os.path.join(extras_folder, 'uv_curvature', f'{id}_{i}_{garment}.npz'), uv_curvature=uv_curvature[i].astype(np.uint8))
# save_to_h5file(h5file, id, i, garment, uv_curvature[i, :, :, :], 'uv_curvature', 'uint8')
if 'body_posed' in data_to_convert:
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])
np.savez_compressed(os.path.join(extras_folder, 'uv_body_posed' + orient_mod, f'{id}_{i}_{garment}.npz'), uv_body_posed=uv_body_posed[i])
# save_to_h5file(h5file, id, i, garment, uv_body_posed[i, :, :, :], 'uv_body_posed', 'float32')
# pass
pass
h5file_curv.close()
data/convert_to_uv_cape.py
-530
View File
@@ -1,530 +0,0 @@
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
data/convert_vto.py
+259
View File
@@ -0,0 +1,259 @@
import numpy as np
import os
import sys
import pickle
import plotly.express as ex
import plotly.graph_objects as go
import torch
import open3d as o3d
import trimesh
# from vedo import trimesh as trimesh_plot
from scipy.spatial import KDTree
from scipy.sparse import lil_matrix
from scipy.optimize import minimize
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from smpl.smpl_torch import SMPL
# import mask faces
from custom_utils.custom_utils import _maskFaces
smpl = SMPL(path='smpl')
def energy_function(params, source_points, target_points, indices):
num_points = source_points.shape[0]
num_dims = source_points.shape[1]
# Reshape the parameters
weights = params[:num_points * num_dims].reshape((num_points, num_dims))
translation = params[num_points * num_dims:num_points * num_dims + num_dims]
control_points = params[num_points * num_dims + num_dims:].reshape((-1, num_dims))
# Compute the transformed source points
transformed_points = source_points + np.dot(weights, control_points.T)
# Compute the TPS energy
tps_energy = 0.5 * np.sum(np.square(target_points[indices] - transformed_points))
return tps_energy
def non_rigid(source_points, target_points):
tree = KDTree(target_points)
distances, indices = tree.query(source_points)
# Initialize the optimization parameters
num_points = source_points.shape[0]
num_dims = source_points.shape[1]
initial_params = np.zeros(num_points * num_dims + num_dims + num_points)
# Define the optimization bounds
bounds = [(None, None)] * (num_points * num_dims)
bounds += [(None, None)] * num_dims
bounds += [(None, None)] * (num_points * num_dims)
# Perform the optimization
result = minimize(energy_function, initial_params, method='L-BFGS-B',
args=(source_points, target_points, indices), bounds=bounds)
# Get the optimized parameters
optimized_params = result.x
# Reshape the parameters
optimized_weights = optimized_params[:num_points * num_dims].reshape((num_points, num_dims))
optimized_translation = optimized_params[num_points * num_dims:num_points * num_dims + num_dims]
optimized_control_points = optimized_params[num_points * num_dims + num_dims:].reshape((-1, num_dims))
transformed_points = source_points + np.dot(optimized_weights, optimized_control_points.T)
return transformed_points
def register_smpl():
# generate body with zero pose and shape
pose = torch.zeros((1, 72))
shape = torch.zeros((1, 11))
smpl_out = smpl(pose, shape)[0][0]
with open('/data-net/code/vto-dataset/tshirt/simulations/tshirt_shape00_55_27.pkl', "rb") as f:
data = pickle.load(f)
garment = data['vertices'][100]
faces = data['faces']
# rotate garment x degrees along y axis
# convert degrees to radians
degrees = 107
theta = degrees * np.pi / 180
rot = np.array([[np.cos(theta), 0, np.sin(theta)],
[0, 1, 0],
[-np.sin(theta), 0, np.cos(theta)]])
garment = np.dot(rot, garment.T).T
# rotate garment x degrees along x axis
# convert degrees to radians
degrees = 5
theta = degrees * np.pi / 180
rot = np.array([[1, 0, 0],
[0, np.cos(theta), -np.sin(theta)],
[0, np.sin(theta), np.cos(theta)]])
garment = np.dot(rot, garment.T).T
# rotate garment x degrees along z axis
# convert degrees to radians
degrees = -5
theta = degrees * np.pi / 180
rot = np.array([[np.cos(theta), -np.sin(theta), 0],
[np.sin(theta), np.cos(theta), 0],
[0, 0, 1]])
garment = np.dot(rot, garment.T).T
garment[:,0] -= 0.06
garment[:,1] -= 0.03
garment[:,2] -= 0.015
# visualize the body and garment
# fig = go.Figure(data=[go.Mesh3d(x=smpl_out[:,0], y=smpl_out[:,1], z=smpl_out[:,2], i=smpl.f[:,0], j=smpl.f[:,1], k=smpl.f[:,2])])
# fig.add_trace(go.Mesh3d(x=garment[:,0], y=garment[:,1], z=garment[:,2], i=faces[:,0], j=faces[:,1], k=faces[:,2]))
# fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
# fig.show()
# non_rigid(garment, smpl_out.numpy())
# Load the source and target meshes
source_mesh = o3d.geometry.TriangleMesh(o3d.utility.Vector3dVector(garment), o3d.utility.Vector3iVector(faces))
target_mesh = o3d.geometry.TriangleMesh(o3d.utility.Vector3dVector(smpl_out), o3d.utility.Vector3iVector(smpl.f))
# flip source and target mesh
source_mesh, target_mesh = target_mesh, source_mesh
source_mesh.compute_vertex_normals()
source_mesh = source_mesh.filter_smooth_laplacian(number_of_iterations=10)
# plot the smoothed mesh with plotly
# fig = go.Figure(data=[go.Mesh3d(x=np.asarray(source_mesh.vertices)[:,0], y=np.asarray(source_mesh.vertices)[:,1], z=np.asarray(source_mesh.vertices)[:,2], i=np.asarray(source_mesh.triangles)[:,0], j=np.asarray(source_mesh.triangles)[:,1], k=np.asarray(source_mesh.triangles)[:,2])])
# fig.show()
# Perform initial alignment (translation and/or scaling)
# You can modify this step based on your specific needs
# For example, you can manually set an initial transformation
init_transformation = [[1.0, 0.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.0, 1.0]]
# Create the ICP object and set the parameters
icp = o3d.pipelines.registration.TransformationEstimationPointToPoint()
criteria = o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=2000)
# Perform ICP registration
reg_result = o3d.pipelines.registration.registration_icp(o3d.geometry.PointCloud(source_mesh.vertices),
o3d.geometry.PointCloud(target_mesh.vertices),
0.075,
init_transformation, icp, criteria)
# Get the aligned source mesh after registration
aligned_mesh = source_mesh.transform(reg_result.transformation)
# Visualize the aligned mesh and target mesh with plotly
fig = go.Figure(data=[go.Mesh3d(x=np.asarray(aligned_mesh.vertices)[:,0], y=np.asarray(aligned_mesh.vertices)[:,1], z=np.asarray(aligned_mesh.vertices)[:,2], i=np.asarray(aligned_mesh.triangles)[:,0], j=np.asarray(aligned_mesh.triangles)[:,1], k=np.asarray(aligned_mesh.triangles)[:,2])])
fig.add_trace(go.Mesh3d(x=np.asarray(target_mesh.vertices)[:,0], y=np.asarray(target_mesh.vertices)[:,1], z=np.asarray(target_mesh.vertices)[:,2], i=np.asarray(target_mesh.triangles)[:,0], j=np.asarray(target_mesh.triangles)[:,1], k=np.asarray(target_mesh.triangles)[:,2]))
fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
fig.show()
# Get the vertices of the aligned source mesh and target mesh
aligned_source_vertices = np.asarray(aligned_mesh.vertices)
target_vertices = np.asarray(target_mesh.vertices)
# Perform nearest neighbor search to find corresponding vertex indices
kd_tree = o3d.geometry.KDTreeFlann(target_mesh)
correspondences = []
matched = set()
for i in range(len(aligned_source_vertices)):
k = 1
_, idx, _ = kd_tree.search_knn_vector_3d(aligned_source_vertices[i], k)
# while idx[-1] in matched:
# k += 1
# _, idx, _ = kd_tree.search_knn_vector_3d(aligned_source_vertices[i], k)
# matched.add(idx[-1])
correspondences.append((i, idx[-1]))
# correspondences is a list of tuples where each tuple contains the corresponding vertex indices
# For example, correspondences = [(0, 10), (1, 5), (2, 15), ...]
# Visualize the correspondences with plotly
# fig = go.Figure(data=[go.Scatter3d(x=aligned_source_vertices[:,0], y=aligned_source_vertices[:,1], z=aligned_source_vertices[:,2], mode='markers', marker=dict(color='red', size=1))])
# fig.add_trace(go.Scatter3d(x=target_vertices[:,0], y=target_vertices[:,1], z=target_vertices[:,2], mode='markers', marker=dict(color='blue', size=1)))
# for corr in correspondences:
# fig.add_trace(go.Scatter3d(x=[aligned_source_vertices[corr[0],0], target_vertices[corr[1],0]], y=[aligned_source_vertices[corr[0],1], target_vertices[corr[1],1]], z=[aligned_source_vertices[corr[0],2], target_vertices[corr[1],2]], mode='lines', line=dict(color='green', width=1)))
# fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
# fig.show()
registered = np.zeros((14475, 3))
mask = np.zeros((14475), dtype=bool)
for i in range(len(correspondences)):
registered[correspondences[i][1]] = aligned_source_vertices[correspondences[i][0]]
mask[correspondences[i][1]] = True
# F = np.load(f'smpl/smpl_faces.npy').astype(np.int32)
# registered_masked, masked_faces = _maskFaces(registered, np.asarray(target_mesh.triangles), mask)
registered_masked, masked_faces = registered, np.asarray(target_mesh.triangles)
# plot aligned source mesh with the target topology
fig = go.Figure(data=[go.Mesh3d(x=registered_masked[:,0], y=registered_masked[:,1], z=registered_masked[:,2], i=masked_faces[:,0], j=masked_faces[:,1], k=masked_faces[:,2])])
fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
fig.show()
return smpl
def convert_vto(folder, out_folder):
os.makedirs(out_folder, exist_ok=True)
os.makedirs(os.path.join(out_folder, 'mesh_posed_no_orient'), exist_ok=True)
os.makedirs(os.path.join(out_folder, 'mesh_posed_body_no_orient'), exist_ok=True)
os.makedirs(os.path.join(out_folder, 'mesh_mask'), exist_ok=True)
os.makedirs(os.path.join(out_folder, 'pose_shape'), exist_ok=True)
F = np.load(f'smpl/smpl_faces.npy').astype(np.int32)
files = os.listdir(folder)
# files = ['/data-net/code/vto-dataset/tshirt/simulations/tshirt_shape00_55_27.pkl']
for file in files:
# read pickle files and get vetices and faces from them
with open(os.path.join(folder, file), "rb") as f:
data = pickle.load(f)
vertices = data['vertices']
faces = data['faces']
pose = data['pose']
shape = data['shape']
trans = data['translation']
id = file.split('.')[0]
# pose[:, :3] = 0
id = id.replace('_', '-')
subject = id.split('-')[2]
sequence = '_'.join(id.split('-')[2:])
# cmu_data = np.load(f'/data-net/code/vto-garment-collisions/assets/CMU/{subject}/{sequence}_poses.npz')
# shape = np.tile(cmu_data['betas'], (shape.shape[0], 1))
# pose = cmu_data['poses']
body = smpl(torch.tensor(pose), torch.tensor(np.concatenate((shape, np.zeros((shape.shape[0],1),dtype=np.float32)),axis=1)))[0].numpy()
pose_shape = np.concatenate((pose, shape, np.zeros((shape.shape[0],1),dtype=np.float32)),axis=1)
# for i in range(vertices.shape[0]):
# np.savez_compressed(os.path.join(out_folder, 'mesh_posed_body_no_orient', f'{id}_{i}.npz'), mesh_posed_body_no_orient=body[i])
# np.save(os.path.join(out_folder, 'pose_shape', f'{id}_{i}.npy'), pose_shape[i])
# plot mesh with plotly
frame = 0
fig = go.Figure(data=[go.Mesh3d(x=vertices[frame,:,0], y=vertices[frame,:,1], z=vertices[frame,:,2], i=faces[:,0], j=faces[:,1], k=faces[:,2])])
fig.add_trace(go.Mesh3d(x=body[frame,:,0], y=body[frame,:,1], z=body[frame,:,2], i=F[:,0], j=F[:,1], k=F[:,2]))
fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
fig.show()
pass
pass
if __name__ == '__main__':
# register_smpl()
convert_vto('/data-net/code/vto-dataset/dress/simulations', '/data-net/code/vto-dataset/dress_converted')
pass
data/dataset.py
+67 -142
View File
@@ -1,7 +1,7 @@
import os
import sys
# os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'
# os.environ["CUDA_VISIBLE_DEVICES"] = '6'
import math
import time
@@ -14,14 +14,13 @@ 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):
@@ -61,6 +60,7 @@ 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,65 +74,50 @@ class StaticCache():
self.uv_masks[res] = down.cpu().numpy()
def read(self, id, garment):
def read(self, id):
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, f'uv_static{res_string}', f'{id}_{garment}.npz')); uv_static = uv_static['uv_static']; data['uv_static' + res_string] = uv_static
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, 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
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', f'{id}_{garment}.npy'))
data['mesh_static'] = np.zeros((14475, 3), dtype=np.float32)
data['mesh_static'][mesh_mask] = mesh_static.astype(np.float32)
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, frame: Frame) -> dict:
if (frame.id, frame.garment) not in self.store:
item = self.read(frame.id, frame.garment)
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[(frame.id, frame.garment)] = stored_data
self.store[id] = stored_data
data = { key: value.to(self.device) for key, value in self.store[(frame.id, frame.garment)].items()}
data = { key: value.to(self.device) for key, value in self.store[id].items()}
return data
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):
class Dataset(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
@@ -140,20 +125,13 @@ class DatasetCloth3D(torch.utils.data.Dataset):
self.resolutions = resolutions
self.transforms = tv.transforms.ToTensor()
self.frames = read_frames_file(os.path.join(folder, frames_file), shuffle=True)
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
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)
@@ -164,59 +142,69 @@ class DatasetCloth3D(torch.utils.data.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] }
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()
mesh_mask = self.static_cache[frame]['mesh_mask'].cpu().numpy().astype(bool)
if 'id' in self.signature:
item['id'] = frame.id
if 'frame_nr' in self.signature:
item['frame_nr'] = frame.frame_nr
item['garment'] = frame.garment
if 'pose_shape' in self.signature:
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)
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_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
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 = 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
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 = 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 = 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
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)
for res in self.resolutions:
res_string = ''
if res != 512:
res_string = '_' + str(res)
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_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, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed']
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_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
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, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_posed = uv_posed['uv_posed']
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
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()}
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):
@@ -226,101 +214,38 @@ class DatasetCloth3D(torch.utils.data.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] }
data = { **data, **self.static_cache[frame.id] }
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 = { 'pose_shape',
signature = { 'id',
'frame_nr',
'pose_shape',
'uv_posed',
'uv_body_posed',
'mesh_body_posed',
'mesh_posed',
'mesh_unposed',
'uv_unposed',
'uv_normals',
'mesh_normals',
'uv_curvature',
'mesh_curvature',
'uv_static',
'mesh_static',
'mesh_mask',
'uv_mask',
'uv_shaped_body',
}
dataset = Dataset('/data-local/127/cloth3d_processed', signature, resolutions=[512], device=device, frames_file='frames_train.txt')
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
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
@@ -349,7 +274,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_normals'][0].numpy().argmax(-1) / 6.0).show()
# 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)
data/dataset_cape.py
-306
View File
@@ -1,306 +0,0 @@
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()
data/dataset_vto.py
+40
View File
@@ -0,0 +1,40 @@
import os
import sys
import argparse
parser = argparse.ArgumentParser(description='')
parser.add_argument('--debug', action='store_true')
args = vars(parser.parse_args())
if args['debug']:
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
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
from torch.utils.data import Dataset
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from data.utils import Frame
class Dataset(Dataset):
def __init__(self, folder):
pass
def __len__(self):
return len(self.frames)
def __getitem__(self, idx):
frame : Frame = self.frames[idx]
if __name__ == '__main__':
from torch.utils.data import DataLoader
data/downscale.py
-223
View File
@@ -1,223 +0,0 @@
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()
data/investigate.py
+14 -63
View File
@@ -4,26 +4,11 @@ import random
import numpy as np
import plotly.graph_objects as go
from pymeshfix import MeshFix
import trimesh
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, read_curvature_frame, writeOBJ
# from custom_utils.custom_utils import _maskFaces
from custom_utils.custom_utils import _maskFaces
def _maskFaces(V, F, mask, return_unfixed=False):
V_mask = mask
F_mask = np.take(V_mask, F.reshape(-1), axis=0)
F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
V_ind = np.zeros(V_mask.shape)
V_ind[V_mask] = np.arange(V_mask.sum())
F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
V_new = V[V_mask, :]
if return_unfixed:
return V_new, F_new, F[F_mask, :]
else:
return V_new, F_new
def get_frames(folder, frames_file):
frames = []
@@ -148,7 +133,7 @@ def check(folder, frame):
v = mesh_posed[mask,:]
is_skirt = (garment == 'Dress' or garment == 'Skirt')
skirt_mod = '_skirt' if is_skirt else ''
F = np.load(f'smpl/smpl_faces{skirt_mod}.npy').astype(np.int32)
F = np.load(f'custom_utils/smpl_faces{skirt_mod}.npy').astype(np.int32)
_, faces, unfixed = _maskFaces(np.zeros((14475,3)), F, mask, True)
meshfix = MeshFix(v, faces)
@@ -195,22 +180,7 @@ def check(folder, frame):
faces_3 = mask_faces(F, mask_2, mask_2_raw, repeated, mesh_posed)
# writeOBJ('temp/investigate_3.obj', v_2, faces_3)
# plot(v_2, faces_3)
# save stl
# mesh = trimesh.Trimesh(vertices=v_2, faces=faces_3, process=False, validate=False)
# mesh.export('temp/investigate_00761.stl')
# # read stl file
# mesh = trimesh.load('temp/fixed_00761.stl', process=False, validate=False)
# mesh = trimesh.load('temp/fixed_00761.ply', process=False, validate=False)
# plotly plot of mesh with v_2 vertices and faces from trimesh mesh
# faces_4 = faces_3 #np.asarray(mesh.faces)
# go.Figure(data = [go.Mesh3d(x=v_2[:, 0], y=v_2[:, 1], z=v_2[:, 2], i=faces_4[:, 0], j=faces_4[:, 1], k=faces_4[:, 2],)
# ]).show()
plot(v_2, faces_3)
# plot repeated
@@ -219,9 +189,9 @@ def check(folder, frame):
r.extend(rr)
x = mesh_posed
y = mesh_posed[r,:]
# go.Figure(data = [go.Scatter3d(x=x[:, 0], y=x[:, 1], z=x[:, 2], marker=dict(size=3, line=dict(width=0),), mode='markers'),
# go.Scatter3d(x=y[:, 0], y=y[:, 1], z=y[:, 2], marker=dict(size=5, line=dict(width=0),), mode='markers'),
# ]).show()
go.Figure(data = [go.Scatter3d(x=x[:, 0], y=x[:, 1], z=x[:, 2], marker=dict(size=3, line=dict(width=0),), mode='markers'),
go.Scatter3d(x=y[:, 0], y=y[:, 1], z=y[:, 2], marker=dict(size=5, line=dict(width=0),), mode='markers'),
]).show()
# unq, idx, count = np.unique(v, axis=0, return_counts=True, return_index=True)
@@ -237,49 +207,30 @@ def check(folder, frame):
# faces_4 = mask_faces(faces.astype(np.int64), mask_3, mask_2_raw, repeated, v)
# plot(v[mask_3,:], faces_4)
return mask_2_raw, faces_3
pass
def remask(folder, frame, mask, out_folder):
id, frame_nr = frame.split('_')
frame_nr = int(frame_nr)
if frame_nr == 0:
mesh_static = np.load(os.path.join(folder, 'mesh_static', f'{id}.npz'))['mesh_static']
mesh_static_new = np.zeros((14475,3), dtype=np.float32)
mesh_static_new[mask,:] = mesh_static[mask,:]
np.savez_compressed(os.path.join(out_folder, 'mesh_static', f'{id}.npz'), mesh_static=mesh_static_new)
# mesh_posed = read_posed_garment_frame(folder, id, garment, frame_nr, mask=None, shape=None, pose=info['poses'][:,frame_nr], deorient=True, unshape=None)
mesh_posed = np.load(os.path.join(folder, 'mesh_posed_no_orient', f'{id}_{frame_nr}.npz'))['mesh_posed_no_orient']
mesh_posed_new = np.zeros((14475,3), dtype=np.float32)
mesh_posed_new[mask,:] = mesh_posed[mask,:]
np.savez_compressed(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{frame_nr}.npz'), mesh_posed_no_orient=mesh_posed)
if __name__ == "__main__":
folder = '/data-local/127/train'
frames = [f'00761_{i}' for i in range(300)]
folder = '/home/workspace/data/dataset/cloth3d/train'
frames = get_frames(folder, 'frames_all.txt')
root = '/'.join(folder.split('/')[:-1])
type = folder.split('/')[-1]
use_no_orient = False
no_shape = False
is_skirt = False
is_skirt = True
skirt_modifier = '_skirt' if is_skirt else ''
orient_mod = '_no_orient' if use_no_orient else ''
shape_mod = '_no_shape' if no_shape else ''
mask = np.load('/data-net/dataset/00761/mesh_mask/00761.npy')
faces = np.load('/data-net/dataset/00761/faces/00761.npy')
for frame in frames:
# frame = random.choice(frames)
# frame = '00761 0 Tshirt'
while True:
frame = random.choice(frames)
frame = '05936 0 Dress'
remask(folder, frame, mask, '/data-net/dataset/00761')
check(folder, frame)
pass
data/preprocessing.py
+5 -633
View File
@@ -1,39 +1,22 @@
import os
import sys
import argparse
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", 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"
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
import math
import shutil
import warnings
import tables as tb
import numpy as np
import scipy.io as sio
import trimesh as tm
from tqdm import tqdm
from tables import NaturalNameWarning
warnings.filterwarnings('ignore', category=NaturalNameWarning)
import plotly.graph_objects as go
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, readPC2Frame
from data.cloth3d_io import readPC2, loadInfo
from data.utils import Frame
from smpl.smpl_torch import SMPL
from custom_utils.custom_utils import _maskFaces
@@ -149,623 +132,12 @@ def create_extras_file(folder):
h5file.close()
def deorient(verts, pose):
rodr = smpl.rodrigues(torch.tensor(pose.T, device=smpl.device).reshape(-1, 1, 3)).reshape(verts.shape[0], -1, 3, 3).cpu().numpy()
rotated = np.zeros_like(verts)
for frame in range(verts.shape[0]):
r = R.from_matrix(rodr[frame, 0, :, :])
rotated[frame] = r.apply(verts[frame], inverse=True)
return rotated
def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_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)
for id in tqdm(ids):
# id = '01467'
garment_dir = os.path.join(folder, id)
if not os.path.isdir(garment_dir):
continue
garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
info = loadInfo(os.path.join(folder, id, 'info.mat'))
for garment in garments:
if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
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
# import plotly.graph_objects as go
# frame = 5
# 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 = 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)
for id in tqdm(ids):
garment_dir = os.path.join(folder, id)
if not os.path.isdir(garment_dir):
continue
info = loadInfo(os.path.join(folder, id, 'info.mat'))
shape = np.concatenate([info['shape'], np.asarray([info['gender']])]).astype(np.float32)
np.save(os.path.join(out_folder, 'pose', f'{id}.npy'), info['poses'].T.astype(np.float32))
np.save(os.path.join(out_folder, 'shape', f'{id}.npy'), shape)
def create_mesh_static(folder, out_folder):
ids = sorted(os.listdir(folder))
os.makedirs(os.path.join(out_folder, 'mesh_static'), exist_ok=True)
for id in tqdm(ids):
garment_dir = os.path.join(folder, id)
if not os.path.isdir(garment_dir):
continue
garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
for garment in garments:
if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
continue
verts_static, _, _, _ = readOBJ(os.path.join(garment_dir, f'{garment}.obj'))
reg_data = sio.loadmat(os.path.join(folder_reg, id, garment + '_weights_fixed.mat'))
verts_static = (verts_static[reg_data['F'], :] * reg_data['W'][:, :, None]).sum(axis=1).astype(np.float32)
# rotate 90 degrees around x axis
verts_static = np.matmul(verts_static, np.array([[1,0,0],[0,0,-1],[0,1,0]]))
# fig = go.Figure(data=[go.Scatter3d(x=verts_static[:,0], y=verts_static[:,2], z=verts_static[:,1], mode='markers', marker=dict(size=1))])
# fig.update_layout(scene_aspectmode='cube')
# fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
# fig.show()
np.save(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.npy'), verts_static.astype(np.float32))
def create_mesh_body_posed_no_orient(folder, out_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):
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'))
pose, shape, gender = info['poses'], info['shape'], info['gender']
pose = np.transpose(pose, [1, 0]).astype(np.float32)
pose[:, :3] = 0
beta = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32)
T = smpl(torch.tensor(pose, device=smpl.device), torch.tensor(np.tile(beta[np.newaxis, :], [pose.shape[0], 1]), device=smpl.device))[0]
T = T.cpu().numpy()
writePC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), T, True)
# verts = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V']
# # plot as scatter plot with plotly
# import plotly.graph_objects as go
# frame = 5
# 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()
pass
def calculate_normals(vertices, faces):
normals = []
for i in range(vertices.shape[0]):
mesh = tm.Trimesh(vertices=vertices[i], faces=faces, process=False, validate=False)
normals.append(mesh.vertex_normals)
return np.stack(normals)
def create_mesh_normals(folder, out_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)
if not os.path.isdir(garment_dir):
continue
garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
for garment in garments:
if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
continue
verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
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))
reg_data = sio.loadmat(os.path.join(folder_reg, id, garment + '_weights_fixed.mat'))
normals = (normals[:, reg_data['F'], :] * reg_data['W'][None, :, :, None]).sum(axis=2)
writePC2(os.path.join(out_folder, 'mesh_normals', f'{id}_{garment}.pc16'), normals, True)
# 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[frame,:,0],
# y=verts[frame,:,2],
# z=verts[frame,:,1],
# mode='markers',
# marker=dict(
# size=1,
# color=normals[frame,:,:],
# 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[frame,:,0],
# y=moved_verts[frame,:,2],
# z=moved_verts[frame,:,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 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 = sorted(os.listdir(folder))
for id in ids:
garment_dir = os.path.join(folder, id)
if not os.path.isdir(garment_dir):
continue
garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
for garment in garments:
vertices = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V']
mask = np.load(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy'))
if mask.sum() == vertices.shape[1]:
print(id, garment)
pass
def create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=0, end=1):
ids = sorted(os.listdir(folder))[start:end]
os.makedirs(os.path.join(out_folder, 'mesh_unposed_no_shape'), exist_ok=True)
F = np.load('smpl/smpl_faces.npy')
# F_skirt = np.load('smpl/smpl_faces_skirt.npy')
for id in tqdm(ids):
garment_dir = os.path.join(folder, id)
if not os.path.isdir(garment_dir):
continue
garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
# info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
pose = np.load(os.path.join(out_folder, 'pose', f'{id}.npy'))
shape = np.tile(np.load(os.path.join(out_folder, 'shape', f'{id}.npy'))[None,:], [pose.shape[0],1])
pose = torch.tensor(pose, dtype=torch.float32, device='cuda')
shape = torch.tensor(shape, dtype=torch.float32, device='cuda')
for garment in garments:
if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
continue
mesh_static = np.load(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.npy'))
mesh_posed = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V']
mask = np.load(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy'))
# mask_idx = np.where(mask)[0]
pose[:,:3] = 0
# pose = pose[:10]
# shape = shape[:10]
# mesh_posed = mesh_posed[:10]
unposed = np.zeros((14475, 3), dtype=np.float32)
unposed[mask] = mesh_static
unposed = torch.tensor(np.tile(unposed[None,:,:], [mesh_posed.shape[0],1,1]), dtype=torch.float32, device='cuda')
unposed.requires_grad = True
mesh_posed = torch.tensor(mesh_posed, dtype=torch.float32, device='cuda')
optimizer = torch.optim.Adam([unposed], lr=0.01)
# verts_1 = unposed[0].cpu().detach().numpy()[mask]
# _, faces = _maskFaces(np.zeros((14475, 3)), F, mask)
# fig = go.Figure(data=[go.Mesh3d(x=verts_1[:,0], y=verts_1[:,2], z=verts_1[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', 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()
# verts_2 = mesh_posed[0].cpu().detach().numpy()
# fig = go.Figure(data=[go.Mesh3d(x=verts_2[:,0], y=verts_2[:,2], z=verts_2[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', 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()
# zero_body = smpl(torch.zeros_like(pose), torch.zeros_like(shape))[0]
# verts_3 = smpl(pose, shape, zero_body)[0][0].cpu().detach().numpy()
# # verts_3 = zero_body[0].cpu().detach().numpy()
# faces_3 = F
# fig = go.Figure(data=[go.Mesh3d(x=verts_3[:,0], y=verts_3[:,2], z=verts_3[:,1], i=faces_3[:,0], j=faces_3[:,1], k=faces_3[:,2], color='lightpink', 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()
for i in range(5000):
optimizer.zero_grad()
pred = smpl(pose, shape, unposed)[0]
loss : torch.Tensor = ((pred[:,mask,:] - mesh_posed)**2).sum()
# Backward pass
loss.backward()
optimizer.step()
with torch.inference_mode():
euclidean_error = torch.sqrt(((pred[:,mask,:] - mesh_posed)**2).sum(-1)).mean()
# Print the loss
if i % 50 == 0:
print(f"it: {i} loss: {loss:.8f} euclidean_error: {euclidean_error:.8f}")
if euclidean_error < 0.0000001:
break
if euclidean_error < 0.0000005 and i > 2000:
break
if euclidean_error < 0.000001 and i > 4000:
break
unposed = unposed.detach().cpu().numpy()
gender = int(shape[0].cpu().numpy()[-1] > 0)
shape_offsets = torch.tensordot(shape[:,:10], smpl.shapedirs[gender], dims=[[1], [2]])
shape_offsets = shape_offsets.cpu().numpy()[:,smpl.to_keep,:]
unshaped = (unposed - shape_offsets)[:,mask,:]
writePC2(os.path.join(out_folder, 'mesh_unposed_no_shape', f'{id}_{garment}.pc16'), unshaped, True)
# verts = unposed[:,mask,:]
# verts = unshaped
# _, faces = _maskFaces(np.zeros((14475, 3)), F, mask)
# frame = 9
# 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.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()
# verts = pred.detach().cpu().numpy()[:,mask,:]
# frame = 9
# 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.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
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 = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), True)['V']
mesh_body_posed = mesh_body_posed[int(mesh_body_posed.shape[0]/2)+frame.frame_nr]
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
def calculate_average(folder, out_folder, frames_file='frames_train_skirt.txt'):
from custom_utils.custom_utils import UV2Mesh
import plotly.express as ex
import plotly.graph_objects as go
frames = []
with open(os.path.join(out_folder, frames_file), 'r') as f:
for line in f:
line = line.strip()
id, frame_nr, garment = line.split(' ')
frames.append(Frame(id, int(frame_nr), garment))
uv_sum = np.zeros((512, 256, 3), dtype=np.float32)
mask_sum = np.zeros((512, 256), dtype=np.float32)
for frame in tqdm(frames):
uv = np.load(os.path.join(out_folder, 'uv_unposed_no_shape', f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))['uv_unposed']
mask = np.load(os.path.join(out_folder, 'uv_mask', f'{frame.id}_{frame.garment}.npz'))['uv_mask']
uv_sum += uv
mask_sum += mask
mask_sum[mask_sum == 0] = 1
avg_garment = uv_sum / mask_sum[:, :, None]
avg_mesh = UV2Mesh()(torch.from_numpy(avg_garment.transpose(2,0,1))[None,:,:,:]).numpy()[0]
F = np.load('smpl/smpl_faces.npy')
fig = go.Figure(data=[go.Mesh3d(x=avg_mesh[:,0], y=avg_mesh[:,2], z=avg_mesh[:,1], i=F[:,0], j=F[:,1], k=F[:,2], color='lightpink', 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()
np.save('custom_data/avg/train_skirt_no_shape_avg_uv.npy', avg_garment)
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'
smpl = SMPL(path='smpl', device='cuda')
folder = '/home/workspace/data/dataset/cloth3d/train'
# clean_dataset(folder)
# generate_frames_file(folder, 'frames_all.txt')
# 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)
calculate_average(folder, out_folder)
# 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)
create_extras_file(folder)
data/registration.py
-461
View File
@@ -1,461 +0,0 @@
import os
import sys
import random
import numpy as np
import plotly.graph_objects as go
import torch
from torch.nn.functional import mse_loss
from pytorch3d.structures import Meshes
from pytorch3d.loss import mesh_laplacian_smoothing, mesh_edge_loss, mesh_normal_consistency, chamfer_distance
from pymeshfix import MeshFix
# os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
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,
read_curvature_frame,
writeOBJ,
readOBJ,
quads2tris
)
from smpl.smpl_torch import SMPL
# from custom_utils.custom_utils import _maskFaces
device = 'cuda'
smpl = SMPL(device, 'smpl')
import torch.nn.functional as F
from pytorch3d.structures import Meshes
def find_closest_vertices(mesh1: Meshes, mesh2: Meshes, vertex_indices: torch.Tensor) -> torch.Tensor:
"""
Given two PyTorch3D meshes and a set of vertex indices for mesh1,
find the closest vertices from mesh2 to these vertices.
Args:
mesh1: A PyTorch3D Meshes object representing the first mesh.
mesh2: A PyTorch3D Meshes object representing the second mesh.
vertex_indices: A tensor of shape (N,) representing the indices of the vertices in mesh1.
Returns:
A tensor of shape (N,) representing the indices of the closest vertices in mesh2.
"""
# Extract the vertices of mesh1 and mesh2.
vertices1 = mesh1.verts_packed()
vertices2 = mesh2.verts_packed()
# Extract the vertices of mesh1 that correspond to the given indices.
vertices1_selected = vertices1[vertex_indices]
# Compute the pairwise distances between the selected vertices of mesh1 and all vertices of mesh2.
distances = F.pairwise_distance(vertices1_selected.unsqueeze(1), vertices2.unsqueeze(0), p=2)
# Find the indices of the closest vertices in mesh2.
closest_indices = torch.argmin(distances, dim=1)
return closest_indices
def _maskFaces(V, F, mask, return_unfixed=False):
V_mask = mask
F_mask = np.take(V_mask, F.reshape(-1), axis=0)
F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
V_ind = np.zeros(V_mask.shape)
V_ind[V_mask] = np.arange(V_mask.sum())
F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
V_new = V[V_mask, :]
if return_unfixed:
return V_new, F_new, F[F_mask, :]
else:
return V_new, F_new
def get_frames(folder, frames_file):
frames = []
with open(os.path.join(folder, frames_file), 'r') as f:
for line in f:
frames.append(line.strip())
return frames
def plot(mesh, faces):
fig = go.Figure()
fig.add_trace(go.Mesh3d(
x=mesh[:,0],
y=mesh[:,1],
z=mesh[:,2],
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),)
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
)
helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,1],),
zaxis = dict(nticks=4, range=[-1,1],),)
fig.update_layout(scene = helper,
)
fig.show()
def build_lookup_table(repeated):
# iterate through the list and create a dict from each sub array in the list where each dict has one element of the sub array as key and the other elements as value as a list and repeat for each element
lookup = {}
for i in range(len(repeated)):
for j in range(len(repeated[i])):
if len(repeated[i]) == 1:
continue
lookup[repeated[i][j]] = [repeated[i][k] for k in range(len(repeated[i]))] # if k != j
return lookup
from data.cloth3d_io import readPC2Frame
import scipy.io as sio
import scipy
import scipy.spatial as spatial
def local_read_posed_garment_frame(folder, id, garment, frame, pose=None, deorient=False):
pc16_path = os.path.join(folder, id, garment + '.pc16')
verts = readPC2Frame(pc16_path, frame, True)
map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1))
if deorient:
R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
# R = spatial.transform.Rotation.from_euler([0,0,-zrot])
verts = R.apply(verts)
return verts, map
def local_read_posed_body_frame(folder, id, frame, pose=None, deorient=False):
pc16_path = os.path.join(folder, id, 'Body_posed.pc16')
verts = readPC2Frame(pc16_path, frame, True)
if deorient:
R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
# R = spatial.transform.Rotation.from_euler([0,0,-zrot])
verts = R.apply(verts)
return verts
def edge_loss_custom(x, y, edges):
x_e = x[:,edges[:,0],:] - x[:,edges[:,1],:]
x_e = torch.sum(x_e ** 2, -1)
x_e = torch.sqrt(x_e + 1e-9)
y_e = y[:,edges[:,0],:] - y[:,edges[:,1],:]
y_e = torch.sum(y_e ** 2, -1)
y_e = torch.sqrt(y_e + 1e-9)
d_e = y_e - x_e
loss = torch.mean(d_e ** 2, dim=-1)[0]
return loss
def crotch(mesh, body, crotch_idx, nr=75):
crotch_middle_on_body = crotch_idx
# create kdtree for body and mesh
body_tree = spatial.cKDTree(body)
mesh_tree = spatial.cKDTree(mesh)
# find closest point on mesh from crotch middle on body along the z axis
crotch_middle_on_mesh = find_closest_point_along_z(body[crotch_middle_on_body], mesh)
crotch_body = body_tree.query(body[crotch_middle_on_body], k=nr)[1]
crotch_mesh = mesh_tree.query(mesh[crotch_middle_on_mesh], k=nr)[1]
body_verts = body[crotch_body]
mesh_verts = mesh[crotch_mesh] - (mesh[crotch_middle_on_mesh] - body[crotch_middle_on_body])
mesh_verts = spatial.cKDTree(mesh_verts)
closest_mesh_verts = mesh_verts.query(body_verts, k=1)[1]
crotch_mesh = crotch_mesh[closest_mesh_verts]
return crotch_body, crotch_mesh
def find_closest_point_along_z(point, verts, eps=0.01):
min_dist = np.inf
closest_point = None
for i, v in enumerate(verts):
dist = np.linalg.norm(point - v)
if dist < min_dist and np.abs(point[0] - v[0]) < eps and np.abs(point[1] - v[1]) < eps:
min_dist = dist
closest_point = i
return closest_point
def mesh_edge_lengths(mesh: Meshes) -> torch.Tensor:
"""
Given a PyTorch3D Meshes object, compute the lengths of all edges in the mesh.
Args:
mesh: A PyTorch3D Meshes object representing the mesh.
Returns:
A tensor of shape (E,) representing the lengths of all edges in the mesh,
where E is the total number of edges in the mesh.
"""
# Extract the vertices and faces of the mesh.
vertices = mesh.verts_packed()
faces = mesh.faces_packed()
# Compute the pairwise distances between all pairs of vertices that form an edge.
edge_indices = faces.view(-1, 3)
edge_vertices = vertices[edge_indices]
edge_lengths = torch.norm(edge_vertices[:, 1:] - edge_vertices[:, :-1], dim=2)
return edge_lengths
def register(folder, frame):
id, frame_nr, garment = frame.split()
frame_nr = int(frame_nr)
info = loadInfo(os.path.join(folder, id, 'info.mat'))
raw_mask = read_garment_mask(folder, id, garment, return_raw=True)
# convert it to indices
mask = raw_mask > 0
mesh_posed, map = local_read_posed_garment_frame(folder, id, garment, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
body_posed = local_read_posed_body_frame(folder, id, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
template, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj'))
faces = quads2tris(faces)
# boundaries = raw_mask == 2
# # invert map and use it to invert boudaries
# map_inv = np.zeros(map.max()+1, dtype=np.int32)
# map_inv[map] = np.arange(len(map))
# boundaries = boundaries[map_inv]
# # get the boundary indices
# mesh1_boundaries = torch.tensor(np.where(boundaries)[0]).to(device)
# pose = torch.zeros(1, 72).to('cuda')
shape = torch.zeros(1, 11).to(device)
shape[:, :10] = torch.tensor(info['shape']).to(device)
shape[:, -1] = torch.tensor(info['gender']).to(device)
# for star pose
# pose[:, 5] = 3.1415/4
# pose[:,8] = -3.1415/4
STAR_POSE = np.zeros((24, 3), np.float32)
STAR_POSE[1,2] = .15
STAR_POSE[2,2] = -.15
pose = torch.tensor(STAR_POSE.flatten()[None,:]).to(device)
star_body = smpl(pose, shape, remove=True)[0][0].cpu().numpy() - smpl(pose, shape, remove=True)[4][0].cpu().numpy()
# rotate star body with 90 degrees around x axis
R = spatial.transform.Rotation.from_euler('x', np.pi/2)
star_body = R.apply(star_body).astype(np.float32)
is_skirt = (garment == 'Dress' or garment == 'Skirt')
skirt_mod = '_skirt' if is_skirt else ''
F = np.load(f'smpl/smpl_faces{skirt_mod}.npy').astype(np.int32)
##################### mask body #####################
new_crotch = 4729
# star_body = star_body[mask]
# new_crotch = np.zeros((14475), bool)
# new_crotch[4729] = True
# new_crotch = new_crotch[mask]
# new_crotch = np.where(new_crotch)[0][0]
# _, F = _maskFaces(np.zeros((14475, 3)), F, mask)
# plot(template, faces)
# v = mesh_posed[mask,:]
##################### boundaries #####################
# meshfix = MeshFix(template, faces)
# holes = meshfix.extract_holes()
# new_border = []
# for border in np.array(holes.points):
# # find the index of the vertex in the original mesh
# ind = np.argwhere(np.all(template == border, axis=1))
# new_border.extend(ind.ravel())
# mesh1_boundaries = torch.tensor(new_border).to(device)
# np.save('temp/body_boundaries.npy', new_border)
mesh1_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
# meshfix = MeshFix(template, faces)
# holes = meshfix.extract_holes()
# new_border = []
# for border in np.array(holes.points):
# # find the index of the vertex in the original mesh
# ind = np.argwhere(np.all(template == border, axis=1))
# new_border.extend(ind.ravel())
# mesh1_boundaries = torch.tensor(new_border).to(device)
# np.save('temp/body_boundaries.npy', new_border)
# mesh1_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
##################### crotch #####################
crotch_body, crotch_mesh = crotch(template, star_body, crotch_idx=new_crotch, nr=50)
# export star_body as obj
# writeOBJ('temp/star_body.obj', star_body, F)
template = torch.tensor(template).to(device)
template_orig = template.clone()
faces = torch.tensor(faces).to(device)
star_body = torch.tensor(star_body).to(device)
F = torch.tensor(F).to(device)
# verts = template.cpu().numpy()
# star_body = star_body.cpu().numpy()
# faces = faces.detach().cpu().numpy()
# F = F.detach().cpu().numpy()
# fig = go.Figure(data=[
# go.Mesh3d(
# x=verts[:,0], y=verts[:,1], z=verts[:,2],
# i=faces[:,0], j=faces[:,1], k=faces[:,2],
# color='lightpink', opacity=0.50),
# go.Mesh3d(
# x=star_body[:,0], y=star_body[:,1], z=star_body[:,2],
# i=F[:,0], j=F[:,1], k=F[:,2],
# color='lightblue', opacity=0.50),
# go.Scatter3d(
# x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
# mode='markers', marker=dict(size=2, color='red')),
# go.Scatter3d(
# x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
# mode='markers', marker=dict(size=2, color='blue')),
# ])
# # make it cuve and limit to -1 and -1
# fig.update_layout(scene = dict(
# xaxis = dict(range=[-1,1],),
# yaxis = dict(range=[-1,1],),
# zaxis = dict(range=[-1,1],),),)
# fig.show()
# Define the meshes to optimize
mesh1 = Meshes(verts=template[None,:,:], faces=faces[None,:,:])
mesh2 = Meshes(verts=star_body[None,:,:], faces=F[None,:,:])
mesh2_boundaries = find_closest_vertices(mesh1, mesh2, mesh1_boundaries)
avg_edge_length = mesh_edge_lengths(mesh1).mean()
# Define the optimizer and learning rate
mesh1.verts_packed().requires_grad = True
optimizer = torch.optim.Adam([mesh1.verts_packed()], lr=0.0001)
# optimizer = torch.optim.SGD([mesh1.verts_packed()], lr=0.001)
# Define the loss functions
def laplacian_smoothing_loss(mesh):
return mesh_laplacian_smoothing(mesh, method="uniform") # cotcurv uniform cot
def edge_loss(mesh):
return mesh_edge_loss(mesh, target_length=avg_edge_length)
def normal_consistency_loss(mesh):
return mesh_normal_consistency(mesh)
# Optimization loop
for i in range(10000):
optimizer.zero_grad()
# find boundaries with the find holes function?
# Forward pass
mesh1_out = Meshes(verts=[mesh1.verts_packed()], faces=[mesh1.faces_packed()])
ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list()), torch.stack(mesh2.verts_list()))
bound_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,mesh1_boundaries,:], torch.stack(mesh2.verts_list())[:,mesh2_boundaries,:])
# crotch_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:], torch.stack(mesh2.verts_list())[:,crotch_body,:],)
crotch_ch_loss = torch.mean((torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:] - torch.stack(mesh2.verts_list())[:,crotch_body,:])**2, dim=2).mean()
loss: torch.Tensor = ch_loss * 10 + (bound_ch_loss + crotch_ch_loss) * 5
# loss: torch.Tensor = ch_loss + ch_norm_loss + (bound_ch_loss + bound_ch_norm_loss) * 1
loss += laplacian_smoothing_loss(mesh1_out) * 0.05
loss += edge_loss(mesh1_out) * 100
# loss += edge_loss_custom(torch.stack(mesh1_out.verts_list()), template_orig[None,:,:], mesh1.edges_packed()) * 100
loss += normal_consistency_loss(mesh1_out) * 0.1
# Backward pass
loss.backward()
optimizer.step()
# Print the loss
if i % 10 == 0:
print(f"Loss: {loss.item()} chamfer: {ch_loss.item()} ")
# create a plotly plot with the body and the mesh1_out
verts = mesh1_out.verts_packed().detach().cpu().numpy()
star_body = star_body.detach().cpu().numpy()
faces = faces.detach().cpu().numpy()
F = F.detach().cpu().numpy()
mesh1_boundaries = mesh1_boundaries.detach().cpu().numpy()
mesh2_boundaries = mesh2_boundaries.detach().cpu().numpy()
fig = go.Figure(data=[
go.Mesh3d(
x=verts[:,0], y=verts[:,1], z=verts[:,2],
i=faces[:,0], j=faces[:,1], k=faces[:,2],
color='lightpink', opacity=0.80),
go.Mesh3d(
x=star_body[:,0], y=star_body[:,1], z=star_body[:,2],
i=F[:,0], j=F[:,1], k=F[:,2],
color='lightblue', opacity=0.50),
go.Scatter3d(
x=verts[mesh1_boundaries,0], y=verts[mesh1_boundaries,1], z=verts[mesh1_boundaries,2],
mode='markers', marker=dict(size=2, color='red')),
go.Scatter3d(
x=star_body[mesh2_boundaries,0], y=star_body[mesh2_boundaries,1], z=star_body[mesh2_boundaries,2],
mode='markers', marker=dict(size=2, color='blue')),
go.Scatter3d(
x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
mode='markers', marker=dict(size=2, color='red')),
go.Scatter3d(
x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
mode='markers', marker=dict(size=2, color='blue')),
])
# make it cuve and limit to -1 and -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/code/cloth3d/train'
# folder = r'D:\Work\code\cloth3d\train'
# frames = [f'00761_{i}' for i in range(300)]
# folder = r'Z:\Work\Huni\Barcelona\Dataset\cloth3d\train'
frames = get_frames(folder, 'frames_all.txt')
root = '/'.join(folder.split('/')[:-1])
type = folder.split('/')[-1]
use_no_orient = False
no_shape = False
is_skirt = False
skirt_modifier = '_skirt' if is_skirt else ''
orient_mod = '_no_orient' if use_no_orient else ''
shape_mod = '_no_shape' if no_shape else ''
# mask = np.load('/data-net/dataset/00761/mesh_mask/00761.npy')
# faces = np.load('/data-net/dataset/00761/faces/00761.npy')
for frame in frames:
frame = random.choice(frames)
# frame = '00761 0 Tshirt'
frame = '01893 256 Jumpsuit'
register(folder, frame)
break
pass
data/registration_post.py
-482
View File
@@ -1,482 +0,0 @@
import os
import sys
import random
import numpy as np
import plotly.graph_objects as go
import torch
from torch.nn.functional import mse_loss
from pytorch3d.structures import Meshes
from pytorch3d.loss import mesh_laplacian_smoothing, mesh_edge_loss, mesh_normal_consistency, chamfer_distance
from pymeshfix import MeshFix
# os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
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,
read_curvature_frame,
writeOBJ,
readOBJ,
quads2tris
)
from smpl.smpl_torch import SMPL
# from custom_utils.custom_utils import _maskFaces
device = 'cuda'
smpl = SMPL(device, 'smpl')
import torch.nn.functional as F
from pytorch3d.structures import Meshes
def find_closest_vertices(mesh1: Meshes, mesh2: Meshes, vertex_indices: torch.Tensor) -> torch.Tensor:
"""
Given two PyTorch3D meshes and a set of vertex indices for mesh1,
find the closest vertices from mesh2 to these vertices.
Args:
mesh1: A PyTorch3D Meshes object representing the first mesh.
mesh2: A PyTorch3D Meshes object representing the second mesh.
vertex_indices: A tensor of shape (N,) representing the indices of the vertices in mesh1.
Returns:
A tensor of shape (N,) representing the indices of the closest vertices in mesh2.
"""
# Extract the vertices of mesh1 and mesh2.
vertices1 = mesh1.verts_packed()
vertices2 = mesh2.verts_packed()
# Extract the vertices of mesh1 that correspond to the given indices.
vertices1_selected = vertices1[vertex_indices]
# Compute the pairwise distances between the selected vertices of mesh1 and all vertices of mesh2.
distances = F.pairwise_distance(vertices1_selected.unsqueeze(1), vertices2.unsqueeze(0), p=2)
# Find the indices of the closest vertices in mesh2.
closest_indices = torch.argmin(distances, dim=1)
return closest_indices
def _maskFaces(V, F, mask, return_unfixed=False):
V_mask = mask
F_mask = np.take(V_mask, F.reshape(-1), axis=0)
F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
V_ind = np.zeros(V_mask.shape)
V_ind[V_mask] = np.arange(V_mask.sum())
F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
V_new = V[V_mask, :]
if return_unfixed:
return V_new, F_new, F[F_mask, :]
else:
return V_new, F_new
def get_frames(folder, frames_file):
frames = []
with open(os.path.join(folder, frames_file), 'r') as f:
for line in f:
frames.append(line.strip())
return frames
def plot(mesh, faces):
fig = go.Figure()
fig.add_trace(go.Mesh3d(
x=mesh[:,0],
y=mesh[:,1],
z=mesh[:,2],
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),)
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
)
helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,1],),
zaxis = dict(nticks=4, range=[-1,1],),)
fig.update_layout(scene = helper,
)
fig.show()
def build_lookup_table(repeated):
# iterate through the list and create a dict from each sub array in the list where each dict has one element of the sub array as key and the other elements as value as a list and repeat for each element
lookup = {}
for i in range(len(repeated)):
for j in range(len(repeated[i])):
if len(repeated[i]) == 1:
continue
lookup[repeated[i][j]] = [repeated[i][k] for k in range(len(repeated[i]))] # if k != j
return lookup
from data.cloth3d_io import readPC2Frame
import scipy.io as sio
import scipy
import scipy.spatial as spatial
def local_read_posed_garment_frame(folder, id, garment, frame, pose=None, deorient=False):
pc16_path = os.path.join(folder, id, garment + '.pc16')
verts = readPC2Frame(pc16_path, frame, True)
map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1))
if deorient:
R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
# R = spatial.transform.Rotation.from_euler([0,0,-zrot])
verts = R.apply(verts)
return verts, map
def local_read_posed_body_frame(folder, id, frame, pose=None, deorient=False):
pc16_path = os.path.join(folder, id, 'Body_posed.pc16')
verts = readPC2Frame(pc16_path, frame, True)
if deorient:
R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
# R = spatial.transform.Rotation.from_euler([0,0,-zrot])
verts = R.apply(verts)
return verts
def edge_loss_custom(x, y, edges):
x_e = x[:,edges[:,0],:] - x[:,edges[:,1],:]
x_e = torch.sum(x_e ** 2, -1)
x_e = torch.sqrt(x_e + 1e-9)
y_e = y[:,edges[:,0],:] - y[:,edges[:,1],:]
y_e = torch.sum(y_e ** 2, -1)
y_e = torch.sqrt(y_e + 1e-9)
d_e = y_e - x_e
loss = torch.mean(d_e ** 2, dim=-1)[0]
return loss
def crotch(mesh, body, crotch_idx, nr=75):
crotch_middle_on_body = crotch_idx
# create kdtree for body and mesh
body_tree = spatial.cKDTree(body)
mesh_tree = spatial.cKDTree(mesh)
# find closest point on mesh from crotch middle on body along the z axis
crotch_middle_on_mesh = find_closest_point_along_z(body[crotch_middle_on_body], mesh)
crotch_body = body_tree.query(body[crotch_middle_on_body], k=nr)[1]
crotch_mesh = mesh_tree.query(mesh[crotch_middle_on_mesh], k=nr)[1]
body_verts = body[crotch_body]
mesh_verts = mesh[crotch_mesh] - (mesh[crotch_middle_on_mesh] - body[crotch_middle_on_body])
mesh_verts = spatial.cKDTree(mesh_verts)
closest_mesh_verts = mesh_verts.query(body_verts, k=1)[1]
crotch_mesh = crotch_mesh[closest_mesh_verts]
return crotch_body, crotch_mesh
def find_closest_point_along_z(point, verts, eps=0.01):
min_dist = np.inf
closest_point = None
for i, v in enumerate(verts):
dist = np.linalg.norm(point - v)
if dist < min_dist and np.abs(point[0] - v[0]) < eps and np.abs(point[1] - v[1]) < eps:
min_dist = dist
closest_point = i
return closest_point
def mesh_edge_lengths(mesh: Meshes) -> torch.Tensor:
"""
Given a PyTorch3D Meshes object, compute the lengths of all edges in the mesh.
Args:
mesh: A PyTorch3D Meshes object representing the mesh.
Returns:
A tensor of shape (E,) representing the lengths of all edges in the mesh,
where E is the total number of edges in the mesh.
"""
# Extract the vertices and faces of the mesh.
vertices = mesh.verts_packed()
faces = mesh.faces_packed()
# Compute the pairwise distances between all pairs of vertices that form an edge.
edge_indices = faces.view(-1, 3)
edge_vertices = vertices[edge_indices]
edge_lengths = torch.norm(edge_vertices[:, 1:] - edge_vertices[:, :-1], dim=2)
return edge_lengths
# try to optimize the already registered garment ??????????
def register(folder, frame):
id, frame_nr, garment = frame.split()
frame_nr = int(frame_nr)
info = loadInfo(os.path.join(folder, id, 'info.mat'))
raw_mask = read_garment_mask(folder, id, garment, return_raw=True)
# convert it to indices
mask = raw_mask > 0
mesh_posed, map = local_read_posed_garment_frame(folder, id, garment, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
body_posed = local_read_posed_body_frame(folder, id, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
template, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj'))
faces = quads2tris(faces)
boundaries = raw_mask == 2
registered_boundaries = torch.tensor(np.where(boundaries[mask])[0]).to(device)
# invert map and use it to invert boudaries
# map_inv = np.zeros(map.max()+1, dtype=np.int32)
# map_inv[map] = np.arange(len(map))
# boundaries = boundaries[map_inv]
# # get the boundary indices
# mesh1_boundaries = torch.tensor(np.where(boundaries)[0]).to(device)
# # pose = torch.zeros(1, 72).to('cuda')
# shape = torch.zeros(1, 11).to(device)
# shape[:, :10] = torch.tensor(info['shape']).to(device)
# shape[:, -1] = torch.tensor(info['gender']).to(device)
# # for star pose
# # pose[:, 5] = 3.1415/4
# # pose[:,8] = -3.1415/4
# STAR_POSE = np.zeros((24, 3), np.float32)
# STAR_POSE[1,2] = .15
# STAR_POSE[2,2] = -.15
# pose = torch.tensor(STAR_POSE.flatten()[None,:]).to(device)
# star_body = smpl(pose, shape, remove=True)[0][0].cpu().numpy() - smpl(pose, shape, remove=True)[4][0].cpu().numpy()
# # rotate star body with 90 degrees around x axis
# R = spatial.transform.Rotation.from_euler('x', np.pi/2)
# star_body = R.apply(star_body).astype(np.float32)
is_skirt = (garment == 'Dress' or garment == 'Skirt')
skirt_mod = '_skirt' if is_skirt else ''
F = np.load(f'smpl/smpl_faces{skirt_mod}.npy').astype(np.int32)
##################### mask body #####################
new_crotch = 4729
# star_body = star_body[mask]
new_crotch = np.zeros((14475), bool)
new_crotch[4729] = True
new_crotch = new_crotch[mask]
new_crotch = np.where(new_crotch)[0][0]
_, F = _maskFaces(np.zeros((14475, 3)), F, mask)
registered = template[map][mask]
# plot(template, faces)
# v = mesh_posed[mask,:]
##################### boundaries #####################
# meshfix = MeshFix(template, faces)
# holes = meshfix.extract_holes()
# new_border = []
# for border in np.array(holes.points):
# # find the index of the vertex in the original mesh
# ind = np.argwhere(np.all(template == border, axis=1))
# new_border.extend(ind.ravel())
# mesh1_boundaries = torch.tensor(new_border).to(device)
# np.save('temp/body_boundaries.npy', new_border)
template_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
# meshfix = MeshFix(template, faces)
# holes = meshfix.extract_holes()
# new_border = []
# for border in np.array(holes.points):
# # find the index of the vertex in the original mesh
# ind = np.argwhere(np.all(template == border, axis=1))
# new_border.extend(ind.ravel())
# mesh1_boundaries = torch.tensor(new_border).to(device)
# np.save('temp/body_boundaries.npy', new_border)
# mesh1_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
##################### crotch #####################
# crotch_body, crotch_mesh = crotch(template, star_body, crotch_idx=new_crotch, nr=50)
# export star_body as obj
# writeOBJ('temp/star_body.obj', star_body, F)
template = torch.tensor(template).to(device)
template_orig = template.clone()
faces = torch.tensor(faces).to(device)
registered = torch.tensor(registered).to(device)
F = torch.tensor(F).to(device)
# verts = template.cpu().numpy()
# star_body = star_body.cpu().numpy()
# faces = faces.detach().cpu().numpy()
# F = F.detach().cpu().numpy()
# fig = go.Figure(data=[
# go.Mesh3d(
# x=verts[:,0], y=verts[:,1], z=verts[:,2],
# i=faces[:,0], j=faces[:,1], k=faces[:,2],
# color='lightpink', opacity=0.50),
# go.Mesh3d(
# x=star_body[:,0], y=star_body[:,1], z=star_body[:,2],
# i=F[:,0], j=F[:,1], k=F[:,2],
# color='lightblue', opacity=0.50),
# go.Scatter3d(
# x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
# mode='markers', marker=dict(size=2, color='red')),
# go.Scatter3d(
# x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
# mode='markers', marker=dict(size=2, color='blue')),
# ])
# # make it cuve and limit to -1 and -1
# fig.update_layout(scene = dict(
# xaxis = dict(range=[-1,1],),
# yaxis = dict(range=[-1,1],),
# zaxis = dict(range=[-1,1],),),)
# fig.show()
# Define the meshes to optimize
mesh1 = Meshes(verts=registered[None,:,:], faces=F[None,:,:])
mesh2 = Meshes(verts=template[None,:,:], faces=faces[None,:,:])
mesh1_boundaries = registered_boundaries
mesh2_boundaries = template_boundaries
# mesh2_boundaries = find_closest_vertices(mesh1, mesh2, mesh1_boundaries)
avg_edge_length = mesh_edge_lengths(mesh1).mean()
# Define the optimizer and learning rate
mesh1.verts_packed().requires_grad = True
optimizer = torch.optim.Adam([mesh1.verts_packed()], lr=0.0001)
# optimizer = torch.optim.SGD([mesh1.verts_packed()], lr=0.001)
# Define the loss functions
def laplacian_smoothing_loss(mesh):
return mesh_laplacian_smoothing(mesh, method="uniform") # cotcurv uniform cot
def edge_loss(mesh):
return mesh_edge_loss(mesh, target_length=avg_edge_length)
def normal_consistency_loss(mesh):
return mesh_normal_consistency(mesh)
# Optimization loop
for i in range(10000):
optimizer.zero_grad()
# find boundaries with the find holes function?
# Forward pass
mesh1_out = Meshes(verts=[mesh1.verts_packed()], faces=[mesh1.faces_packed()])
ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list()), torch.stack(mesh2.verts_list()))
bound_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,mesh1_boundaries,:], torch.stack(mesh2.verts_list())[:,mesh2_boundaries,:])
# crotch_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:], torch.stack(mesh2.verts_list())[:,crotch_body,:],)
# crotch_ch_loss = torch.mean((torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:] - torch.stack(mesh2.verts_list())[:,crotch_body,:])**2, dim=2).mean()
loss: torch.Tensor = ch_loss * 3 + (bound_ch_loss) * 20
# loss: torch.Tensor = ch_loss + ch_norm_loss + (bound_ch_loss + bound_ch_norm_loss) * 1
lap_loss = laplacian_smoothing_loss(mesh1_out)
e_loss = edge_loss(mesh1_out)
# loss += edge_loss_custom(torch.stack(mesh1_out.verts_list()), template_orig[None,:,:], mesh1.edges_packed()) * 100
norm_loss = normal_consistency_loss(mesh1_out)
loss += lap_loss * 0.05 + e_loss * 100 + norm_loss * 0.1
# Backward pass
loss.backward()
optimizer.step()
# Print the loss
if i % 10 == 0:
print(f"it: {i} : loss: {loss:.5f} chamfer: {ch_loss:.6f} bound_chamfer: {bound_ch_loss:.6f} laplacian: {lap_loss:.6f} edge: {e_loss:.6f} normal: {norm_loss:.5f}")
# create a plotly plot with the body and the mesh1_out
registered = mesh1_out.verts_packed().detach().cpu().numpy()
template = mesh2.verts_packed().detach().cpu().numpy()
faces = faces.detach().cpu().numpy()
F = F.detach().cpu().numpy()
# faces, F = F, faces
# verts, star_body = star_body, verts
mesh1_boundaries, mesh2_boundaries = mesh1_boundaries.detach().cpu().numpy(), mesh2_boundaries.detach().cpu().numpy()
# mesh2_boundaries = mesh1_boundaries.detach().cpu().numpy()
fig = go.Figure(data=[
go.Mesh3d(
x=template[:,0], y=template[:,1], z=template[:,2],
i=faces[:,0], j=faces[:,1], k=faces[:,2],
color='lightblue', opacity=0.80),
go.Mesh3d(
x=registered[:,0], y=registered[:,1], z=registered[:,2],
i=F[:,0], j=F[:,1], k=F[:,2],
color='lightpink', opacity=0.50),
go.Scatter3d(
x=registered[mesh1_boundaries,0], y=registered[mesh1_boundaries,1], z=registered[mesh1_boundaries,2],
mode='markers', marker=dict(size=2, color='red')),
go.Scatter3d(
x=template[mesh2_boundaries,0], y=template[mesh2_boundaries,1], z=template[mesh2_boundaries,2],
mode='markers', marker=dict(size=2, color='blue')),
# go.Scatter3d(
# x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
# mode='markers', marker=dict(size=2, color='red')),
# go.Scatter3d(
# x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
# mode='markers', marker=dict(size=2, color='blue')),
])
# make it cuve and limit to -1 and -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()
# create a plotly wireframe plot of registered
fig = go.Figure()
for edge in mesh1_out.edges_packed().detach().cpu().numpy():
fig.add_trace(go.Scatter3d(
x=registered[edge,0], y=registered[edge,1], z=registered[edge,2],
mode='lines', line=dict(color='red', width=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/code/cloth3d/train'
# folder = r'D:\Work\code\cloth3d\train'
# frames = [f'00761_{i}' for i in range(300)]
# folder = r'Z:\Work\Huni\Barcelona\Dataset\cloth3d\train'
frames = get_frames(folder, 'frames_all.txt')
root = '/'.join(folder.split('/')[:-1])
type = folder.split('/')[-1]
use_no_orient = False
no_shape = False
is_skirt = False
skirt_modifier = '_skirt' if is_skirt else ''
orient_mod = '_no_orient' if use_no_orient else ''
shape_mod = '_no_shape' if no_shape else ''
# mask = np.load('/data-net/dataset/00761/mesh_mask/00761.npy')
# faces = np.load('/data-net/dataset/00761/faces/00761.npy')
for frame in frames:
frame = random.choice(frames)
# frame = '00761 0 Tshirt'
frame = '01893 256 Jumpsuit'
register(folder, frame)
break
pass
data/registration_weights.py
-1071
View File
File diff suppressed because it is too large Load Diff
data/utils.py
+4 -4
View File
@@ -1,10 +1,10 @@
import tables as tb
class Frame:
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 __init__(self, id, frame_nr, garment=''):
self.id = id
self.frame_nr = frame_nr
self.garment = garment
def __hash__(self):
return hash((self.id, self.frame_nr, self.garment))
docker/Dockerfile2
+4 -10
View File
@@ -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.1-cuda11.7-cudnn8-devel
# FROM pytorch/pytorch:2.0.0-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 install -y sudo git apt-utils htop
RUN apt-get update && apt-get upgrade -y && 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,12 +25,6 @@ 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
@@ -39,8 +33,8 @@ WORKDIR /workspace
ENV PATH="$PATH:/home/$USER_NAME/.local/bin"
RUN pip install --upgrade pip
COPY requirements.txt /tmp/
RUN pip install torch-cluster -f https://data.pyg.org/whl/torch-2.0.0+cu117.html
RUN pip install --requirement /tmp/requirements.txt
RUN pip install torch-cluster -f https://data.pyg.org/whl/torch-2.0.0+cu117.html
# defining the volumes that will be mounted
docker/docker-compose.yml
-23
View File
@@ -34,29 +34,6 @@ services:
- '/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]
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:
docker/requirements.txt
+3
View File
@@ -1,10 +1,13 @@
torch==2.0.0
numpy
matplotlib
opencv-python
wandb
scipy
plotly
pandas
cmapy
pytorch3d
tables
numpy-stl
torchmetrics
losses/losses.py
+5 -5
View File
@@ -18,12 +18,12 @@ from torch_cluster import knn
smpl_to_keep_orig = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
smpl_to_keep_orig = smpl_to_keep_orig[smpl_to_keep_orig < 6890]
# to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
F_coll = torch.tensor(np.load("smpl/smpl_faces.npy"), device='cuda')
F_coll_skirt = torch.tensor(np.load("smpl/smpl_faces_skirt.npy"), device='cuda')
to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
F_coll = torch.tensor(np.load("smpl/smpl_faces.npy")[to_keep,:], device='cuda')
F_coll_skirt = torch.tensor(np.load("smpl/smpl_faces_skirt.npy")[to_keep,:], device='cuda')
F_coll_orig = torch.tensor(np.load("smpl/smpl_faces_orig.npy")[:6890,:], device='cuda')
F_norm = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[None,:,:], [4,1,1]), device='cuda')
F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[None,:,:], [4,1,1]), device='cuda')
F_norm = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
F_norm_orig = torch.tensor(np.tile(np.load("smpl/smpl_faces_orig.npy")[:6890,:][None,:,:], [4,1,1]), device='cuda')
F_norm_orig = F_norm_orig[:,smpl_to_keep_orig,:]
models/base.py
+2 -2
View File
@@ -35,10 +35,10 @@ class VariationalEncoder(nn.Module):
class AEEncoder(nn.Module):
def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7, activation='tanh', is_label=False, nr_of_labels=None, device='cuda'):
def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7, activation='tanh', is_label=False, nr_of_labels=None, device='cuda', input_features=3):
super(AEEncoder, self).__init__()
self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_labels, device=device)
self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_labels, device=device, input_features=input_features)
if activation == 'tanh':
self.act = nn.Tanh()
elif activation == 'sigmoid':
models/cond_encoder.py
+1 -1
View File
@@ -8,7 +8,7 @@ class AutoEncoder(nn.Module):
def __init__(self, config):
super(AutoEncoder, self).__init__()
self.encoder : AEEncoder = AEEncoder(config['encoder_block_settings'], config['kernel_size'], config['activation'], is_label=config['is_label'], device=config['device'], nr_of_labels=7)
self.encoder : AEEncoder = AEEncoder(config['encoder_block_settings'], config['kernel_size'], config['activation'], is_label=config['is_label'], device=config['device'], nr_of_labels=7, input_features=config['input_features'])
self.decoder : Decoder = Decoder(config['decoder_block_settings'], kernel_size=config['kernel_size'], is_label=config['is_label'], device=config['device'], nr_of_classes=6)
models/cvae_model.py
+14 -13
View File
@@ -18,6 +18,7 @@ class CVAE(nn.Module):
super(CVAE, self).__init__()
self.size = config['resolution']
if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'], device=config['device'])
if config['static_model'] != '':
load_model(self.static_enc, config['static_model'] + '_enc')
@@ -26,16 +27,16 @@ class CVAE(nn.Module):
if config['body_model'] != '':
load_model(self.body_enc, config['body_model'] + '_enc')
if config['use_normals']:
self.norm_enc : AEEncoder = AEEncoder(config['cond_ae_norm']['encoder_block_settings'], config['cond_ae_norm']['kernel_size'], config['cond_ae_norm']['activation'], device=config['device'])
if config['normals_model'] != '':
load_model(self.norm_enc, config['normals_model'] + '_enc')
if config['use_curvature']:
self.curv_enc : AEEncoder = AEEncoder(config['cond_ae_curv']['encoder_block_settings'], config['cond_ae_curv']['kernel_size'], config['cond_ae_curv']['activation'], device=config['device'], is_label=True, nr_of_labels=7)
if config['curvature_model'] != '':
load_model(self.curv_enc, config['curvature_model'] + '_enc')
if config['freeze_encoders']:
for param in self.static_enc.parameters(): param.requires_grad = False
for param in self.body_enc.parameters(): param.requires_grad = False
if config['use_normals']:
for param in self.norm_enc.parameters(): param.requires_grad = False
if config['use_curvature']:
for param in self.curv_enc.parameters(): param.requires_grad = False
self.encoder = VariationalEncoder(config['cvae']['encoder_block_settings'], config['cvae']['kernel_size'], config['cvae']['input_features'], device=config['device'])
self.decoder : Decoder = Decoder(config['cvae']['decoder_block_settings'], kernel_size=config['cvae']['kernel_size'], device=config['device'])
@@ -47,7 +48,7 @@ class CVAE(nn.Module):
else:
self.avg_garment = None
self.use_normals = config['use_normals']
self.use_curvature = config['use_curvature']
self.use_offsets_over_body = not config['use_skinning']
if self.use_offsets_over_body:
@@ -97,9 +98,9 @@ class CVAE(nn.Module):
def encode_conds(self, uv_template, uv_body_posed, uv_curv):
z_body = self.body_enc(uv_body_posed)
z_static = self.static_enc(uv_template)
if self.use_normals:
z_norm = self.norm_enc(uv_curv)
z_cond = torch.concat([z_static, z_body, z_norm], dim=1)
if self.use_curvature:
z_curv = self.curv_enc(uv_curv)
z_cond = torch.concat([z_static, z_body, z_curv], dim=1)
else:
z_cond = torch.concat([z_static, z_body], dim=1)
@@ -107,13 +108,13 @@ class CVAE(nn.Module):
def forward(self, inputs):
uv_unposed, uv_template, uv_body_posed, uv_norm, uv_mask = inputs
x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_norm], dim=1)
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], dim=1)
mean, logvar = self.encoder(x)
z: torch.Tensor = self.reparameterize(mean, logvar)
z = z.view((mean.shape[0], *self.latent_shape))
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_norm)
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
z = torch.concat([z, z_cond], dim=1)
x = self.decoder(z)
models/cvae_norm_model.py
-97
View File
@@ -1,97 +0,0 @@
import os
import time
import random
import torch
from torch import nn
import numpy as np
import scipy as sp
# import plotly.graph_objects as go
# import plotly.express as ex
from utils.checkpoints import load_model
from models.base import AEEncoder, VariationalEncoder, Decoder
class CVAENorm(nn.Module):
def __init__(self, config):
super(CVAENorm, self).__init__()
self.size = config['resolution']
self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'], device=config['device'])
if config['static_model'] != '':
load_model(self.static_enc, config['static_model'] + '_enc')
self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], config['cond_ae_body']['kernel_size'], config['cond_ae_body']['activation'], device=config['device'])
if config['body_model'] != '':
load_model(self.body_enc, config['body_model'] + '_enc')
if config['freeze_encoders']:
for param in self.static_enc.parameters(): param.requires_grad = False
for param in self.body_enc.parameters(): param.requires_grad = False
self.encoder = VariationalEncoder(config['cvae']['encoder_block_settings'], config['cvae']['kernel_size'], config['cvae']['input_features'], device=config['device'])
self.decoder : Decoder = Decoder(config['cvae']['decoder_block_settings'], kernel_size=config['cvae']['kernel_size'], device=config['device'])
self.latent_shape = (config['cvae']['latent_dim'], 8, 4)
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None, add_noise=False):
if noise is None:
eps = tf.random.normal(shape=(uv_template.shape[0], 8, 4, self.latent_dim))
else:
eps = noise
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=False)
if add_noise:
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
z = tf.concat([eps, z_cond], axis=-1)
x = self.decoder(z, training=False)
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
if self.use_offsets_over_body:
x = x + uv_body_posed
x = x * uv_mask[:,:,:,tf.newaxis]
return x
def reparameterize(self, mean: torch.Tensor, logvar: torch.Tensor, training: bool = True, fix_noise: bool = False):
if training == True:
if not fix_noise:
eps = torch.randn(size=mean.shape).to(mean.device)
else:
eps = torch.randn(size=(1, mean.shape[1]))
eps = torch.tile(eps, [mean.shape[0],1])
out = eps * torch.exp(logvar * 0.5) + mean
else:
out = torch.exp(logvar * 0.5) + mean
return out
def encode_conds(self, uv_template, uv_body_posed):
z_body = self.body_enc(uv_body_posed)
z_static = self.static_enc(uv_template)
z_cond = torch.concat([z_static, z_body], dim=1)
return z_cond
def forward(self, inputs):
uv_unposed, uv_template, uv_body_posed, uv_mask = inputs
x = torch.concat([uv_unposed, uv_template, uv_body_posed], dim=1)
mean, logvar = self.encoder(x)
z: torch.Tensor = self.reparameterize(mean, logvar)
z = z.view((mean.shape[0], *self.latent_shape))
z_cond = self.encode_conds(uv_template, uv_body_posed)
z = torch.concat([z, z_cond], dim=1)
x = self.decoder(z)
x = x * uv_mask[:,None,:,:]
return x, mean, logvar
models/direct.py
+29
View File
@@ -0,0 +1,29 @@
import torch
from torch import nn
from models.convnext_orig import ConvNeXt, CNBlock, CNBlockConfig
from typing import Any, Callable, List, Optional, Sequence
class DirectModel(nn.Module):
def __init__(self, block_setting: List[CNBlockConfig], kernel_size: int = 7, input_features: int = 16, device: str = 'cuda'):
super(DirectModel, self).__init__()
self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, input_features=input_features, device=device)
self.decoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, device=device, reversed=True)
layers = [CNBlock(block_setting[-1].input_channels, kernel_size=kernel_size) for _ in range(5)]
self.mean_branch = nn.Sequential(*layers).to(device)
layers = [CNBlock(block_setting[-1].input_channels, kernel_size=kernel_size) for _ in range(5)]
self.var_branch = nn.Sequential(*layers).to(device)
self.flatten_mean = nn.Flatten()
self.flatten_var = nn.Flatten()
def forward(self, inputs):
body, template = inputs
x = self.encoder(x)
return mean
run.py
+1 -1
View File
@@ -1,6 +1,6 @@
import os
import sys
from configs.cvae_norm_1 import config
from configs.direct_1 import config
import argparse
# parser = argparse.ArgumentParser(description='')
run_parallel.py
+5 -8
View File
@@ -4,27 +4,24 @@ import os
if __name__ == '__main__':
# folder = '/home/workspace/data/dataset/cloth3d/train'
folder = '/home/workspace/data/dataset/cloth3d/train'
# folder = '/home/workspace/local-data/cape_processed/train/male'
# folder = '/home/workspace/local-data/valid_skirt'
# folder = '/home/workspace/local-data/valid'
# ids = os.listdir(folder)
# ids = os.listdir(folder + '/mesh_posed_no_orient')#[280000:]
skirt = 0
obj_folder = '/data-net/dataset/cloth3d/train'
folder = '/data-local/127/train'
count = 0
step = 100
count = int(500 * 1e3)
step = 10000
# !!!!!!! change skirt parameter in script
while count <1200:
while count < 510 * 1e3:
# os.system('nohup python convert_to_uv_cape.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
# os.system('nohup python convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_unposed_onehot_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
# os.system('nohup python convert_to_uv_local.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_local_unposed_onehot_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_local_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
# os.system('nohup python data/convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_2_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv/train_2_' + str(count) + '_' + str(count+step) + '_again.out' + ' 2>&1 &')
os.system(f'nohup python data/calculate_normals.py -f {folder} -f_obj {obj_folder} -s {count} -e {(count+step)} -n normals_train_calc_{count}_{(count+step)} > nohup/normals/normals_train_calc_{count}_{(count+step)}.out' + ' 2>&1 &')
os.system('nohup python data/convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_2_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv/train_2_' + str(count) + '_' + str(count+step) + '_again.out' + ' 2>&1 &')
# os.system('nohup python convert_to_uv_local_single.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
# os.system('nohup python convert_to_uv_local_single.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_valid_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_valid_skirt_local_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
# os.system('nohup python downscale.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + ' > nohup/downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
scripts/ae_body_train.py
+15 -13
View File
@@ -11,12 +11,10 @@ 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.iccvw.ae_body_cape_m import config
else:
config_module = importlib.import_module(args['config'])
config = config_module.config
from configs.ae_body_cape_m import config
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
if args['debug']:
@@ -34,7 +32,10 @@ import plotly.graph_objects as go
import plotly.express as ex
import torch.optim as optim
from losses.losses import edge_loss
if config['is_skirt']:
from losses.losses import edge_loss_skirt as edge_loss
else:
from losses.losses import edge_loss
from data.dataset import Dataset
# from torch.utils.data import DataLoader
@@ -48,7 +49,6 @@ 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=False, is_smpl_orig=config['is_cape'])
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], 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.npy")
F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
garment_unposed, faces = mesh_pred.cpu().numpy(), F
gt = [mesh_orig[i].cpu().numpy() for i in range(pred.shape[0])]
@@ -175,7 +175,9 @@ if __name__ == '__main__':
if not args['debug']:
init_wandb()
uv_mask_body = np.load('custom_utils/mask_256x512.npy')
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 = 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'])
@@ -196,15 +198,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/iccvw/' + config['name'] + '_enc', top_n=5)
ckpt_saver_dec = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_dec', top_n=5)
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)
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=10000, log_graph=True)
wandb.watch(model, log='all', log_freq=1000, log_graph=True)
train()
scripts/ae_norm_train.py
-221
View File
@@ -1,221 +0,0 @@
import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
import argparse
import importlib
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
parser.add_argument('--config', type=str)
parser.add_argument('--debug', action='store_true')
args = vars(parser.parse_args())
if args['debug']:
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']:
os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
import time
import wandb
import random
import numpy as np
import torch
from torch import nn
# from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure as MSSSIM
from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, get_summary_str
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 data.dataset import Dataset
# from torch.utils.data import DataLoader
from custom_utils.custom_utils import plot_3d
from models.cond_encoder import AutoEncoder
from utils.checkpoints import CheckpointSaver, load_model
from custom_utils.custom_utils import Accumulator
from utils.utils import InfiniteDataLoader as DataLoader
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_norm_train.py')
artifact.add_file('losses/losses.py')
artifact.add_file('models/base.py')
artifact.add_file('models/cond_encoder.py')
artifact.add_file('models/convnext.py')
artifact.add_file('data/dataset.py')
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')]
for config_file_name in config_file_names:
file_name = '/'.join(config_file_name.split('.')) + '.py'
if os.path.exists(file_name):
artifact.add_file(file_name)
run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
def train_step(batch):
model.train()
pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.4
ms_ssim = torch.mean(_MSSSIM(pred + 1.4, batch['uv_normals'] + 1.4))
loss = 1e3 * (1 - ms_ssim)
optimizer.zero_grad()
loss.backward()
optimizer.step()
model.eval()
with torch.inference_mode():
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_normals']
l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_normals'] - pred)), axis=1))
metric = torch.mean(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1)))
return { 'train/loss': loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, }
def valid_step(batch):
model.eval()
with torch.inference_mode():
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.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)))
return { 'valid/loss': loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, }
def plot_everything(batch: dict, type: str):
model.eval()
pred = model([batch['uv_normals'], batch['uv_mask']]).detach() * 1.4
# mesh_pred = uv2mesh(pred)
# mesh_orig = batch['mesh_normals']
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.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
# gt = [mesh_orig[i].cpu().numpy() for i in range(pred.shape[0])]
# body = uv2mesh(batch['uv_body_posed']).cpu().numpy()
# garment_posed = garment_unposed
# collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
# collision_faces = [np.concatenate((F, faces + (3638 if config['is_cape'] else 14475)), axis=0) for i in range(len(garment_posed))]
# collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
# figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces, collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
figs = {}
return { **figs, **{ type + '/images': imgs } }
def train():
valid_iterator = iter(valid_dataloader)
step = 0
for batch in train_dataloader:
train_loss = train_step(batch)
train_accumulator.update(train_loss)
if step % 25 == 0:
data_valid = next(valid_iterator)
valid_loss = valid_step(data_valid)
valid_accumulator.update(valid_loss)
if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid')
to_log = { **to_log, **train_figs, **valid_figs }
if step == 0:
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
if not args['debug']:
wandb.log(to_log, step=step)
train_accumulator.reset()
valid_accumulator.reset()
print(step)
if step % 10000 == 0:
ckpt_saver_enc(model.encoder, step, optimizer, scaler=None)
ckpt_saver_dec(model.decoder, step, optimizer, scaler=None)
step += 1
class Scale(torch.nn.Module):
def __init__(self, scale_factor):
super().__init__()
self.scale_factor = scale_factor
def forward(self, x):
return x * self.scale_factor
if __name__ == '__main__':
if not args['debug']:
init_wandb()
skirt_modifier = '_skirt' if config['is_skirt'] else ''
dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'], is_cape=config['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()
valid_accumulator = Accumulator()
model : AutoEncoder = AutoEncoder(config['cond_ae'])
optimizer = optim.Adam(model.parameters(), lr=config['lr'])
# _MSSSIM = MSSSIM(data_range=2).to(config['device'])
_MSSSIM = MS_SSIM(data_range=2.8, 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')
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=20000, log_graph=True)
train()
scripts/ae_static_train.py
+10 -15
View File
@@ -10,13 +10,11 @@ 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.iccvw.ae_static_1 import config
else:
config_module = importlib.import_module(args['config'])
config = config_module.config
from configs.ae_static_cape_m import config
# from configs.ae_static_orig_2 import config
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
if args['debug']:
@@ -51,7 +49,6 @@ 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')
@@ -94,7 +91,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) ** 2, axis=-1))) / torch.sum(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'])
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, }
@@ -208,20 +205,18 @@ if __name__ == '__main__':
_MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
if config['pretrained_model'] != '':
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}')
load_model(model.encoder, config['pretrained_model'] + '_enc')
load_model(model.decoder, config['pretrained_model'] + '_dec')
ckpt_saver_enc = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_enc', top_n=5)
ckpt_saver_dec = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_dec', top_n=5)
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)
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=10000, log_graph=True)
wandb.watch(model, log='all', log_freq=1000, log_graph=True)
train()
scripts/cvae.py
+17 -30
View File
@@ -10,13 +10,10 @@ 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.iccvw.abl_all import config
else:
config_module = importlib.import_module(args['config'])
config = config_module.config
pass
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
if args['debug']:
@@ -44,7 +41,6 @@ from data.dataset import Dataset
from models.cond_encoder import AutoEncoder
from utils.checkpoints import CheckpointSaver, load_model
from utils.utils import InfiniteDataLoader as DataLoader
from utils.utils import text_to_html
from utils.postprocess import fix_collision
from losses.losses import collision_loss, normal_loss
from pytorch_msssim import MS_SSIM
@@ -54,8 +50,7 @@ from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary
def init_wandb():
wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
run = wandb.init(project='iccvw', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('scripts/cvae.py')
artifact.add_file('losses/losses.py')
@@ -88,7 +83,7 @@ def train_step(batch):
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
model.train()
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']])
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
mesh_pred = uv2mesh(pred)
if config['use_skinning']:
@@ -144,7 +139,7 @@ def valid_step(batch):
model.eval()
with torch.inference_mode():
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']])
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
mesh_pred = uv2mesh(pred)
if config['use_skinning']:
@@ -193,16 +188,15 @@ def plot_everything(batch: dict, type: str):
model.eval()
with torch.inference_mode():
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']])
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
mesh_pred = uv2mesh(pred).detach().cpu().numpy()
orig_img = torch.permute(uv_orig, [0,2,3,1])
uv_static_img = torch.permute(batch['uv_static'], [0,2,3,1])
uv_body_posed_img = torch.permute(batch['uv_body_posed'], [0,2,3,1])
uv_normals_img = (torch.permute(batch['uv_normals'], [0,2,3,1]) + 1.1) / 2.2
pred_img = torch.permute(pred, [0,2,3,1])
img = torch.concat((uv_static_img, uv_body_posed_img, uv_normals_img, orig_img, pred_img), axis=2)
img = torch.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)]
F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
@@ -259,14 +253,7 @@ def plot_everything(batch: dict, type: str):
mean_hist = wandb.Histogram(sequence=mean.cpu().numpy().mean(axis=0))
logvar_hist = wandb.Histogram(sequence=logvar.cpu().numpy().mean(axis=0))
batch_id = batch['id']
batch_frame_nr = batch['frame_nr'].cpu().numpy()
batch_garment = batch['garment']
batch_info = [f'batch_info: id: {batch_id[i]}, frame_nr: {batch_frame_nr[i]}, garment: {batch_garment[i]}' for i in range(len(batch_id))]
batch_info = text_to_html(batch_info)
batch_info = wandb.Html(batch_info)
return { **figs, **figs_coll, **{ f'{type}/batch_info': batch_info }, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
return { **figs, **figs_coll, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
def sample(folder):
@@ -275,10 +262,10 @@ def sample(folder):
count = 0
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
uv_normals = tf.one_hot(batch['uv_normals'], 7)
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
nr_of_samples=2
pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1])
@@ -294,7 +281,7 @@ def sample(folder):
mask = batch['mesh_mask'].numpy().astype(bool)[0]
_, faces = _maskFaces(np.ones((config['smpl_res'], 3)), F, mask)
pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
mesh_pred_zero = uv2mesh(pred_zero)
mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
@@ -466,11 +453,11 @@ def train():
valid_accumulator.reset()
print(step)
if train_loss['train/euclid_error'] < config['balancing_activation_limit']:
if train_loss['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
g_balancing_coeff += 1 / 10000
if step % 25000 == 0:
ckpt_saver(model, step, optimizer, None)
if step % 10000 == 0:
ckpt_saver(model, step, optimizer)
step += 1
@@ -496,7 +483,7 @@ if __name__ == '__main__':
_MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
g_balancing_coeff = config['starting_balancing_coeff']
ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=10)
ckpt_saver = CheckpointSaver('checkpoints/iccv/' + config['name'], top_n=5)
step = 0
if config['pretrained_model'] != '':
@@ -504,6 +491,6 @@ if __name__ == '__main__':
if not args['debug']:
# model = torch.compile(model)
wandb.watch(model, log='all', log_freq=10000, log_graph=True)
wandb.watch(model, log='all', log_freq=1000, log_graph=True)
train()
scripts/cvae_norm.py
-400
View File
@@ -1,400 +0,0 @@
import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
import importlib
import argparse
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
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.cvae_norm_1 import config
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
if args['debug']:
os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
import time
import argparse
import wandb
import tqdm
import random
import numpy as np
import plotly.graph_objects as go
import plotly.express as ex
from PIL import Image
if config['is_skirt']:
from losses.losses import edge_loss_skirt as edge_loss
else:
from losses.losses import edge_loss
import torch
import torch.nn as nn
import torch.optim as optim
from data.dataset import Dataset
from models.cond_encoder import AutoEncoder
from utils.checkpoints import CheckpointSaver, load_model
from utils.utils import InfiniteDataLoader as DataLoader
from utils.postprocess import fix_collision
from losses.losses import collision_loss, normal_loss
from pytorch_msssim import MS_SSIM
from models.cvae_norm_model import CVAENorm
from smpl.smpl_torch import SMPL
from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary_html, maskFacesBatch, get_summary_str, plot_3d_collision
def init_wandb():
run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('scripts/cvae_norm.py')
artifact.add_file('losses/losses.py')
artifact.add_file('models/base.py')
artifact.add_file('models/cond_encoder.py')
artifact.add_file('models/convnext.py')
artifact.add_file('models/cvae_norm_model.py')
artifact.add_file('data/dataset.py')
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')]
for config_file_name in config_file_names:
file_name = '/'.join(config_file_name.split('.')) + '.py'
if os.path.exists(file_name):
artifact.add_file(file_name)
run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
def kl_loss(mean: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
kl_loss = 1 + logvar - torch.square(mean) - torch.exp(logvar)
kl_loss = torch.mean(kl_loss)
kl_loss *= -0.5
return kl_loss
def train_step(batch):
model.train()
pred, mean, logvar = model([batch['uv_normals'], batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']])
pred = pred * 1.1
mesh_pred = uv2mesh(pred)
l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_normals'] - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_normals'] - mesh_pred) ** 2, axis=-1), axis=-1))
ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_normals'] + 1.1))
kl = kl_loss(mean, logvar)
reconstr_loss = l2_loss + l2_mesh_loss + (1 - ms_ssim)
loss = reconstr_loss + g_balancing_coeff * kl
optimizer.zero_grad()
loss.backward()
optimizer.step()
model.eval()
with torch.inference_mode():
mean = torch.mean(mean)
logvar = torch.mean(logvar)
metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_normals'] - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
return { 'train/loss': loss, 'train/euclid_error': metric,
'train/reconstr_loss': reconstr_loss,
'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
'train/l2_loss': l2_loss, 'train/l2_mesh_loss': l2_mesh_loss,
'train/ms_ssim': ms_ssim,
}
def valid_step(batch):
model.eval()
with torch.inference_mode():
pred, mean, logvar = model([batch['uv_normals'], batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']])
pred = pred * 1.1
mesh_pred = uv2mesh(pred)
l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_normals'] - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_pred - batch['mesh_normals']) ** 2, axis=-1), axis=-1))
ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_static'] + 1.1))
kl = kl_loss(mean, logvar)
reconstr_loss = l2_loss + l2_mesh_loss + (1 - ms_ssim)
loss = reconstr_loss + g_balancing_coeff * kl
mean = torch.mean(mean)
logvar = torch.mean(logvar)
metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_normals'] - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
return { 'valid/loss': loss, 'valid/euclid_error': metric,
'valid/reconstr_loss': reconstr_loss, 'valid/l2_mesh_loss': l2_mesh_loss,
'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar,
'valid/ms_ssim': ms_ssim,
}
def plot_everything(batch: dict, type: str):
model.eval()
with torch.inference_mode():
pred, mean, logvar = model([batch['uv_normals'], batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']])
pred = pred * 1.1
# mesh_pred = uv2mesh(pred).detach().cpu().numpy()
uv_static_img = torch.permute(batch['uv_static'], [0,2,3,1])
uv_body_posed_img = torch.permute(batch['uv_body_posed'], [0,2,3,1])
orig_img = (torch.permute(batch['uv_normals'], [0,2,3,1]) + 1.1) / 2.2
pred_img = torch.permute(pred, [0,2,3,1])
img = torch.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
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")
# mask = batch['mesh_mask'].cpu().numpy().astype(bool)
mean_hist = wandb.Histogram(sequence=mean.cpu().numpy().mean(axis=0))
logvar_hist = wandb.Histogram(sequence=logvar.cpu().numpy().mean(axis=0))
return { **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
def sample(folder):
# run = wandb.init(project='vae_plots', name=folder.split('/')[-1])
load_model(model,folder)
count = 0
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
uv_normals = tf.one_hot(batch['uv_normals'], 7)
nr_of_samples=2
pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1])
mesh_pred = uv2mesh(pred)
mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy()
mesh_orig = batch['mesh_posed'][0].numpy()
mesh_orig_reposed = smpl.run_garment(pose[:1], shape[:1], batch['mesh_unposed'][:1]).numpy()[0]
body = batch['mesh_body_posed'][0].numpy()
body_color = np.ones_like(body, dtype=np.int32)*100
F = np.load("custom_utils/smpl_faces.npy")
mask = batch['mesh_mask'].numpy().astype(bool)[0]
_, faces = _maskFaces(np.ones((config['smpl_res'], 3)), F, mask)
pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
mesh_pred_zero = uv2mesh(pred_zero)
mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
meshes = []
meshes.append(go.Mesh3d(
# 8 vertices of a cube
x=mesh_orig_reposed[mask,0],
y=mesh_orig_reposed[mask,2],
z=mesh_orig_reposed[mask,1],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
))
for i in range(nr_of_samples):
x = mesh_posed[i][mask, :].copy()
# x[:, 2] += 1 * (i + 1)
x[:, 2] += 1
meshes.append(go.Mesh3d(
# 8 vertices of a cube
x=x[:,0],
y=x[:,2],
z=x[:,1],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
))
for i in range(nr_of_samples):
x = mesh_posed_zero[i][mask, :].copy()
# x[:, 2] += 1 * (i + 1)
x[:, 2] += 2
meshes.append(go.Mesh3d(
# 8 vertices of a cube
x=x[:,0],
y=x[:,2],
z=x[:,1],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
))
# for i in range(5):
# x = body.copy()
# x[:, 2] += 1 * i
# meshes.append(go.Mesh3d(
# # 8 vertices of a cube
# x=x[:,0],
# y=x[:,2],
# z=x[:,1],
# # i, j and k give the vertices of triangles
# i = F[:,0],
# j = F[:,1],
# k = F[:,2],
# vertexcolor=body_color
# ))
fig = go.Figure(data=meshes)
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1))
# fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
# yaxis = dict(nticks=4, range=[-1,nr_of_samples+1],),
yaxis = dict(nticks=4, range=[-1,3],),
zaxis = dict(nticks=4, range=[-1,1],),)
)
# wandb.log({'plot_' + str(count): fig})
fig.show()
pred = pred.numpy()
frames = [Image.fromarray(((pred[i] + 1) * 128).astype(np.int8), mode='RGB') for i in range(pred.shape[0])]
frames[0].save("temp/garment_1.gif", format="GIF", append_images=frames[1:],
save_all=True, duration=500, loop=0)
if count > 10:
break
print(count)
count += 1
print()
def test():
str_split = config['name'].split('_')
test_idx = str_split.index('test')
checkpoint_name = '_'.join(str_split[:test_idx])
load_model(model, 'checkpoints/' + checkpoint_name)
dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False, do_repeat=False)
static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False)
step = 0
losses = []
for batch in tqdm.tqdm(dataset_test):
batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
test_loss = test_step(batch)
losses.append(test_loss)
to_log = { **{ 'current_step': step }, **test_loss }
if step % 1000 == 0:
figs = plot_everything(batch, 'test')
to_log = { **to_log, **figs }
wandb.log(to_log, step = step)
# print(step)
step += 1
# averages = dict()
path_to_plotly_html = "temp/plotly_figure.html"
histograms = dict()
for key in test_loss.keys():
accumulated = [losses[i][key].numpy() for i in range(len(losses))]
fig = ex.histogram(accumulated, title='Average: ' + str(np.asarray(accumulated).mean()))
fig.write_html(path_to_plotly_html, auto_play = False)
histograms['all/' + key.split('/')[-1] + '_hist'] = wandb.Html(path_to_plotly_html)
# averages['all/' + key.split('/')[-1] + '_avg'] = np.asarray(accumulated).mean()
wandb.log({ **histograms }, step = step)
wandb.finish()
pass
def train():
global g_balancing_coeff
global step
valid_iterator = iter(valid_dataloader)
for batch in train_dataloader:
train_loss = train_step(batch)
train_accumulator.update(train_loss)
if step % 25 == 0:
data_valid = next(valid_iterator)
valid_loss = valid_step(data_valid)
valid_accumulator.update(valid_loss)
if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid')
to_log = { **to_log, **train_figs, **valid_figs }
if step == 0:
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
if not args['debug']:
wandb.log(to_log, step=step)
train_accumulator.reset()
valid_accumulator.reset()
print(step)
if train_loss['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
g_balancing_coeff += 1 / 10000
if step % 10000 == 0:
ckpt_saver(model, step, optimizer, None)
step += 1
if __name__ == '__main__':
if not args['debug']:
init_wandb()
smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
skirt_modifier = '_skirt' if config['is_skirt'] else ''
dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'])
dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'])
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()
valid_accumulator = Accumulator()
model = CVAENorm(config)
optimizer = optim.Adam(model.parameters(), lr=config['lr'])
_MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
g_balancing_coeff = config['starting_balancing_coeff']
ckpt_saver = CheckpointSaver('checkpoints/iccv/' + config['name'], top_n=5)
step = 0
if config['pretrained_model'] != '':
step = load_model(model, config['pretrained_model'])
if not args['debug']:
# model = torch.compile(model)
wandb.watch(model, log='all', log_freq=10000, log_graph=True)
train()
scripts/direct_train.py
+281
View File
@@ -0,0 +1,281 @@
import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
import importlib
import argparse
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
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.direct_1 import config
# from configs.ae_static_orig_2 import config
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
if args['debug']:
os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
import time
import argparse
import wandb
import random
import numpy as np
import torch
# from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure as MSSSIM
from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, plot_3d_collision
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 data.dataset import Dataset
# from torch.utils.data import DataLoader
from losses.losses import collision_loss, normal_loss
from utils.postprocess import fix_collision
from custom_utils.custom_utils import plot_3d
from utils.checkpoints import CheckpointSaver, load_model
from custom_utils.custom_utils import Accumulator
from utils.utils import InfiniteDataLoader as DataLoader
from models.cond_encoder import AutoEncoder as Model
from smpl.smpl_torch import SMPL
from pytorch_msssim import MS_SSIM
def init_wandb():
run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('scripts/direct_train.py')
artifact.add_file('losses/losses.py')
artifact.add_file('models/base.py')
artifact.add_file('models/cond_encoder.py')
artifact.add_file('models/convnext.py')
artifact.add_file('data/dataset.py')
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')]
for config_file_name in config_file_names:
file_name = '/'.join(config_file_name.split('.')) + '.py'
if os.path.exists(file_name):
artifact.add_file(file_name)
run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
def train_step(batch):
model.train()
# with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp):
if config['use_curvature']:
input = torch.cat([batch['uv_curvature'], batch['uv_static']], axis=1)
else:
input = torch.cat([batch['uv_body_posed'], batch['uv_static']], axis=1)
pred = model([input, batch['uv_mask']])
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed']
shape = batch['pose_shape'][:, 72:]
gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
# l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
loss = l2_mesh_loss + (e_loss + c_loss + n_loss) * config['reg_loss_coeff']
optimizer.zero_grad()
loss.backward()
optimizer.step()
# scaler.scale(loss).backward()
# scaler.step(optimizer)
# scaler.update()
# optimizer.zero_grad()
model.eval()
with torch.inference_mode():
ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_unposed'] + 1))
metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/euclid_error': metric, 'train/collision_ratio': c_ratio,
'train/ms_ssim': ms_ssim, 'train/collision_loss': c_loss, 'train/l2_mesh_loss': l2_mesh_loss,
'train/normal_loss': n_loss }
def valid_step(batch):
model.eval()
with torch.inference_mode():
if config['use_curvature']:
input = torch.cat([batch['uv_curvature'], batch['uv_static']], axis=1)
else:
input = torch.cat([batch['uv_body_posed'], batch['uv_static']], axis=1)
pred = model([input, batch['uv_mask']])
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed']
shape = batch['pose_shape'][:, 72:]
gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
# l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_unposed'] + 1))
loss = l2_mesh_loss + (e_loss + c_loss + n_loss) * config['reg_loss_coeff']
metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/euclid_error': metric, 'valid/collision_ratio': c_ratio,
'valid/ms_ssim': ms_ssim, 'valid/collision_loss': c_loss, 'valid/l2_mesh_loss': l2_mesh_loss,
'valid/normal_loss': n_loss }
def plot_everything(batch: dict, type: str):
model.eval()
with torch.inference_mode():
if config['use_curvature']:
input = torch.cat([batch['uv_curvature'], batch['uv_static']], axis=1)
else:
input = torch.cat([batch['uv_body_posed'], batch['uv_static']], axis=1)
pred = model([input, batch['uv_mask']]).detach()
mesh_pred = uv2mesh(pred).cpu().numpy()
mesh_orig = batch['mesh_unposed'].cpu().numpy()
orig_img = torch.permute(batch['uv_unposed'], [0,2,3,1])
uv_static_img = torch.permute(batch['uv_static'], [0,2,3,1])
uv_body_posed_img = torch.permute(batch['uv_body_posed'], [0,2,3,1])
pred_img = torch.permute(pred, [0,2,3,1])
img = torch.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
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")
mask = batch['mesh_mask'].cpu().numpy().astype(bool)
shape = batch['pose_shape'][:, 72:]
gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])]).cpu().numpy()
garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
garment_unposed = [garment_unposed[i] + shape_offsets[i, mask[i], :] for i in range(mask.shape[0])]
gt = [mesh_orig[i][mask[i], :] + shape_offsets[i, mask[i], :] for i in range(pred.shape[0])]
for i in range(mesh_orig.shape[0]):
mesh_pred[i, mask[i], :] += shape_offsets[i, mask[i], :]
mesh_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], torch.tensor(mesh_pred, device=config['device']))[0].cpu().numpy()
gt_posed = [batch['mesh_posed'].cpu().numpy()[i, mask[i], :] for i in range(mask.shape[0])]
mesh_posed_fixed = fix_collision(torch.tensor(mesh_posed, device=config['device']), batch['mesh_body_posed'], torch.tensor(mask, device=config['device'])).cpu().numpy()
body = batch['mesh_body_posed'].cpu().numpy()
garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
collision_faces = [np.concatenate((F, faces[i] + config['smpl_res']), axis=0) for i in range(len(garment_posed))]
collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
garment_posed_fixed = [mesh_posed_fixed[i, mask[i], :] for i in range(body.shape[0])]
collision_fixed = [np.concatenate((body[i], garment_posed_fixed[i]), axis=0) for i in range(len(garment_posed))]
collision_gt = [np.concatenate((body[i], gt_posed[i]), axis=0) for i in range(len(garment_posed))]
figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i]) for i in range(4) }
return { **figs, **figs_coll, **{ type + '/images': imgs, } }
def train():
valid_iterator = iter(valid_dataloader)
step = 0
for batch in train_dataloader:
train_loss = train_step(batch)
train_accumulator.update(train_loss)
if step % 25 == 0:
data_valid = next(valid_iterator)
valid_loss = valid_step(data_valid)
valid_accumulator.update(valid_loss)
if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid')
to_log = { **to_log, **train_figs, **valid_figs }
if step == 0:
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
if not args['debug']:
wandb.log(to_log, step=step)
train_accumulator.reset()
valid_accumulator.reset()
print(step)
if step % 10000 == 0:
ckpt_saver_enc(model.encoder, step, optimizer, scaler)
ckpt_saver_dec(model.decoder, step, optimizer, scaler)
step += 1
if __name__ == '__main__':
if not args['debug']:
init_wandb()
use_amp = True
skirt_modifier = '_skirt' if config['is_skirt'] else ''
# dataset = Dataset('/data-local/127/train', config['signature'], resolutions=[512], device=device, frames_file='frames_jumpsuit.txt')
dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'], is_cape=config['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()
valid_accumulator = Accumulator()
model : Model = Model(config)
optimizer = optim.Adam(model.parameters(), lr=config['lr'])
scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
# _MSSSIM = MSSSIM(data_range=2).to(config['device'])
_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')
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)
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)
train()
scripts/upscaler_train.py
+15 -25
View File
@@ -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"] = '2'
os.environ["CUDA_VISIBLE_DEVICES"] = '6'
import time
@@ -49,36 +49,26 @@ 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 is_skirt:
train_frames = 'frames_train_skirt.txt'
valid_frames = 'frames_valid_skirt.txt'
if not is_skirt:
train_dir = '/data-local/127/train'
valid_dir = '/data-local/127/valid'
else:
train_frames = 'frames_train.txt'
valid_frames = 'frames_valid.txt'
if is_cape:
train_frames = None
valid_frames = None
train_dir = '/data-local/128/train_skirt'
valid_dir = '/data-local/127/valid_skirt'
if is_cape:
if is_male:
train_dir = '/data-local/127_2/cape_processed/train/male'
valid_dir = '/data-local/127_2/cape_processed/test/male'
train_dir = '/data-local/127/cape_processed/train/male'
valid_dir = '/data-local/127/cape_processed/test/male'
else:
train_dir = '/data-local/127_2/cape_processed/train/female'
valid_dir = '/data-local/127_2/cape_processed/test/female'
train_dir = '/data-local/127/cape_processed/train/female'
valid_dir = '/data-local/127/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():
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'])
run = wandb.init(project='upscale', name=name, notes='', dir='/data-local/127/temp', tags=['upscaler', 'iccv'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('scripts/upscaler_train.py')
artifact.add_file('losses/losses.py')
@@ -265,8 +255,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=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)
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)
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()
@@ -289,13 +279,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/iccvw/' + name, top_n=5)
ckpt_saver = CheckpointSaver('checkpoints/iccv/torch/' + name, top_n=5)
print(f'started training: {name}')
if not args['debug']:
# model = torch.compile(model)
wandb.watch(model, log='all', log_freq=10000, log_graph=True)
wandb.watch(model, log='all', log_freq=1000, log_graph=True)
train()
smpl/smpl_torch.py
+3 -7
View File
@@ -163,9 +163,6 @@ class SMPL(Module):
v_shaped = torch.stack(v_shaped, axis=-1)
# J = torch.matmul(self.J_regressor, v_shaped)
if verts is None:
J = torch.stack([torch.matmul(self.J_regressor[g.to(torch.long)], v_shaped[i, :, :, g.to(torch.long)]) for i, g in enumerate(betas[:,-1])])
else:
J = torch.stack([torch.matmul(self.J_regressor[torch.greater(g, 0.5).to(torch.long)], v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)]) for i, g in enumerate(betas[:,-1])])
J0 = torch.reshape(J[:,0,:], (-1, 1, 3))
J = J - J0
@@ -209,10 +206,9 @@ class SMPL(Module):
)
)
joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.device)], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
weights = self.weights if verts is not None else self.weights_all
T = torch.stack([torch.tensordot(results[i], weights[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])])
T = torch.stack([torch.tensordot(results[i], self.weights_all[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])])
rest_shape_h = torch.cat(
(v_posed, torch.ones((batch_num, v_posed.shape[1], 1), dtype=torch.float32).to(self.device)), dim=2
)
@@ -234,4 +230,4 @@ class SMPL(Module):
if remove and verts is None:
result = result[:,self.to_keep,:]
return result, joints, v_shaped, v_posed
return result, joints
utils/checkpoints.py
-1
View File
@@ -39,6 +39,5 @@ def load_model(model: torch.nn.Module, checkpoint, optimizer: torch.nn.Module=No
model.load_state_dict(checkpoint['model'])
if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer'].state_dict())
if scaler is not None and checkpoint['scaler'] is not None: scaler.load_state_dict(checkpoint['scaler'])
print(f'Loaded model from {checkpoints[-1]} at iteration {iteration}')
return int(iteration) if iteration != '' else 0
utils/generate_figures.py
-126
View File
@@ -1,126 +0,0 @@
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
utils/postprocess.py
+1 -1
View File
@@ -22,7 +22,7 @@ from losses.losses import _nearest_neighbour
""" COLLISION FUNCTIONS START """
# F = np.load("smpl/smpl_faces.npy")
to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
F_fix_coll = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[None,:,:], [4,1,1]), device='cuda')
F_fix_coll = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
smpl_to_keep_orig_post = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
smpl_to_keep_orig_post = smpl_to_keep_orig_post[smpl_to_keep_orig_post < 6890]
F_fix_coll_orig = torch.tensor(np.tile(np.load("smpl/smpl_faces_orig.npy")[:6890,:][None,:,:], [4,1,1]), device='cuda')[:,smpl_to_keep_orig_post,:]
utils/utils.py
-6
View File
@@ -43,9 +43,3 @@ class InfiniteDataLoader(DataLoader):
self.dataset_iterator = super().__iter__()
batch = next(self.dataset_iterator)
return batch
def text_to_html(lines):
html = ''
for line in lines:
html += "<pre>" + line + "</pre>"
return html