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

Compare commits

base: luni/cvpr_2023:dbbcd5f14ec3d622d60e77640c2c67928f48dff8
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:iccv_review
Branches Tags
luni/cvpr_2023:main luni/cvpr_2023:iccv_review luni/cvpr_2023:vto luni/cvpr_2023:extrapolation luni/cvpr_2023:skip_conn
8 Commits
dbbcd5f14e .. iccv_review
Author SHA1 Message Date
luni 32c3d1103f train with fixed reg losses 2023-06-02 09:36:01 +00:00
luni 263754d755 fix postprocessing 2023-06-02 09:18:01 +00:00
luni 4f869c864c training with remasked sequence 2023-06-01 21:08:27 +00:00
luni e0825722e3 collision and normal loss fix!!! 2023-06-01 21:08:13 +00:00
luni e23d097398 direct training experiments 2023-06-01 12:55:42 +00:00
luni f3c1e3308e direct training with normals 2023-05-31 21:40:18 +00:00
luni 4e759fcccc Merge branch 'main' into iccv_review 2023-05-31 15:08:59 +02:00
luni 53217c6e5d direct training 2023-05-23 22:57:08 +02:00
31 changed files with 1197 additions and 3702 deletions
Expand all files Collapse all files
configs/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/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/direct_3.py
+82
View File
@@ -0,0 +1,82 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '0'
config['name'] = 'direct_3'
config['notes'] = 'with normals input'
config['tags'] = ['torch', 'direct']
config['script'] = 'scripts/direct_train.py'
config['pretrained_model'] = ''
config['use_curvature'] = False
config['use_normals'] = True
# assert that only use curvature or normals
assert not (config['use_curvature'] and config['use_normals'])
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'] = 9
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_normals',
'mesh_mask',
'mesh_body_posed',
'mesh_posed',
'uv_unposed',
'mesh_unposed',
'uv_static',
'uv_mask',
'uv_body_posed',
}
configs/direct_4.py
+82
View File
@@ -0,0 +1,82 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '5'
config['name'] = 'direct_5'
config['notes'] = 'with normals input but only single garment sequence'
config['tags'] = ['torch', 'direct']
config['script'] = 'scripts/direct_train.py'
config['pretrained_model'] = ''
config['use_curvature'] = False
config['use_normals'] = True
# assert that only use curvature or normals
assert not (config['use_curvature'] and config['use_normals'])
config['batch_size'] = 4
config['train_dir'] = '/data-local/127/train'
config['valid_dir'] = '/data-local/127/train'
config['test_dir'] = '/data-net/dataset/test_ready'
config['train_frames'] = 'frames_00761.txt'
config['valid_frames'] = 'frames_00761.txt'
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = True
config['is_cape'] = False
config['remove_global_orientation'] = True
config['device'] = 'cuda'
config['input_features'] = 9
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_normals',
'mesh_mask',
'mesh_body_posed',
'mesh_posed',
'uv_unposed',
'mesh_unposed',
'uv_static',
'uv_mask',
'uv_body_posed',
}
configs/direct_6.py
+83
View File
@@ -0,0 +1,83 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '4'
config['name'] = 'direct_6'
config['notes'] = 'with normals input but only single garment sequence with msssim and l2 mesh loss only'
config['tags'] = ['torch', 'direct']
config['script'] = 'scripts/direct_train.py'
config['pretrained_model'] = ''
config['use_curvature'] = False
config['use_normals'] = True
config['use_skinning'] = True
# assert that only use curvature or normals
assert not (config['use_curvature'] and config['use_normals'])
config['batch_size'] = 4
config['train_dir'] = '/data-local/127/train'
config['valid_dir'] = '/data-local/127/train'
config['test_dir'] = '/data-net/dataset/test_ready'
config['train_frames'] = 'frames_00761.txt'
config['valid_frames'] = 'frames_00761.txt'
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = True
config['is_cape'] = False
config['remove_global_orientation'] = True
config['device'] = 'cuda'
config['input_features'] = 9
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_normals',
'mesh_mask',
'mesh_body_posed',
'mesh_posed',
'uv_unposed',
'mesh_unposed',
'uv_static',
'uv_mask',
'uv_body_posed',
}
configs/direct_7.py
+87
View File
@@ -0,0 +1,87 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '0'
config['name'] = 'direct_9'
config['notes'] = 'with normals input whole dataset with msssim and l2 mesh loss and reg losses and without skinning'
config['tags'] = ['torch', 'direct']
config['script'] = 'scripts/direct_train.py'
config['pretrained_model'] = ''
config['use_curvature'] = False
config['use_normals'] = True
config['use_skinning'] = False
# assert that only use curvature or normals
assert not (config['use_curvature'] and config['use_normals'])
config['batch_size'] = 4
config['train_dir'] = '/data-local/127/train'
config['valid_dir'] = '/data-local/127/train'
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'] = 9
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_normals',
'mesh_mask',
'mesh_body_posed',
'mesh_posed',
'uv_static',
'uv_mask',
'uv_body_posed',
}
if config['use_skinning']:
config['signature'].add('uv_unposed')
config['signature'].add('mesh_unposed')
else:
config['signature'].add('uv_posed')
configs/direct_8.py
+87
View File
@@ -0,0 +1,87 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '6'
config['name'] = 'direct_8'
config['notes'] = 'with normals input but only single garment sequence with msssim and l2 mesh loss only and without skinning and remasked'
config['tags'] = ['torch', 'direct']
config['script'] = 'scripts/direct_train.py'
config['pretrained_model'] = ''
config['use_curvature'] = False
config['use_normals'] = True
config['use_skinning'] = False
# assert that only use curvature or normals
assert not (config['use_curvature'] and config['use_normals'])
config['batch_size'] = 4
config['train_dir'] = '/data-net/dataset/00761'
config['valid_dir'] = '/data-net/dataset/00761'
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'] = True
config['is_cape'] = False
config['remove_global_orientation'] = True
config['device'] = 'cuda'
config['input_features'] = 9
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_normals',
'mesh_mask',
'mesh_body_posed',
'mesh_posed',
'uv_static',
'uv_mask',
'uv_body_posed',
}
if config['use_skinning']:
config['signature'].add('uv_unposed')
config['signature'].add('mesh_unposed')
else:
config['signature'].add('uv_posed')
custom_data/zero_body.npy
BIN
View File
Binary file not shown.
custom_utils/custom_utils.py
+1
View File
@@ -43,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]
data/calculate_normals.py
+2 -2
View File
@@ -62,8 +62,8 @@ def get_frames(folder):
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("-f", "--folder", help = "end of range", default='')
parser.add_argument("-f_obj", "--folder_object", help = "end of range", default='/')
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')
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
+148 -125
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_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.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 = 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')):
# 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')
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
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/dataset.py
+71 -131
View File
@@ -21,7 +21,6 @@ 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
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) -> dict:
if frame 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] = stored_data
self.store[id] = stored_data
data = { key: value.to(self.device) for key, value in self.store[frame].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 Dataset(Dataset):
def __init__(self, folder, signature, resolutions=[512], frames_file='', device='cuda', use_cache=False, is_cape=False):
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 Dataset(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)
@@ -167,56 +145,73 @@ class Dataset(Dataset):
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)
item['id'] = frame.id
item['frame_nr'] = frame.frame_nr
item['garment'] = frame.garment
if 'id' in self.signature:
item['id'] = frame.id
if 'frame_nr' in self.signature:
item['frame_nr'] = frame.frame_nr
if 'pose_shape' in self.signature:
pose = 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
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 = 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_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_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
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 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,93 +221,38 @@ class Dataset(Dataset):
self.cache[frame] = data
data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in data.items()}
data = { **data, **self.static_cache[frame] }
data = { **data, **self.static_cache[frame.id] }
return data
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
@@ -341,7 +281,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/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 -453
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")
# parser.add_argument("-n", "--name", help = "end of range")
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
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,443 +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 = 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']
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
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 = 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 = 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 = 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 = 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')
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'])
_, 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 verify_registration_ans_mask(folder, out_folder):
ids = 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')
if __name__ == '__main__':
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)
# create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder)
# verify_registration_ans_mask(folder, out_folder)
# create_pose_shape(folder, out_folder)
# create_mesh_static(folder, out_folder)
# create_mesh_body_posed_no_orient(folder, out_folder)
# create_mesh_normals(folder, out_folder)
# create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end)
create_frames_file(folder, out_folder)
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
+1 -1
View File
@@ -1,7 +1,7 @@
import tables as tb
class Frame:
def __init__(self, id: str, frame_nr: int, garment: str=''):
def __init__(self, id, frame_nr, garment=''):
self.id = id
self.frame_nr = frame_nr
self.garment = garment
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
+13 -13
View File
@@ -27,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'])
@@ -48,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:
@@ -98,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)
@@ -108,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_6 import config
import argparse
# parser = argparse.ArgumentParser(description='')
scripts/cvae.py
+11 -13
View File
@@ -13,8 +13,7 @@ args = vars(parser.parse_args())
config_module = importlib.import_module(args['config'])
config = config_module.config
if args['debug']:
from configs.abl_norm import config
pass
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
if args['debug']:
@@ -51,7 +50,7 @@ from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary
def init_wandb():
run = wandb.init(project='iccv', 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')
@@ -84,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']:
@@ -140,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']:
@@ -189,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")
@@ -264,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])
@@ -283,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()
@@ -459,7 +457,7 @@ def train():
g_balancing_coeff += 1 / 10000
if step % 10000 == 0:
ckpt_saver(model, step, optimizer, None)
ckpt_saver(model, step, optimizer)
step += 1
@@ -493,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
+315
View File
@@ -0,0 +1,315 @@
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_7 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_2/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)
elif config['use_normals']:
input = torch.cat([batch['uv_normals'], batch['uv_body_posed'], 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)
if config['use_skinning']:
mesh_orig = batch['mesh_unposed']
uv_orig = batch['uv_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]
else:
mesh_orig = batch['mesh_posed']
uv_orig = batch['uv_posed']
mesh_pred_posed = mesh_pred
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))
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]))
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, uv_orig + 1))
loss = l2_mesh_loss + (1 - ms_ssim) * 1e3 + (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():
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)
elif config['use_normals']:
input = torch.cat([batch['uv_normals'], batch['uv_body_posed'], 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)
if config['use_skinning']:
mesh_orig = batch['mesh_unposed']
uv_orig = batch['uv_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]
else:
mesh_orig = batch['mesh_posed']
uv_orig = batch['uv_posed']
mesh_pred_posed = mesh_pred
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, uv_orig + 1))
loss = l2_mesh_loss + (1 - ms_ssim) * 1e3 + (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)
elif config['use_normals']:
input = torch.cat([batch['uv_normals'], batch['uv_body_posed'], 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()
if config['use_skinning']:
mesh_orig = batch['mesh_unposed'].cpu().numpy()
uv_orig = batch['uv_unposed']
else:
mesh_orig = batch['mesh_posed'].cpu().numpy()
uv_orig = batch['uv_posed']
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)
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)
if config['use_skinning']:
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()
else:
garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
mesh_posed = mesh_pred
gt = [mesh_orig[i][mask[i], :] for i in range(pred.shape[0])]
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_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))]
# faces_ = np.load('/data-net/dataset/00761/faces/00761.npy')
# faces = [faces_ 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))]
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=10000, log_graph=True)
train()
smpl/smpl_torch.py
+4 -8
View File
@@ -163,12 +163,9 @@ 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
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
R_cube_big = self.rodrigues(pose.reshape(-1, 1, 3)).reshape(batch_num, -1, 3, 3)
if verts is not None:
@@ -211,8 +208,7 @@ 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]
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[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
)