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

save all

Browse Source
This commit is contained in:
luni
2023-08-24 11:36:22 +00:00
parent 59f67a2495
commit afd482522d
36 changed files with 2194 additions and 291 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.vscode/launch.json
Vendored
+1 -1
View File
@@ -12,7 +12,7 @@
"console": "integratedTerminal",
"justMyCode": true, // true false
"cwd": "${workspaceFolder}",
"args": ["--config", "configs.pyr_32", "--debug"]
"args": ["--debug"]
}
]
}
configs/iccvw/pyr_64_128_256_512_e2e_var.py
+302
View File
@@ -0,0 +1,302 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '6'
config['name'] = 'pyr_iccvw_64_128_256_512_e2e_var_2'
config['notes'] = ''
config['tags'] = ['pyr', 'iccv', 'torch']
config['script'] = 'scripts/cvae_pyramid_e2e.py'
config['batch_size'] = 4
config['is_skirt'] = False
config['is_cape'] = False
config['device'] = 'cuda'
config['use_skinning'] = True # and deshaping
config['use_extra_losses'] = True # edge, collision, 3d mesh
config['use_norm_loss'] = True
config['use_normals'] = True
config['pretrained_model'] = 'checkpoints/iccvw/pyr_iccvw_64_128_256_512_e2e_2'
config['pretrained_model_iteration'] = '250000'
config['pyramid'] = dict()
config['pyramid']['resolutions'] = [64, 128, 256, 512]
config['pyramid']['pretrained_pyramid'] = ''
config['pyramid']['pretrained_pyramid_iteration'] = ''
config['pyramid']['final_upscale_enhancer'] = ''
config['train_dir'] = '/data-local/127/cloth3d_processed'
config['valid_dir'] = '/data-local/127/cloth3d_processed'
config['test_dir'] = '/data-local/127_2/cloth3d_processed_test'
if config['is_skirt']:
config['train_frames'] = 'frames_train_skirt.txt'
config['valid_frames'] = 'frames_valid_skirt.txt'
config['test_frames'] = 'frames_test_skirt.txt'
else:
config['train_frames'] = 'frames_train.txt'
config['valid_frames'] = 'frames_valid.txt'
config['test_frames'] = 'frames_test.txt'
config['dataset_in_memory'] = False
config['resolution'] = None
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['lr'] = 0.0001
config['starting_balancing_coeff'] = 1
config['balancing_activation_limit'] = 0.01
config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
config['avg_garment_path'] = ''
config['bottleneck_scale'] = 2
config['signature']: set = { 'id',
'frame_nr',
'pose_shape',
'uv_body_posed',
'mesh_body_posed',
'mesh_posed',
'mesh_unposed',
'uv_unposed',
'uv_normals',
'uv_static',
'mesh_mask',
'uv_mask',
}
config['pyramid']['is_label'] = False
# config['pyramid']['add_noise'] = False
# config['pyramid']['noise_amount'] = 0.0
config['pyramid']['loss_weights'] = [1,1,1,1,1]
for res in config['pyramid']['resolutions']:
if res != 512:
res_modifier = '_' + str(res)
config['signature'].add('uv_body_posed' + res_modifier)
config['signature'].add('uv_unposed' + res_modifier)
config['signature'].add('uv_normals' + res_modifier)
config['signature'].add('uv_static' + res_modifier)
config['signature'].add('uv_mask' + res_modifier)
config['static_model'] = 'checkpoints/iccvw/static_iccvw_4'
config['body_model'] = 'checkpoints/iccvw/body_iccvw_1'
config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1'
config['freeze_encoders'] = True
config['pyramid']['levels'] = []
for res in config['pyramid']['resolutions']:
if res == 512:
level = dict()
level['resolution'] = 512
level['bottleneck_size'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 7
level['upscale_enhancer'] = ''
level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 3),
CNBlockConfig(128, level['bottleneck_size'], 9),
CNBlockConfig(level['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(level['bottleneck_size'] * 2, 128, 3),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 3),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
]
level['encoder_block_settings'] = block_setting
level['decoder_block_settings'] = block_setting_reverse
config['pyramid']['levels'].append(level)
elif res == 256:
level = dict()
level['resolution'] = 256
level['kernel_size'] = 5
level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_256'
level['upscale_enhancer_freeze'] = True
level['bottleneck_size'] = 64 * config['bottleneck_scale']
level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
block_setting = [
CNBlockConfig(64, 128, 3),
CNBlockConfig(128, 128, 3),
CNBlockConfig(128, level['bottleneck_size'], 9),
CNBlockConfig(level['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(level['bottleneck_size'] * 2, 128, 3),
CNBlockConfig(128, 128, 9),
CNBlockConfig(128, 64, 3),
CNBlockConfig(64, 3, 3),
]
level['encoder_block_settings'] = block_setting
level['decoder_block_settings'] = block_setting_reverse
config['pyramid']['levels'].append(level)
elif res == 128:
level = dict()
level['resolution'] = 128
level['kernel_size'] = 5
level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_128'
level['upscale_enhancer_freeze'] = True
level['bottleneck_size'] = 64 * config['bottleneck_scale']
level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
block_setting = [
CNBlockConfig(64, 128, 9),
CNBlockConfig(128, level['bottleneck_size'], 9),
CNBlockConfig(level['bottleneck_size'], None, 9),
]
block_setting_reverse = [
CNBlockConfig(level['bottleneck_size'] * 2, 128, 9),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 3, 9),
]
level['encoder_block_settings'] = block_setting
level['decoder_block_settings'] = block_setting_reverse
config['pyramid']['levels'].append(level)
elif res == 64:
level = dict()
level['resolution'] = 64
level['kernel_size'] = 3
level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_64'
level['upscale_enhancer_freeze'] = True
level['bottleneck_size'] = 64 * config['bottleneck_scale']
level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
block_setting = [
CNBlockConfig(64, level['bottleneck_size'], 9),
CNBlockConfig(level['bottleneck_size'], None, 9),
]
block_setting_reverse = [
CNBlockConfig(level['bottleneck_size'] * 2, 64, 9),
CNBlockConfig(64, 3, 9),
]
level['encoder_block_settings'] = block_setting
level['decoder_block_settings'] = block_setting_reverse
config['pyramid']['levels'].append(level)
elif res == 32:
level = dict()
level['resolution'] = 32
level['kernel_size'] = 3
level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_32'
level['upscale_enhancer_freeze'] = True
level['bottleneck_size'] = 64 * config['bottleneck_scale']
level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
block_setting = [
CNBlockConfig(level['bottleneck_size'], None, 16),
]
block_setting_reverse = [
CNBlockConfig(level['bottleneck_size'] * 2, 3, 16),
]
level['encoder_block_settings'] = block_setting
level['decoder_block_settings'] = block_setting_reverse
config['pyramid']['levels'].append(level)
else:
raise NotImplementedError
####################### cond_ae_norm ############################
config['cond_ae_norm'] = dict()
config['cond_ae_norm']['device'] = config['device']
config['cond_ae_norm']['is_label'] = False
# config['cond_ae_norm']['latent_dim'] = config['latent_dim']
config['cond_ae_norm']['kernel_size'] = 7
config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 9),
CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
]
config['cond_ae_norm']['encoder_block_settings'] = block_setting
config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_norm']['activation'] = 'tanh'
####################### cond_ae_body ############################
config['cond_ae_body'] = dict()
config['cond_ae_body']['device'] = config['device']
config['cond_ae_body']['is_label'] = False
# config['cond_ae_body']['latent_dim'] = config['latent_dim']
config['cond_ae_body']['kernel_size'] = 7
config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 9),
CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
]
config['cond_ae_body']['encoder_block_settings'] = block_setting
config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_body']['activation'] = 'tanh'
####################### cond_ae_static ############################
config['cond_ae_static'] = dict()
config['cond_ae_static']['device'] = config['device']
config['cond_ae_static']['is_label'] = False
# config['cond_ae_static']['latent_dim'] = config['latent_dim']
config['cond_ae_static']['kernel_size'] = 7
config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(8, 16, 3),
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 8, 3),
CNBlockConfig(8, 3, 3),
]
config['cond_ae_static']['encoder_block_settings'] = block_setting
config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_static']['activation'] = 'tanh'
configs/iccvw/pyr_64_128_256_512_skirt_e2e_var.py
+302
View File
@@ -0,0 +1,302 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '0'
config['name'] = 'pyr_iccvw_64_128_256_512_skirt_e2e_var_3'
config['notes'] = ''
config['tags'] = ['pyr', 'iccv', 'torch']
config['script'] = 'scripts/cvae_pyramid_e2e.py'
config['batch_size'] = 4
config['is_skirt'] = True
config['is_cape'] = False
config['device'] = 'cuda'
config['use_skinning'] = True # and deshaping
config['use_extra_losses'] = True # edge, collision, 3d mesh
config['use_norm_loss'] = True
config['use_normals'] = True
config['pretrained_model'] = 'checkpoints/iccvw/pyr_iccvw_64_128_256_512_skirt_e2e'
config['pretrained_model_iteration'] = '200000'
config['pyramid'] = dict()
config['pyramid']['resolutions'] = [64, 128, 256, 512]
config['pyramid']['pretrained_pyramid'] = ''
config['pyramid']['pretrained_pyramid_iteration'] = ''
config['pyramid']['final_upscale_enhancer'] = ''
config['train_dir'] = '/data-local/127/cloth3d_processed'
config['valid_dir'] = '/data-local/127/cloth3d_processed'
config['test_dir'] = '/data-local/127_2/cloth3d_processed_test'
if config['is_skirt']:
config['train_frames'] = 'frames_train_skirt.txt'
config['valid_frames'] = 'frames_valid_skirt.txt'
config['test_frames'] = 'frames_test_skirt.txt'
else:
config['train_frames'] = 'frames_train.txt'
config['valid_frames'] = 'frames_valid.txt'
config['test_frames'] = 'frames_test.txt'
config['dataset_in_memory'] = False
config['resolution'] = None
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['lr'] = 0.0001
config['starting_balancing_coeff'] = 1
config['balancing_activation_limit'] = 0.01
config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
config['avg_garment_path'] = ''
config['bottleneck_scale'] = 2
config['signature']: set = { 'id',
'frame_nr',
'pose_shape',
'uv_body_posed',
'mesh_body_posed',
'mesh_posed',
'mesh_unposed',
'uv_unposed',
'uv_normals',
'uv_static',
'mesh_mask',
'uv_mask',
}
config['pyramid']['is_label'] = False
# config['pyramid']['add_noise'] = False
# config['pyramid']['noise_amount'] = 0.0
config['pyramid']['loss_weights'] = [1,1,1,1,1]
for res in config['pyramid']['resolutions']:
if res != 512:
res_modifier = '_' + str(res)
config['signature'].add('uv_body_posed' + res_modifier)
config['signature'].add('uv_unposed' + res_modifier)
config['signature'].add('uv_normals' + res_modifier)
config['signature'].add('uv_static' + res_modifier)
config['signature'].add('uv_mask' + res_modifier)
config['static_model'] = 'checkpoints/iccvw/static_iccvw_1_skirt'
config['body_model'] = 'checkpoints/iccvw/body_iccvw_1_skirt'
config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1_skirt'
config['freeze_encoders'] = True
config['pyramid']['levels'] = []
for res in config['pyramid']['resolutions']:
if res == 512:
level = dict()
level['resolution'] = 512
level['bottleneck_size'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 7
level['upscale_enhancer'] = ''
level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 3),
CNBlockConfig(128, level['bottleneck_size'], 9),
CNBlockConfig(level['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(level['bottleneck_size'] * 2, 128, 3),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 3),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
]
level['encoder_block_settings'] = block_setting
level['decoder_block_settings'] = block_setting_reverse
config['pyramid']['levels'].append(level)
elif res == 256:
level = dict()
level['resolution'] = 256
level['kernel_size'] = 5
level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_256_skirt'
level['upscale_enhancer_freeze'] = True
level['bottleneck_size'] = 64 * config['bottleneck_scale']
level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
block_setting = [
CNBlockConfig(64, 128, 3),
CNBlockConfig(128, 128, 3),
CNBlockConfig(128, level['bottleneck_size'], 9),
CNBlockConfig(level['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(level['bottleneck_size'] * 2, 128, 3),
CNBlockConfig(128, 128, 9),
CNBlockConfig(128, 64, 3),
CNBlockConfig(64, 3, 3),
]
level['encoder_block_settings'] = block_setting
level['decoder_block_settings'] = block_setting_reverse
config['pyramid']['levels'].append(level)
elif res == 128:
level = dict()
level['resolution'] = 128
level['kernel_size'] = 5
level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_128_skirt'
level['upscale_enhancer_freeze'] = True
level['bottleneck_size'] = 64 * config['bottleneck_scale']
level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
block_setting = [
CNBlockConfig(64, 128, 9),
CNBlockConfig(128, level['bottleneck_size'], 9),
CNBlockConfig(level['bottleneck_size'], None, 9),
]
block_setting_reverse = [
CNBlockConfig(level['bottleneck_size'] * 2, 128, 9),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 3, 9),
]
level['encoder_block_settings'] = block_setting
level['decoder_block_settings'] = block_setting_reverse
config['pyramid']['levels'].append(level)
elif res == 64:
level = dict()
level['resolution'] = 64
level['kernel_size'] = 3
level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_64_skirt'
level['upscale_enhancer_freeze'] = True
level['bottleneck_size'] = 64 * config['bottleneck_scale']
level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
block_setting = [
CNBlockConfig(64, level['bottleneck_size'], 9),
CNBlockConfig(level['bottleneck_size'], None, 9),
]
block_setting_reverse = [
CNBlockConfig(level['bottleneck_size'] * 2, 64, 9),
CNBlockConfig(64, 3, 9),
]
level['encoder_block_settings'] = block_setting
level['decoder_block_settings'] = block_setting_reverse
config['pyramid']['levels'].append(level)
elif res == 32:
level = dict()
level['resolution'] = 32
level['kernel_size'] = 3
level['upscale_enhancer'] = 'checkpoints/iccvw/iccv_32_skirt'
level['upscale_enhancer_freeze'] = True
level['bottleneck_size'] = 64 * config['bottleneck_scale']
level['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
block_setting = [
CNBlockConfig(level['bottleneck_size'], None, 16),
]
block_setting_reverse = [
CNBlockConfig(level['bottleneck_size'] * 2, 3, 16),
]
level['encoder_block_settings'] = block_setting
level['decoder_block_settings'] = block_setting_reverse
config['pyramid']['levels'].append(level)
else:
raise NotImplementedError
####################### cond_ae_norm ############################
config['cond_ae_norm'] = dict()
config['cond_ae_norm']['device'] = config['device']
config['cond_ae_norm']['is_label'] = False
# 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'
data/convert_to_uv.py
+5 -1
View File
@@ -399,10 +399,11 @@ 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("-f", "--folder", help = "end of range", default='/data-local/127_2/cloth3d_processed_test')
parser.add_argument("-s", "--start", help = "start of range", default=0, type=int)
parser.add_argument("-e", "--end", help = "end of range", default=100, type=int)
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
parser.add_argument('--debug', action='store_true')
args = parser.parse_args()
F = {'non-skirt': np.load("smpl/smpl_faces.npy"), 'skirt': np.load("smpl/smpl_faces_skirt.npy")}
@@ -442,6 +443,7 @@ if __name__ == '__main__':
#########################################################################################
folder = args.folder
# data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_normals', 'body_posed'] # 'mesh_unposed',
data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_normals', 'body_posed', 'mesh_unposed'] # 'mesh_unposed',
ids = sorted(os.listdir(folder + '/mesh_body_posed_no_orient'))[args.start:args.end]
@@ -460,6 +462,8 @@ if __name__ == '__main__':
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)
# ids = ['00438']
for id in tqdm.tqdm(ids):
garments = glob.glob(os.path.join(folder, 'mesh_mask', f'{id}_*'))
data/convert_to_uv_cape.py
+33 -34
View File
@@ -401,7 +401,7 @@ def get_shapedirs():
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/test/female')
parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/train/female')
parser.add_argument("-s", "--start", help = "start of range", default=0, type=int)
parser.add_argument("-e", "--end", help = "end of range", default=100, type=int)
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
@@ -448,11 +448,11 @@ if __name__ == '__main__':
folder = args.folder
# data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_normals', 'body_posed', 'mesh_unposed'] # 'mesh_unposed',
data_to_convert = ['mesh_static', 'mesh_normals'] # 'mesh_unposed',
data_to_convert = ['mesh_posed', 'mesh_static'] # 'mesh_unposed',
frames = sorted(os.listdir(folder + '/mesh_posed_body_no_orient'))#[args.start:args.end]
frames = sorted(os.listdir(folder + '/mesh_posed_body_no_orient'))
frames = [[frame.split('.')[0], frame.split('.')[1]] for frame in frames]
ids = list(set([frame[0] for frame in frames]))
ids = sorted(list(set([frame[0] for frame in frames])))[args.start:args.end]
frames_by_id = {id: [frame[1] for frame in frames if frame[0] == id] for id in ids}
shapedirs = get_shapedirs()
@@ -462,11 +462,11 @@ if __name__ == '__main__':
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)
# os.makedirs(os.path.join(folder, 'uv_static'), exist_ok=True)
# os.makedirs(os.path.join(folder, 'uv_normals'), exist_ok=True)
# os.makedirs(os.path.join(folder, 'uv_unposed_no_shape'), exist_ok=True)
# os.makedirs(os.path.join(folder, 'uv_body_posed_no_orient'), exist_ok=True)
# os.makedirs(os.path.join(folder, 'uv_mask'), exist_ok=True)
mask = np.ones((3638), dtype=bool)
@@ -476,15 +476,14 @@ if __name__ == '__main__':
# 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_unposed' in data_to_convert:
# mesh_unposed = np.zeros((x.shape[0], 14475, 3), dtype=np.float32)
# mesh_unposed[:,mask, :] = readPC2(os.path.join(folder, 'mesh_unposed_no_shape' ,f'{id}_{garment}.pc16'), True)['V']
if 'mesh_static' in data_to_convert:
mesh_static = np.load(os.path.join(folder, 'mesh_static', f'{id}.npz'))['mesh_static']
if 'mesh_normals' in data_to_convert:
mesh_normals = np.stack([np.load(os.path.join(folder, 'mesh_normals', f'{id}.{int(i):06d}.npz'))['mesh_normals'] for i in frames_by_id[id]], axis=0)
if 'body_posed' in data_to_convert and k != 1:
body_posed = readPC2(os.path.join(folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V']
if 'mesh_posed' in data_to_convert:
mesh_posed = np.stack([np.load(os.path.join(folder, 'mesh_posed_no_orient', f'{id}.{int(i):06d}.npz'))['mesh_posed_no_orient'] for i in frames_by_id[id]], axis=0)
start = time.perf_counter()
@@ -493,13 +492,13 @@ if __name__ == '__main__':
uv_static = uv_static[0].astype(np.float32)
if 'mesh_posed' in data_to_convert:
uv_posed = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
if 'mesh_unposed' in data_to_convert:
uv_unposed = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
if 'mesh_normals' in data_to_convert:
uv_normals = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_normals[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
if 'body_posed' in data_to_convert and k != 1:
uv_body_posed = convert_to_uv_batch_body(uv_res, m2uv['non-skirt'], uv_map['non-skirt'], body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border)
uv_posed_no_orient = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
# if 'mesh_unposed' in data_to_convert:
# uv_unposed = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
# if 'mesh_normals' in data_to_convert:
# uv_normals = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_normals[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
# if 'body_posed' in data_to_convert and k != 1:
# uv_body_posed = convert_to_uv_batch_body(uv_res, m2uv['non-skirt'], uv_map['non-skirt'], body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border)
print(time.perf_counter() - start)
@@ -515,16 +514,16 @@ if __name__ == '__main__':
# np.savez_compressed(os.path.join(folder, 'uv_mask', f'{id}_{garment}.npz'), uv_mask=uv_mask_local)
if 'mesh_posed' in data_to_convert:
for i in range(len(uv_posed)):
np.savez_compressed(os.path.join(folder, 'uv_posed_no_orient', f'{id}_{i}_{garment}.npz'), uv_posed=uv_posed[i])
if 'mesh_unposed' in data_to_convert:
for i in range(len(uv_unposed)):
np.savez_compressed(os.path.join(folder, 'uv_unposed_no_shape', f'{id}_{i}_{garment}.npz'), uv_unposed=uv_unposed[i])
if 'mesh_normals' in data_to_convert:
for i, frame_nr in enumerate(frames_by_id[id]):
np.savez_compressed(os.path.join(folder, 'uv_normals', f'{id}.{frame_nr}.npz'), uv_normals=uv_normals[i])
if 'body_posed' in data_to_convert and k != 1:
for i in range(len(uv_body_posed)):
np.savez_compressed(os.path.join(folder, 'uv_body_posed_no_orient', f'{id}_{i}.npz'), uv_body_posed=uv_body_posed[i])
for k, i in enumerate(frames_by_id[id]):
np.savez_compressed(os.path.join(folder, 'uv_posed_no_orient', f'{id}.{int(i):06d}.npz'), uv_posed_no_orient=uv_posed_no_orient[k])
# if 'mesh_unposed' in data_to_convert:
# for i in range(len(uv_unposed)):
# np.savez_compressed(os.path.join(folder, 'uv_unposed_no_shape', f'{id}_{i}_{garment}.npz'), uv_unposed=uv_unposed[i])
# if 'mesh_normals' in data_to_convert:
# for i, frame_nr in enumerate(frames_by_id[id]):
# np.savez_compressed(os.path.join(folder, 'uv_normals', f'{id}.{frame_nr}.npz'), uv_normals=uv_normals[i])
# if 'body_posed' in data_to_convert and k != 1:
# for i in range(len(uv_body_posed)):
# np.savez_compressed(os.path.join(folder, 'uv_body_posed_no_orient', f'{id}_{i}.npz'), uv_body_posed=uv_body_posed[i])
# pass
data/dataset.py
+11 -1
View File
@@ -276,7 +276,7 @@ if __name__ == '__main__':
'uv_mask',
}
dataset = Dataset('/data-local/127/cloth3d_processed', signature, resolutions=[512], device=device, frames_file='frames_train.txt')
dataset = Dataset('/data-local/127_2/cloth3d_processed_test', signature, resolutions=[512], device=device, frames_file='frames_temp.txt')
train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
smpl = SMPL(orig_smpl=False, path='smpl', device='cuda')
@@ -293,6 +293,16 @@ if __name__ == '__main__':
step = 0
for batch in train_dataloader:
batch
from custom_utils.custom_utils import save_stl
fix_data = np.load(os.path.join(dataset.folder, 'mesh_fixed', f'{batch["id"][0]}_{batch["garment"][0]}.npz'))
faces = fix_data['faces']
mask = fix_data['mask']
save_stl(f'temp/{batch["id"][0]}_{batch["frame_nr"][0].cpu().numpy()}_{batch["garment"][0]}_mesh_posed.stl', batch['mesh_posed'][0].cpu().numpy()[mask], faces)
save_stl(f'temp/{batch["id"][0]}_{batch["frame_nr"][0].cpu().numpy()}_{batch["garment"][0]}_mesh_unposed.stl', batch['mesh_unposed'][0].cpu().numpy()[mask], faces)
print(batch['id'][0])
with profiler:
m = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], batch['mesh_unposed'])[0]
data/downscale.py
+5 -23
View File
@@ -11,7 +11,7 @@ import plotly.express as ex
from torchvision.transforms import Resize
from torchvision.transforms.functional import InterpolationMode
os.environ["CUDA_VISIBLE_DEVICES"] = "4"
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from data.utils import Frame
@@ -165,32 +165,19 @@ def downscale_uv_maps_static(folder, type, sizes):
np.savez_compressed(os.path.join(folder, type + '_' + str(size), f'{frame[0]}_{frame[1]}.npz'), uv_static = down)
def downscale_uv_maps_mask(folder, type, size):
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
frames = sorted(os.listdir(os.path.join(folder, type)))[int(args.start):int(args.end)]
for frame in tqdm.tqdm(frames):
id, _ = frame.split('.')
img = np.load(os.path.join(folder, type, frame)); img = img['uv_mask'].astype(np.float32)
down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
down = np.squeeze(down.numpy().astype(bool))
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_mask = down)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127/cloth3d_processed')
parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cloth3d_processed_test')
parser.add_argument("-s", "--start", help = "start of range", default=0)
parser.add_argument("-e", "--end", help = "end of range", default=25000)
parser.add_argument("-e", "--end", help = "end of range", default=300000)
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
parser.add_argument("-d", "--debug", help = "end of range", default=False, action='store_true')
args = parser.parse_args()
# sizes = [32, 64, 128, 256]
device = 'cuda'
sizes = [32, 64, 128, 256]
sizes = [64, 128, 256]
folder = args.folder
resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BICUBIC).to(device) for res in sizes }
resizers_nearest = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST_EXACT).to(device) for res in sizes }
@@ -199,19 +186,14 @@ if __name__ == '__main__':
uv_mask_body = np.load('custom_utils/mask_256x512.npy')
uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
uv_masks_body[512] = uv_mask_body.astype(bool)
for res in [32, 64, 128, 256]:
for res in [64, 128, 256]:
down = resizers_nearest[res](torch.from_numpy(uv_mask_body[None,:,:].astype(np.int32)).to(device))
down = np.squeeze(down.cpu().numpy().astype(bool))
uv_masks_body[res] = down
uv_mask_store = { size: UVMaskStore(folder, size=size) for size in sizes }
uv_mask_store[512] = UVMaskStore(folder, size=512)
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=32)
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=64)
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=128)
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=256)
# downscale_uv_maps_static(folder=folder, type='uv_static', sizes=sizes)
downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes)
data/downscale_cape.py
+27 -7
View File
@@ -23,7 +23,7 @@ def get_img_idxs():
img_idxs[512] = img_idx
img_coords[512] = coords
uv_masks[512] = uv_mask.astype(np.int32)
for res in [32, 64, 128, 256]:
for res in [64, 128, 256]:
down = resizers_nearest[res](torch.from_numpy(uv_mask[None,:,:]).to(device)).cpu().numpy()
down = np.squeeze(down.astype(bool))
coords = np.transpose(np.stack(np.nonzero((1 - down))), (1,0))
@@ -37,7 +37,7 @@ def get_img_idxs():
uv_mask_body = np.load('custom_utils/mask_256x512.npy')
uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
uv_masks_body[512] = uv_mask_body
for res in [32, 64, 128, 256]:
for res in [64, 128, 256]:
down = resizers_nearest[res](torch.from_numpy(uv_mask_body[None,:,:]).to(device)).cpu().numpy()
down = np.squeeze(down.astype(bool))
uv_masks_body[res] = down
@@ -113,6 +113,8 @@ def downscale_uv_normals(folder, type, sizes):
frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
for frame in tqdm.tqdm(frames):
if not frame.startswith('03223_shortshort_chicken_wings'):
continue
img = np.load(os.path.join(folder, type, frame)); img = img['uv_normals']
img = fill(img, uv_masks_body[512])
for size in sizes:
@@ -140,6 +142,23 @@ def downscale_uv_maps_unposed(folder, type, sizes):
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_unposed_no_shape = down)
def downscale_uv_maps_posed(folder, type, sizes):
for size in sizes:
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
for frame in tqdm.tqdm(frames):
img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_no_orient']
mask_full = uv_mask_store[512][None]
img = fill(img, mask_full)
for size in sizes:
mask = uv_mask_store[size][None]
down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
down = down * mask[None,:,:]
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_posed_no_orient = down)
def downscale_uv_maps_static(folder, type, sizes):
for size in sizes:
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
@@ -173,7 +192,7 @@ def downscale_uv_maps_mask(folder, type, size):
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/train/male')
parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/train/female')
parser.add_argument("-s", "--start", help = "start of range", default=0)
parser.add_argument("-e", "--end", help = "end of range", default=100000)
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
@@ -182,7 +201,7 @@ if __name__ == '__main__':
device = 'cuda'
sizes = [32, 64, 128, 256]
sizes = [64, 128, 256]
folder = args.folder
resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BICUBIC).to(device) for res in sizes }
@@ -199,9 +218,10 @@ if __name__ == '__main__':
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=256)
# downscale_uv_maps_static(folder=folder, type='uv_static', sizes=sizes)
downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes)
downscale_uv_normals(folder=folder, type='uv_normals', sizes=sizes)
downscale_uv_maps_posed_body(folder=folder, type='uv_posed_body_no_orient', sizes=sizes)
# downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes)
downscale_uv_maps_posed(folder=folder, type='uv_posed_no_orient', sizes=sizes)
# downscale_uv_normals(folder=folder, type='uv_normals', sizes=sizes)
# downscale_uv_maps_posed_body(folder=folder, type='uv_posed_body_no_orient', sizes=sizes)
print()
data/preprocessing.py
+184 -10
View File
@@ -12,6 +12,7 @@ args = parser.parse_args()
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
import math
import shutil
@@ -76,6 +77,8 @@ def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_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)
ids = ['00021']
for id in tqdm(ids):
# id = '01467'
@@ -87,8 +90,8 @@ def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder):
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
# if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
# continue
verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
# if info['poses'].shape[1] == verts.shape[0]:
@@ -488,6 +491,176 @@ def create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=0, end=1)
pass
F_coll = torch.tensor(np.load("smpl/smpl_faces.npy"), device='cuda')
F_coll_skirt = torch.tensor(np.load("smpl/smpl_faces_skirt.npy"), device='cuda')
F_norm = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[None,:,:], [1,1,1]), device='cuda')
F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[None,:,:], [1,1,1]), device='cuda')
from scipy.spatial import KDTree
from pytorch3d.structures import Meshes
from torch_cluster import knn
def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor):
# flatten V and B along the batch dimension, create a batch_x and batch_y variable to restore the batch dimension after the nearest neighbour search using pytorch_geomtry knn function
batch_size = V.shape[0]
vertices_count = V.shape[1]
batch_x = torch.arange(batch_size, device=V.get_device()).repeat_interleave(vertices_count)
batch_y = torch.arange(batch_size, device=V.get_device()).repeat_interleave(14475)
V = V.reshape(-1, V.shape[-1])
B = B.reshape(-1, B.shape[-1])
idx = knn(B.to(V.dtype), V, k=1, batch_x=batch_y, batch_y=batch_x)[1]
idx = (idx - batch_x * vertices_count).reshape(batch_size, vertices_count)
return idx
def collision_loss(V: torch.Tensor, B: torch.Tensor, thr: float = .004):
F_norm_local = F_norm
mesh = Meshes(B, F_norm_local)
B_vn = torch.stack(mesh.verts_normals_list())
# compute correspondences
idx = _nearest_neighbour(V, B)
# compute loss
D = V - torch.stack([B[i,idx[i],:] for i in range(V.shape[0])])
b_vn = torch.stack([B_vn[i,idx[i],:] for i in range(V.shape[0])])
dot = torch.einsum('abc,abc->ab', D, b_vn)
thresholded = torch.minimum(dot - thr, torch.zeros_like(dot))
loss = torch.sum(thresholded ** 2, axis=-1)
return loss
def area_loss(pred, gt, faces):
"""
Calculates the area loss between predicted and ground truth meshes.
Args:
pred (torch.Tensor): Predicted mesh vertices of shape (batch_size, num_vertices, 3).
gt (torch.Tensor): Ground truth mesh vertices of shape (batch_size, num_vertices, 3).
faces (torch.Tensor): Mesh faces of shape (num_faces, 3).
Returns:
torch.Tensor: Area loss between predicted and ground truth meshes.
"""
# Calculate the area of each face for predicted and ground truth meshes
pred_faces = pred[:, faces, :]
gt_faces = gt[:, faces, :]
pred_edges1 = pred_faces[:, :, 1, :] - pred_faces[:, :, 0, :]
pred_edges2 = pred_faces[:, :, 2, :] - pred_faces[:, :, 0, :]
gt_edges1 = gt_faces[:, :, 1, :] - gt_faces[:, :, 0, :]
gt_edges2 = gt_faces[:, :, 2, :] - gt_faces[:, :, 0, :]
pred_cross = torch.cross(pred_edges1, pred_edges2, dim=2)
gt_cross = torch.cross(gt_edges1, gt_edges2, dim=2)
pred_area = 0.5 * torch.norm(pred_cross, dim=2)
gt_area = 0.5 * torch.norm(gt_cross, dim=2)
# Calculate the area loss
area_loss = torch.mean(torch.abs(pred_area - gt_area))
return area_loss
def create_mesh_unposed_no_shape_static(folder, folder_reg, out_folder, start=0, end=1):
ids = sorted(os.listdir(folder))[start:end]
os.makedirs(os.path.join(out_folder, 'mesh_static_unposed_no_shape'), exist_ok=True)
F = np.load('smpl/smpl_faces.npy')
# F_skirt = np.load('smpl/smpl_faces_skirt.npy')
ids = ['00005']
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')]
shape = np.load(os.path.join(out_folder, 'shape', f'{id}.npy'))[None,:]
shape = torch.tensor(shape, dtype=torch.float32, device='cuda')
smpl = SMPL(path='smpl', device='cuda', batch_size=1)
pose = np.zeros((24, 3), np.float32)
pose[1,2] = .15
pose[2,2] = -.15
pose = torch.tensor(pose.flatten()[None,:], 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, faces, _, _ = readOBJ(os.path.join(folder, id, f'{garment}.obj'))
faces = quads2tris(faces)
mesh_static = np.matmul(mesh_static, np.array([[1,0,0],[0,0,-1],[0,1,0]]))
unposed = mesh_static
unposed = torch.tensor(unposed[None,:,:], dtype=torch.float32, device='cuda')
unposed.requires_grad = True
mesh_static = torch.tensor(mesh_static[None,:,:], dtype=torch.float32, device='cuda')
optimizer = torch.optim.Adam([unposed], lr=0.0001)
for i in range(10000):
optimizer.zero_grad()
body = smpl(pose, shape)[0]
c_loss = collision_loss(unposed, body)
a_loss = area_loss(unposed, mesh_static, faces)
loss : torch.Tensor = c_loss + a_loss
# Backward pass
loss.backward()
optimizer.step()
pose = pose / 1.0005
# Print the loss
if i % 50 == 0:
pose_sum = np.abs(pose.detach().cpu().numpy()).sum().item()
loss_str = loss.detach().cpu().numpy()[0]
c_loss_str = c_loss.detach().cpu().numpy()[0]
a_loss_str = a_loss.detach().cpu().numpy()
print(f"it: {i} loss: {loss_str:.10f} c_loss: {c_loss_str:.10f}, a_loss: {a_loss_str:.8f}, pose: {pose_sum:.8f}")
unposed = unposed.detach().cpu().numpy()
body = body.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)
writePC2(os.path.join(out_folder, 'mesh_static_unposed_no_shape', f'{id}_{garment}.pc16'), unshaped, True)
verts = unposed
# verts = unshaped
frame = 0
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)])
# add body mesh
fig.add_trace(go.Mesh3d(x=body[frame,:,0], y=body[frame,:,2], z=body[frame,:,1], i=F[:,0], j=F[:,1], k=F[:,2], color='lightblue', opacity=1))
fig.update_layout(scene = dict(
xaxis = dict(range=[-1,1],),
yaxis = dict(range=[-1,1],),
zaxis = dict(range=[-1,1],),),)
fig.update_layout(scene_aspectmode='cube')
fig.show()
# 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)
@@ -879,13 +1052,13 @@ def generate_fixed_masks(folder, out_folder):
if __name__ == '__main__':
# folder = '/data-net/dataset/cloth3d/test'
# folder_reg = '/data-local/127/cloth3d_registration/test'
# out_folder = '/data-local/127_2/cloth3d_processed_test'
folder = '/data-net/dataset/cloth3d/test'
folder_reg = '/data-local/127/cloth3d_registration/test'
out_folder = '/data-local/127_2/cloth3d_processed_test'
folder = '/data-net/dataset/cloth3d/train'
folder_reg = '/data-local/127/cloth3d_registration/train'
out_folder = '/data-local/127/cloth3d_processed'
# 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')
@@ -898,13 +1071,14 @@ if __name__ == '__main__':
# remove_irregular_poses(folder, folder_reg, out_folder)
# double_check_irregular_poses(folder, folder_reg, out_folder)
# create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder)
create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder)
# verify_registration_ans_mask(folder, out_folder)
# create_pose_shape(folder, out_folder)
# create_mesh_static(folder, out_folder)
# create_mesh_body_posed_no_orient(folder, out_folder)
# create_mesh_normals(folder, out_folder)
# create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end)
# create_mesh_unposed_no_shape_static(folder, folder_reg, out_folder, start=args.start, end=args.end)
# calculate_average(folder, out_folder)
# create_frames_file(folder, out_folder)
@@ -915,4 +1089,4 @@ if __name__ == '__main__':
# folder = '/data-local/127_2/cape_processed/test/female'
# create_mesh_normals_cape(folder, None)
generate_fixed_masks(folder, out_folder)
# generate_fixed_masks(folder, out_folder)
data/utils.py
+17 -1
View File
@@ -1,7 +1,23 @@
import tables as tb
class Frame:
def __init__(self, id: str, frame_nr: int, garment: str=''):
def __init__(self, id: str ='', frame_nr: int= 0, garment: str='', frame = None, frame_str: str = ''):
if frame is not None:
self.id: str = frame.id
self.frame_nr: int = frame.frame_nr
self.garment: str = frame.garment
assert id == '' and frame_nr == 0 and garment == '' and frame_str == ''
elif frame_str != '':
if ' ' in frame_str:
self.id: str = frame_str.split()[0]
self.frame_nr: int = int(frame_str.split()[1])
self.garment: str = frame_str.split()[2]
else:
self.id: str = frame_str.split('_')[0]
self.frame_nr: int = int(frame_str.split('_')[1])
self.garment: str = frame_str.split('_')[2]
assert id == '' and frame_nr == 0 and garment == '' and frame is None
else:
self.id: str = id
self.frame_nr: int = frame_nr
self.garment: str = garment
docker/Dockerfile2
+4 -1
View File
@@ -9,8 +9,11 @@ ARG USER_ID=11017
ARG GROUP_NAME=guest
ARG GROUP_ID=11000
# fix for tzdata asking for user input
ARG DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC
# Update the system and install required packages
RUN apt-get update && apt-get install -y sudo git apt-utils htop
RUN apt-get update && apt-get install -y sudo git apt-utils htop ffmpeg libsm6 libxext6
# Create a new group with the same GID as the host
RUN groupadd --gid $GROUP_ID $GROUP_NAME
docker/Dockerfile_hood
+48
View File
@@ -0,0 +1,48 @@
# Use an existing image as the base
# FROM nvcr.io/nvidia/pytorch:23.04-py3
# FROM pytorch/pytorch:1.13.1-cuda11.6-cudnn8-devel
FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel
#
# Add the current user to the image
ARG USER_NAME=lhunor
ARG USER_ID=11017
ARG GROUP_NAME=guest
ARG GROUP_ID=11000
# Update the system and install required packages
RUN apt-get update && apt-get install -y sudo git apt-utils htop
# Create a new group with the same GID as the host
RUN groupadd --gid $GROUP_ID $GROUP_NAME
# Create a new user with the same UID and GID as the host
RUN useradd --uid $USER_ID --gid $GROUP_ID --shell /bin/bash --create-home $USER_NAME
# don't require password with sudo, for convenience
# not the safest thing to do, but hopefully okay inside the container
RUN echo '%sudo ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers
# Add the new user to the sudo group
RUN usermod -aG sudo $USER_NAME
# install conda stuff still as root
RUN conda update -n base conda
RUN conda install -c fvcore -c iopath -c conda-forge fvcore iopath -y
RUN conda install -c bottler nvidiacub -y
RUN conda install pytorch3d -c pytorch3d -y
# Set the default user to the new user
USER $USER_NAME
WORKDIR /workspace
ENV PATH="$PATH:/home/$USER_NAME/.local/bin"
RUN pip install --upgrade pip
# defining the volumes that will be mounted
VOLUME [ "/data-net" ]
VOLUME [ "/data-local" ]
VOLUME [ "/home-net" ]
# Start the container with a bash shell
CMD ["/bin/bash"]
docker/docker-compose.yml
+21
View File
@@ -105,3 +105,24 @@ services:
reservations:
devices:
- capabilities: [gpu]
lhunor_hood:
image: lhunor_hood:pt2.0.1
container_name: lhunor_hood_pt2.0.1
ipc: host
tty: true
ulimits:
memlock: -1
stack: 67108864
volumes:
- '/home-net/lhunor:/home-net'
- '/data-net/guest/lhunor/data:/data-net'
- '/data-fast/126-data1/lhunor/data:/data-local/126'
- '/data-fast/127-data1/lhunor/data:/data-local/127'
- '/data-fast/127-data2/lhunor/data:/data-local/127_2'
- '/data-fast/128-data1/lhunor/data:/data-local/128'
deploy:
resources:
reservations:
devices:
- capabilities: [gpu]
losses/losses.py
+54 -4
View File
@@ -32,7 +32,7 @@ def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor):
# flatten V and B along the batch dimension, create a batch_x and batch_y variable to restore the batch dimension after the nearest neighbour search using pytorch_geomtry knn function
batch_size = V.shape[0]
vertices_count = V.shape[1]
batch_x = torch.arange(batch_size, device=V.get_device()).repeat_interleave(V.shape[1])
batch_x = torch.arange(batch_size, device=V.device).repeat_interleave(V.shape[1])
V = V.reshape(-1, V.shape[-1])
B = B.reshape(-1, B.shape[-1])
idx = knn(B.to(V.dtype), V, k=1, batch_x=batch_x, batch_y=batch_x)[1]
@@ -58,7 +58,9 @@ def collision_loss(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: fl
F_norm_local = F_norm
elif B.shape[1] == 3638:
F_norm_local = F_norm_orig
mesh = Meshes(B, F_norm_local)
if V.device != F_norm_local.device:
F_norm_local = F_norm_local.to(V.device)
mesh = Meshes(B, F_norm_local[:B.shape[0],:,:])
B_vn = torch.stack(mesh.verts_normals_list())
# compute correspondences
@@ -74,6 +76,54 @@ def collision_loss(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: fl
return loss, ratio
def _nearest_neighbour_error(V: torch.Tensor, B: torch.Tensor):
# flatten V and B along the batch dimension, create a batch_x and batch_y variable to restore the batch dimension after the nearest neighbour search using pytorch_geomtry knn function
batch_size = V.shape[0]
vertices_count = V.shape[1]
batch_x = torch.arange(batch_size, device=V.device).repeat_interleave(B.shape[1])
batch_y = torch.arange(batch_size, device=V.device).repeat_interleave(V.shape[1])
V = V.reshape(-1, V.shape[-1])
B = B.reshape(-1, B.shape[-1])
idx = knn(B.to(V.dtype), V, k=1, batch_x=batch_x, batch_y=batch_y)[1]
idx = (idx - batch_y * vertices_count).reshape(batch_size, vertices_count)
return idx
# def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor):
# # create an empty tensor to store the indices of the nearest neighbors
# batch_size = V.shape[0]
# vertices_count = V.shape[1]
# idx = torch.zeros((batch_size, vertices_count), dtype=torch.long, device=V.device)
# # loop over the batch dimension and find the nearest neighbor for each point in B
# for i in range(batch_size):
# tree = KDTree(B[i].cpu().numpy())
# _, nearest = tree.query(V[i].cpu().numpy(), k=1)
# idx[i] = torch.from_numpy(nearest)
# return idx
def collision_error(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: float = 0.0005):
if B.shape[1] == 14475:
F_norm_local = F_norm
elif B.shape[1] == 3638:
F_norm_local = F_norm_orig
if V.device != F_norm_local.device:
F_norm_local = F_norm_local.to(V.device)
mesh = Meshes(B, F_norm_local[:B.shape[0],:,:])
B_vn = torch.stack(mesh.verts_normals_list())
# compute correspondences
idx = _nearest_neighbour_error(V, B)
# compute loss
D = V - torch.stack([B[i,idx[i],:] for i in range(V.shape[0])])
b_vn = torch.stack([B_vn[i,idx[i],:] for i in range(V.shape[0])])
dot = torch.einsum('abc,abc->ab', D, b_vn)
thresholded = torch.minimum(dot - thr, torch.zeros_like(dot))
loss = torch.sum(thresholded ** 2, axis=-1)
ratio = torch.sum((thresholded < 0).to(torch.float32)) / V.shape[1]
return loss, ratio
""" COLLISION FUNCTIONS END """
@@ -218,9 +268,9 @@ def normal_loss(pred, gt, mask, is_skirt=False, is_smpl_orig=False):
else:
raise NotImplementedError
mesh = Meshes(pred, f_norm)
mesh = Meshes(pred, f_norm[:pred.shape[0], :, :])
pred_norm = torch.stack(mesh.verts_normals_list())
mesh = Meshes(gt, f_norm)
mesh = Meshes(gt, f_norm[:pred.shape[0], :, :])
gt_norm = torch.stack(mesh.verts_normals_list())
return torch.mean(torch.mean((torch.sum(torch.abs(pred_norm - gt_norm), axis=-1) * mask), axis=-1))
models/pyramid_level.py
+1 -1
View File
@@ -39,7 +39,7 @@ class PyramidLevel(nn.Module):
def sample(self, z_cond, uv_mask, noise=None):
if noise is None:
eps = torch.randn(z_cond.shape).to(uv_mask.device)
eps = torch.normal(mean=0.0, std=1.0, size=z_cond.shape).to(uv_mask.device)
else:
eps = noise
models/pyramid_model.py
+54 -1
View File
@@ -10,6 +10,29 @@ from models.pyramid_level import PyramidLevel
from torchvision.transforms import Resize
from torchvision.transforms.functional import InterpolationMode
# class SampleParallel(nn.Module):
# def __init__(self, levels):
# super().__init__()
# self.levels = nn.ModuleList(levels)
# def forward(self, z_cond, uv_masks, noise=None):
# level_outs = nn.parallel.data_parallel(self._forward_level, [(self.levels[i], z_cond, uv_masks[i], noise) for i in range(len(uv_masks))], dim=0, device_ids=[0, 0, 0, 0], output_device=0)
# # level_outs = nn.parallel.data_parallel(self.levels, [(z_cond, uv_masks[i], noise) for i in range(len(uv_masks))], device_ids=[0, 0, 0, 0], output_device=0)
# out = torch.zeros((uv_masks[0].shape[0], 3, 64, 32)).to(uv_masks[0].device)
# for i in range(len(self.levels)):
# level_out = level_outs[i]
# if i != len(self.levels) - 1:
# out = out + level_out
# out = self.levels[i].enhance(out, uv_masks[i])
# out = self.resizers[self.levels[i+1].size](out)
# out = out * uv_masks[i+1][:,None,:,:]
# else:
# out = out + level_out
# out = out * uv_masks[0][:,None,:,:]
# return out
# def _forward_level(self, level, z_cond, uv_mask, noise=None):
# return level(z_cond, uv_mask, noise=noise)
class Pyramid(nn.Module):
@@ -43,15 +66,22 @@ class Pyramid(nn.Module):
self.use_normals = config['use_normals']
def sample(self, uv_template, uv_body_posed, uv_normals, uv_mask, noise=None):
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_normals)
# uv_masks = [uv_mask[64], uv_mask[128], uv_mask[256], uv_mask[512]]
# sampleP = SampleParallel(self.levels)
# out = sampleP(z_cond, uv_masks, noise=noise)
# return out
assert len(self.levels) == 4
out = torch.zeros((uv_template.shape[0], 3, 64, 32)).to(uv_template.device)
for i in range(len(self.levels)):
level_out = self.levels[i].sample(z_cond, uv_mask[self.levels[i].size], noise=noise)
level_out = self.levels[i].sample(z_cond, uv_mask[self.levels[i].size], noise=noise[i] if noise is not None else None)
if i != len(self.levels) - 1:
out = out + level_out
out = self.levels[i].enhance(out, uv_mask[self.levels[i].size])
@@ -64,6 +94,29 @@ class Pyramid(nn.Module):
return out
def sample_level(self, uv_template, uv_body_posed, uv_normals, uv_mask, prev, level, noise=None):
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_normals)
assert len(self.levels) == 4
i = level
out = prev
level_out = self.levels[i].sample(z_cond, uv_mask[self.levels[i].size], noise=noise[i] if noise is not None else None)
if i != len(self.levels) - 1:
out = out + level_out
out = self.levels[i].enhance(out, uv_mask[self.levels[i].size])
out = self.resizers[self.levels[i+1].size](out)
out = out * uv_mask[self.levels[i+1].size][:,None,:,:]
else:
out = out + level_out
out = out * uv_mask[512][:,None,:,:]
return out
def encode_conds(self, uv_template, uv_body_posed, uv_norm):
z_static = self.static_enc(uv_template)
z_body = self.body_enc(uv_body_posed)
run.py
+8 -7
View File
@@ -1,6 +1,6 @@
import os
import sys
from configs.cvae_norm_1 import config
from configs.iccvw.pyr_64_128_256_512_skirt_e2e_var import config
import argparse
# parser = argparse.ArgumentParser(description='')
@@ -12,14 +12,15 @@ import argparse
if __name__ == '__main__':
import wandb
# wandb.login(key='******************************************')
wandb.login(host='https://wandb.lnrweb.com', key='local-40a577c9eccab708f1aa81027f21a995596a5838')
# wandb.login(host='https://wandb.lnrweb.com', key='local-40a577c9eccab708f1aa81027f21a995596a5838')
# wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
name = config['name']
script = config['script']
config = [key for key in sys.modules.keys() if key.startswith('configs.')][-1]
current_file = '/'.join(__file__.split('/')[-2:])
print(f'nohup python {script} --name lhunor_{name} --config {config} > nohup/torch/{name}.out 2>&1 &')
if args['debug']:
os.system(f'python {script} --name lhunor_{name} --config {config}')
else:
os.system(f'nohup python {script} --name lhunor_{name} --config {config} > nohup/torch/{name}.out 2>&1 &')
print(f'nohup python {script} --name lhunor_{name} --config {config} > nohup/iccvw/{name}.out 2>&1 &')
# if args['debug']:
# os.system(f'python {script} --name lhunor_{name} --config {config}')
# else:
# os.system(f'nohup python {script} --name lhunor_{name} --config {config} > nohup/torch/{name}.out 2>&1 &')
run_parallel.py
+9 -4
View File
@@ -12,19 +12,24 @@ if __name__ == '__main__':
# ids = os.listdir(folder + '/mesh_posed_no_orient')#[280000:]
skirt = 0
obj_folder = '/data-net/dataset/cloth3d/train'
folder = '/data-local/127/train'
folder = '/data-local/127_2/cape_processed/train/male'
# folder = '/data-local/127_2/cloth3d_processed_test'
count = 0
step = 100
step = 35
# !!!!!!! change skirt parameter in script
while count <1200:
while count < 550:
# os.system('nohup python convert_to_uv_cape.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
# os.system('nohup python convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_unposed_onehot_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
# os.system('nohup python convert_to_uv_local.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_local_unposed_onehot_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_local_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
# os.system('nohup python data/convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_2_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv/train_2_' + str(count) + '_' + str(count+step) + '_again.out' + ' 2>&1 &')
os.system(f'nohup python data/calculate_normals.py -f {folder} -f_obj {obj_folder} -s {count} -e {(count+step)} -n normals_train_calc_{count}_{(count+step)} > nohup/normals/normals_train_calc_{count}_{(count+step)}.out' + ' 2>&1 &')
# os.system(f'nohup python data/convert_to_uv_cape.py -f {folder} -s {count} -e {(count+step)} -n to_uv_cape_m_test_{count}_{(count+step)} > nohup/to_uv/cape_m_test_{count}_{(count+step)}.out' + ' 2>&1 &')
# os.system(f'nohup python data/convert_to_uv.py -f {folder} -s {count} -e {(count+step)} -n to_uv_{count}_{(count+step)} > nohup/to_uv/test_{count}_{(count+step)}.out' + ' 2>&1 &')
# os.system(f'nohup python data/downscale.py -f {folder} -s {count} -e {(count+step)} -n down_test_{count}_{(count+step)} > nohup/down/test_{count}_{(count+step)}.out' + ' 2>&1 &')
# os.system(f'nohup python data/downscale_cape.py -f {folder} -s {count} -e {(count+step)} -n down_train_cape_m_{count}_{(count+step)} > nohup/down/train_cape_m_{count}_{(count+step)}.out' + ' 2>&1 &')
os.system(f'nohup python scripts/calculate_chamfer.py -s {count} -e {(count+step)} -n hood_130k_{count}_{(count+step)} > nohup/hood/hood_130k_{count}_{(count+step)}.out' + ' 2>&1 &')
# os.system('nohup python convert_to_uv_local_single.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
# os.system('nohup python convert_to_uv_local_single.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_valid_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_valid_skirt_local_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
# os.system('nohup python downscale.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + ' > nohup/downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
scripts/calculate.py
+34
View File
@@ -0,0 +1,34 @@
import numpy as np
error_1 = dict()
error_2 = dict()
#end2end with reg
# error_1[0] = {'Jumpsuit': 0.00888, 'Trousers': 0.01193, 'Tshirt': 0.01065, 'Top': 0.00725}
# error_1[1] = {'Jumpsuit': 0.00491, 'Trousers': 0.00704, 'Tshirt': 0.00531, 'Top': 0.00393}
# error_1[2] = {'Jumpsuit': 0.00379, 'Trousers': 0.00547, 'Tshirt': 0.00404, 'Top': 0.00309}
# error_1[3] = {'Jumpsuit': 0.00349, 'Trousers': 0.00502, 'Tshirt': 0.00382, 'Top': 0.00281}
# error_2[0] = {'Dress': 0.01009, 'Skirt': 0.01379}
# error_2[1] = {'Dress': 0.0055, 'Skirt': 0.0077}
# error_2[2] = {'Dress': 0.00418, 'Skirt': 0.00582}
# error_2[3] = {'Dress': 0.00359, 'Skirt': 0.00496}
#end2end without reg
error_1[0] = {'Jumpsuit': 0.00927, 'Trousers': 0.01229, 'Tshirt': 0.01113, 'Top': 0.0076}
error_1[1] = {'Jumpsuit': 0.00503, 'Trousers': 0.00721, 'Tshirt': 0.0054, 'Top': 0.004}
error_1[2] = {'Jumpsuit': 0.00402, 'Trousers': 0.0059, 'Tshirt': 0.00421, 'Top': 0.00319}
error_1[3] = {'Jumpsuit': 0.00347, 'Trousers': 0.00528, 'Tshirt': 0.00365, 'Top': 0.00271}
error_2[0] = {'Dress': 0.01005, 'Skirt': 0.01395}
error_2[1] = {'Dress': 0.0056, 'Skirt': 0.00787}
error_2[2] = {'Dress': 0.00441, 'Skirt': 0.00634}
error_2[3] = {'Dress': 0.00379, 'Skirt': 0.0054}
for i in range(4):
errors_mean = dict()
error = {**error_1[i], **error_2[i]}
for category in error.keys():
errors_mean[category] = round(np.mean(error[category]), 5)
print(f'Level {i}: {errors_mean}')
mean = round(np.mean(list(errors_mean.values())), 5)
print(f'Final level {i}: {mean}')
scripts/calculate_area_diff.py
+47
View File
@@ -0,0 +1,47 @@
import trimesh
import numpy as np
cases = [
# ('06461_237_Trousers', 'temp/sampling_pyramid/06461_237_Trousers', [0,8,12,13]),
# ('01171_213_Dress', 'temp/sampling_pyramid/01171_213_Dress', [0,6,10,15]),
('01852_70_Tshirt_lvl_0', 'temp/sampling_pro_pyramid/01852_70_Tshirt', [15,10,1,5,13]),
('01852_70_Tshirt_lvl_1', 'temp/sampling_pro_pyramid/01852_70_Tshirt', [15,14,3,5,0 ]),
('01852_70_Tshirt_lvl_2', 'temp/sampling_pro_pyramid/01852_70_Tshirt', [15,12,0,8,6]),
('01852_70_Tshirt_lvl_3', 'temp/sampling_pro_pyramid/01852_70_Tshirt', [15,2,8,1,0]),
]
for case in cases:
# Load the original mesh
folder = case[1]
original_mesh = trimesh.load(f'{folder}/{case[0]}_{case[2][0]}.stl')
# Load the variation meshes
variation_meshes = []
for i in case[2][1:]:
variation_meshes.append(trimesh.load(f'{folder}/{case[0]}_{i}.stl'))
# Calculate area of original mesh
# original_area = original_mesh.area
# Calculate area difference for each variation
area_diffs = []
for variation_mesh in variation_meshes:
# calculate per vertex euclid distance
euclid = np.linalg.norm(original_mesh.vertices - variation_mesh.vertices, axis=1) * 1000
euclid_diff = np.mean(euclid)
# area_diff = abs(variation_mesh.area - original_area) / original_area
area_diffs.append(euclid_diff)
# Average area differences
# avg_area_diff = np.mean(area_diffs)
variance = np.var(area_diffs)
# Print results
print(f'Case: {case[0]}')
# print('Original area:', original_area)
print('Euclidean differences:', area_diffs)
print('Euclidean variance:', variance)
# print('Average Euclidean difference:', avg_area_diff)
scripts/calculate_chamfer.py
+165 -36
View File
@@ -4,47 +4,40 @@ sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
os.environ["CUDA_VISIBLE_DEVICES"] = '6'
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
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from tqdm import tqdm
from data.dataset import Dataset
from data.utils import Frame
from data.cloth3d_io import readPC2Frame, readPC2
from data.preprocessing import deorient
from utils.checkpoints import CheckpointSaver, load_model
from torch.utils.data import DataLoader
from utils.utils import text_to_html, pickle_load
from utils.postprocess import fix_collision
from losses.losses import collision_loss, normal_loss, kl_loss
from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary_html, maskFacesBatch, get_summary_str, plot_3d_collision, group_inputs_for_levels
from plotly.subplots import make_subplots
import pytorch3d.loss
from scipy.spatial.distance import cdist
from data.cloth3d_io import readPC2Frame, readPC2, readOBJ
from utils.utils import pickle_load
from scipy.spatial import cKDTree
from smpl.smpl_torch import SMPL
from scipy.spatial.transform import Rotation as R
import trimesh as tm
from losses.losses import collision_loss, normal_loss, kl_loss, collision_error
parser = argparse.ArgumentParser()
parser.add_argument("-s", "--start", help = "start of range", default=0, type=int)
parser.add_argument("-e", "--end", help = "end of range", default=550, type=int)
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
parser.add_argument('--debug', action='store_true')
args = parser.parse_args()
def chamfer(points1, points2, points2_tree: cKDTree):
tree1 = cKDTree(points1)
dist1, _ = tree1.query(points2, workers=40)
dist2, _ = points2_tree.query(points1, workers=40)
dist1, _ = tree1.query(points2, workers=10)
dist2, _ = points2_tree.query(points1, workers=10)
chamfer_dist = np.mean(dist1) + np.mean(dist2)
return chamfer_dist
@@ -60,34 +53,116 @@ def deorient(verts, pose, inverse=True):
return rotated
def get_collision_depth(m1, m2):
cm = tm.collision.CollisionManager()
cm.add_object('m1', m1)
cm.add_object('m2', m2)
data = cm.in_collision_internal(return_data=True)[1]
# c = np.asarray([d.index('m2') for d in data])
# count = np.unique(c).shape[0]
depth_sum = np.sum([d.depth for d in data])
return depth_sum
def get_area_diff(m1, m2):
return np.abs(m1.area - m2.area)
device = 'cuda'
def get_collision_loss(pred, body, mask):
c_loss, c_ratio = collision_error(torch.from_numpy(pred).unsqueeze(0).to(device), torch.from_numpy(body).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device))
return c_loss.cpu().numpy()[0], c_ratio.cpu().numpy()
def get_collision_loss_2(pred, body, mask):
_pred = np.zeros((14475, 3), dtype=np.float32)
_pred[mask] = pred
c_loss, c_ratio = collision_loss(torch.from_numpy(_pred).unsqueeze(0).to(device), torch.from_numpy(body).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device), thr=0.0005)
return c_loss.cpu().numpy()[0], c_ratio.cpu().numpy()
def calculate():
errors_hood = []
collisions_hood = []
areas_hood = []
collision_ratios_hood = []
if calculate_ours:
errors_ours = []
garments = sorted(os.listdir(hood_folder))
collisions_ours = []
areas_ours = []
collision_ratios_ours = []
garments = sorted(os.listdir(hood_folder))[args.start:args.end]
garments = [g[:-4] for g in garments]
body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32)
for garment in tqdm(garments):
picker = 'skirt' if garment.split('_')[1] in ['Dress', 'Skirt'] else 'nonskirt'
id, garment_type = garment.split('_')
pred_hood = pickle_load(f'{hood_folder}/{garment}.pkl')['verts']
pred_hood_data = pickle_load(f'{hood_folder}/{garment}.pkl')
pred_hood, pred_hood_faces = pred_hood_data['verts'], pred_hood_data['faces']
gt = readPC2(os.path.join(cloth3d_folder, id, garment_type + '.pc16'), True)['V'][2:]
template, template_faces, _, _ = readOBJ(os.path.join(cloth3d_folder, id, garment_type + '.obj'))
mask = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')['mask']
pose = np.load(f'/data-local/127_2/cloth3d_processed_test/pose/{id}.npy')[2:]
gt = deorient(gt, pose, True)
body = readPC2(f'/data-local/127_2/cloth3d_processed_test/mesh_body_posed_no_orient/{id}.pc16', True)['V'][2:]
fix_data = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')
faces = fix_data['faces']
m_template = tm.Trimesh(vertices=template, faces=template_faces, process=False, validate=False)
for i in range(pred_hood.shape[0]):
gt_tree = cKDTree(gt[i])
try:
pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'].item()['lvl3']
if calculate_ours:
# pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'].item()['lvl3']
pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0']
except:
print('missed one')
continue
# m_body = tm.Trimesh(vertices=body[i], faces=body_faces, process=False, validate=False)
errors_hood.append((f'{id}_{i}_{garment_type}', chamfer(pred_hood[i], gt[i], gt_tree)))
errors_ours.append((f'{id}_{i}_{garment_type}', chamfer(pred_ours[mask,:], gt[i], gt_tree)))
m_hood = tm.Trimesh(vertices=pred_hood[i], faces=pred_hood_faces, process=False, validate=False)
# collision_hood = get_collision_depth(m_body, m_hood)
# collisions_hood.append((f'{id}_{i}_{garment_type}', collision_hood))
area_hood = get_area_diff(m_template, m_hood)
areas_hood.append((f'{id}_{i}_{garment_type}', area_hood))
collision_ratio_hood = get_collision_loss(pred_hood[i], body[i], mask)[1]
collision_ratios_hood.append((f'{id}_{i}_{garment_type}', collision_ratio_hood))
if calculate_ours:
errors_ours.append((f'{id}_{i}_{garment_type}', chamfer(pred_ours[mask], gt[i], gt_tree)))
m_ours = tm.Trimesh(vertices=pred_ours[mask], faces=faces, process=False, validate=False)
# collision_ours = get_collision_depth(m_body, m_ours)
# collisions_ours.append((f'{id}_{i}_{garment_type}', collision_ours))
area_ours = get_area_diff(m_template, m_ours)
areas_ours.append((f'{id}_{i}_{garment_type}', area_ours))
collision_ratio_ours = get_collision_loss(pred_ours[mask], body[i], mask)[1]
# collision_ratio_ours_2 = get_collision_loss_2(pred_ours[mask], body[i], mask)[1]
collision_ratios_ours.append((f'{id}_{i}_{garment_type}', collision_ratio_ours))
# fig = go.Figure(data=[go.Mesh3d(x=pred_ours[mask][:,0], y=pred_ours[mask][:,2], z=pred_ours[mask][:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='blue')])
# fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
# 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', scene_aspectratio=dict(x=1, y=1, z=1))
# fig.show()
# fig = go.Figure(data=[go.Mesh3d(x=gt[i][:,0], y=gt[i][:,2], z=gt[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')])
# fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
# 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', scene_aspectratio=dict(x=1, y=1, z=1))
# fig.show()
# fig = go.Figure(data=[go.Mesh3d(x=pred_hood[i][:,0], y=pred_hood[i][:,2], z=pred_hood[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')])
# fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
# 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', scene_aspectratio=dict(x=1, y=1, z=1))
# fig.show()
# pass
# fig = go.Figure()
# fig.add_trace(go.Scatter3d(x=gt[i][:, 0], y=gt[i][:, 2], z=gt[i][:, 1], mode='markers', marker=dict(size=1)))
@@ -95,6 +170,7 @@ def calculate():
# 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()
# pass
# fig = go.Figure()
# fig.add_trace(go.Scatter3d(x=gt[i][:, 0], y=gt[i][:, 2], z=gt[i][:, 1], mode='markers', marker=dict(size=1)))
@@ -105,30 +181,83 @@ def calculate():
# save errors to a csv file
with open(f'{results_folder}/hood/{name}_hood.csv', 'w') as f:
os.makedirs(f'{results_folder}/hood/{name}_hood', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_hood/chamfer', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_hood/collision', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_hood/area', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_hood/collision_ratio', exist_ok=True)
for frame, error in errors_hood:
frame_str = ','.join(str(frame).split('_'))
f.write(f'{frame_str},{error}\n')
np.save(f'{results_folder}/hood/{name}_hood/chamfer/{frame}.npy', error)
for frame, error in collisions_hood:
np.save(f'{results_folder}/hood/{name}_hood/collision/{frame}.npy', error)
for frame, error in areas_hood:
np.save(f'{results_folder}/hood/{name}_hood/area/{frame}.npy', error)
for frame, error in collision_ratios_hood:
np.save(f'{results_folder}/hood/{name}_hood/collision_ratio/{frame}.npy', error)
os.makedirs(f'{results_folder}/hood/{name}_ours', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_ours/chamfer', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_ours/collision', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_ours/area', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_ours/collision_ratio', exist_ok=True)
with open(f'{results_folder}/hood/{name}_ours.csv', 'w') as f:
for frame, error in errors_ours:
frame_str = ','.join(str(frame).split('_'))
f.write(f'{frame_str},{error}\n')
np.save(f'{results_folder}/hood/{name}_ours/chamfer/{frame}.npy', error)
for frame, error in collisions_ours:
np.save(f'{results_folder}/hood/{name}_ours/collision/{frame}.npy', error)
for frame, error in areas_ours:
np.save(f'{results_folder}/hood/{name}_ours/area/{frame}.npy', error)
for frame, error in collision_ratios_ours:
np.save(f'{results_folder}/hood/{name}_ours/collision_ratio/{frame}.npy', error)
# with open(f'{results_folder}/hood/{name}/chamfer_hood.csv', 'w') as f:
# for frame, error in errors_hood:
# frame_str = ','.join(str(frame).split('_'))
# f.write(f'{frame_str},{error}\n')
# with open(f'{results_folder}/hood/{name}/collision_hood.csv', 'w') as f:
# for frame, error in collisions_hood:
# frame_str = ','.join(str(frame).split('_'))
# f.write(f'{frame_str},{error}\n')
# with open(f'{results_folder}/hood/{name}/area_hood.csv', 'w') as f:
# for frame, error in areas_hood:
# frame_str = ','.join(str(frame).split('_'))
# f.write(f'{frame_str},{error}\n')
# if calculate_ours:
# with open(f'{results_folder}/hood/{name}_chamfer_ours.csv', 'w') as f:
# for frame, error in errors_ours:
# frame_str = ','.join(str(frame).split('_'))
# f.write(f'{frame_str},{error}\n')
# with open(f'{results_folder}/hood/{name}/collision_ours.csv', 'w') as f:
# for frame, error in collisions_ours:
# frame_str = ','.join(str(frame).split('_'))
# f.write(f'{frame_str},{error}\n')
# with open(f'{results_folder}/hood/{name}/area_ours.csv', 'w') as f:
# for frame, error in areas_ours:
# frame_str = ','.join(str(frame).split('_'))
# f.write(f'{frame_str},{error}\n')
if __name__ == '__main__':
name = 'hood_2'
name = 'hood_3_var'
hood_folder = '/data-local/127_2/results/hood/all_2'
hood_folder = '/data-local/127_2/results/hood/all_3'
ours_folder = dict()
ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_200000'
ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_2_250000'
ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_var_3_450000'
ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_var_2_600000'
cloth3d_folder = '/data-net/dataset/cloth3d/test'
results_folder = '/data-local/127_2/results'
out_folder = f'{results_folder}/errors/{name}'
calculate_ours = True
calculate()
scripts/calculate_error_all.py
+8 -6
View File
@@ -10,6 +10,7 @@ parser.add_argument('--config', type=str)
parser.add_argument('--debug', action='store_true')
parser.add_argument('--iteration', type=str, default='200000')
parser.add_argument('-g', '--gpu', type=str, default='0')
parser.add_argument('-f', '--final', action='store_true')
args = vars(parser.parse_args())
if args['debug']:
@@ -68,7 +69,7 @@ def get_normal_loss(pred, gt, mask):
return n_loss.cpu().numpy()
def get_collision_loss(pred, body, mask):
c_loss, c_ratio = collision_loss(torch.from_numpy(pred).unsqueeze(0).to(device), torch.from_numpy(body).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device))
c_loss, c_ratio = collision_loss(torch.from_numpy(pred).unsqueeze(0).to(device), torch.from_numpy(body).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device), thr=0.0005)
return c_loss.cpu().numpy()[0], c_ratio.cpu().numpy()
@@ -115,10 +116,15 @@ def calculate():
errors_edge.append((str(frame), get_edge_loss(pred, gt, new_mask)))
else:
for i in range(4):
levels = [0,1,2,3] if not args['final'] else [3]
for i in levels:
_pred = pred.item()[f'lvl{i}']
c_loss, c_ratio = get_collision_loss(_pred, body, new_mask)
errors_collision.append((f'{frame}_lvl{i}', c_loss))
errors_collision_ratio.append((f'{frame}_lvl{i}', c_ratio))
errors_normal.append((f'{frame}_lvl{i}', get_normal_loss(_pred, gt, new_mask)))
errors_edge.append((f'{frame}_lvl{i}', get_edge_loss(_pred, gt, new_mask)))
# faces = fix_data['faces']
# body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32)
@@ -143,10 +149,6 @@ def calculate():
# fig.show()
errors_collision.append((f'{frame}_lvl{i}', c_loss))
errors_collision_ratio.append((f'{frame}_lvl{i}', c_ratio))
errors_normal.append((f'{frame}_lvl{i}', get_normal_loss(_pred, gt, new_mask)))
errors_edge.append((f'{frame}_lvl{i}', get_edge_loss(_pred, gt, new_mask)))
# save errors to a csv file
scripts/cvae.py
+50 -2
View File
@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
args = vars(parser.parse_args())
if args['debug']:
from configs.iccvw.abl_all_5_var import config
from configs.iccvw.abl_all_5 import config
else:
config_module = importlib.import_module(args['config'])
config = config_module.config
@@ -45,7 +45,7 @@ from data.dataset import Dataset
from models.cond_encoder import AutoEncoder
from utils.checkpoints import CheckpointSaver, load_model
from utils.utils import InfiniteDataLoader as DataLoader
from utils.utils import text_to_html
from utils.utils import text_to_html, save_stl, create_temp_frames
from utils.postprocess import fix_collision
from losses.losses import collision_loss, normal_loss
from pytorch_msssim import MS_SSIM
@@ -352,7 +352,52 @@ def sample():
break
def generate():
from data.utils import Frame
load_model(model, f'checkpoints/iccvw/{config["name"]}')
frame = '01852_70_Tshirt'
# frame = '03849_27_Dress' # 03849_27_Dress 00054_1_Dress
frame = Frame(frame_str=frame)
create_temp_frames(config['train_dir'], frame)
dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=config['dataset_in_memory'])
train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True, )
batch = next(iter(train_dataloader))
if config['use_skinning']:
mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
else:
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
# with record_function("forward"):
model.eval()
with torch.inference_mode():
pred, _, _ = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']])
mesh_pred = uv2mesh(pred)
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])])
# shape_offsets = torch.stack([torch.matmul(shape[i, :10], smpl.shapedirs[gender[i]][:, 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_pred_posed = mesh_pred
fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
mask = fix_data['mask']
faces = fix_data['faces']
pred = mesh_pred_posed.cpu().numpy()[0, mask, :]
body = batch['mesh_body_posed'].cpu().numpy()[0]
body_faces = smpl.f
save_stl(Path('temp') / 'baseline' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_baseline_mesh.stl', pred, faces)
save_stl(Path('temp') / 'baseline' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_baseline_body.stl', body, body_faces)
pass
def train():
@@ -410,6 +455,9 @@ if __name__ == '__main__':
model = CVAE(config)
generate()
exit()
sample()
exit()
scripts/cvae_pyramid.py
+364 -26
View File
@@ -11,7 +11,8 @@ parser.add_argument('--debug', action='store_true')
args = vars(parser.parse_args())
if args['debug']:
from configs.iccvw.pyr_64_128_256_512 import config
# from configs.iccvw.pyr_64_128_256_512 import config
from configs.iccvw.pyr_64_128_256_512_e2e import config
else:
config_module = importlib.import_module(args['config'])
config = config_module.config
@@ -24,7 +25,7 @@ if args['debug']:
import time
import argparse
import wandb
import tqdm
from tqdm import tqdm
import random
import numpy as np
import plotly.graph_objects as go
@@ -501,6 +502,32 @@ def speed_step(levels_input, conds):
return model.sample([levels_input, conds], training=False)
from scipy.spatial import KDTree
import trimesh
def connect_unconnected_vertices(mesh):
# Find unconnected vertices
unconnected_verts = np.where(np.array(mesh.vertex_degree) == 0)[0]
# Create KD-tree from mesh vertices
tree = KDTree(mesh.vertices)
new_faces = mesh.faces.copy()
# Connect unconnected vertices to their two closest neighbors
for vert in unconnected_verts:
# Find the two closest vertices
_, closest_verts = tree.query(mesh.vertices[vert], k=3)
# Create a new face with the unconnected vertex and its two closest neighbors
new_face = [vert, closest_verts[1], closest_verts[2]]
new_faces = np.vstack((new_faces, new_face))
mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=new_faces, process=False)
return mesh
def sample():
from data.utils import Frame
from utils.utils import save_stl
@@ -515,11 +542,10 @@ def sample():
load_model(model, f'checkpoints/iccvw/{config["name"]}')
count = 0
frames = ['01852_70_Tshirt']
frames = ['06461_237_Trousers'] # 03528_132_Tshirt 01852_70_Tshirt 06461_237_Trousers 01893_173_Jumpsuit
frames = ['01171_213_Dress'] # 01039_118_Dress 01171_213_Dress
frames = [Frame(frame_str=frame) for frame in frames]
for frame in frames:
create_temp_frames(config['train_dir'], frame)
dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[64, 128, 256, 512], is_cape=config['is_cape'])
@@ -532,6 +558,8 @@ def sample():
uv_normals = batch['uv_normals']
batch_size = uv_static.size(0)
body = batch['mesh_body_posed'][0].cpu().numpy()
body_faces = smpl.f
for i in range(batch_size):
@@ -541,8 +569,11 @@ def sample():
mask = fix_data['mask']
faces = fix_data['faces']
os.makedirs(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}', exist_ok=True)
gt = batch['mesh_posed'][i].cpu().numpy()[mask]
save_stl(Path('temp') / 'sampling_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt.stl', gt, faces)
save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt.stl', gt, faces)
save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt_body.stl', body, body_faces)
uv_static_tiled = uv_static[i][None,:,:,:].repeat(4, 1, 1, 1)
uv_body_posed_tiled = uv_body_posed[i][None,:,:,:].repeat(4, 1, 1, 1)
@@ -552,6 +583,8 @@ def sample():
body_posed = batch['mesh_body_posed'][i][None,:,:].repeat(4, 1, 1)
mask_tiled = torch.from_numpy(mask[None,:]).to(body_posed.device).repeat(4, 1)
noise = []
for _ in range(4):
with torch.inference_mode():
pred = model.sample(uv_static_tiled, uv_body_posed_tiled, uv_normals_tiled, uv_mask_tiled)
@@ -570,9 +603,12 @@ def sample():
for k in range(4):
_mesh_posed_fixed = mesh_posed_fixed[k, mask, :]
m = tm.Trimesh(_mesh_posed_fixed, faces, process=False, validate=False)
m = connect_unconnected_vertices(m)
tm.smoothing.filter_humphrey(m, iterations=2)
verts = m.vertices
save_stl(Path('temp') / 'sampling_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}.stl', verts, faces)
save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}.stl', verts, faces)
save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}_body.stl', body, body_faces)
counter += 1
pass
@@ -582,6 +618,265 @@ def sample():
break
def sample_pro():
from data.utils import Frame
from utils.utils import save_stl
from pathlib import Path
import trimesh as tm
def create_temp_frames(folder: str, frame: Frame):
with open(Path(folder) / 'frames_temp.txt', 'w') as f:
f.writelines([f'{frame}\n' for _ in range(4)])
model.eval()
load_model(model, f'checkpoints/iccvw/{config["name"]}')
count = 0
# frames = ['06461_237_Trousers'] # 03528_132_Tshirt 01852_70_Tshirt 06461_237_Trousers 01893_173_Jumpsuit
frames = ['01852_70_Tshirt'] # 01039_118_Dress 01171_213_Dress 01852_70_Tshirt
frames = [Frame(frame_str=frame) for frame in frames]
for frame in frames:
create_temp_frames(config['train_dir'], frame)
dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[64, 128, 256, 512], is_cape=config['is_cape'])
dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
batch = next(iter(dataloader))
uv_static = batch['uv_static']
uv_body_posed = batch['uv_body_posed']
uv_mask = { 512: batch['uv_mask'], 256: batch['uv_mask_256'], 128: batch['uv_mask_128'], 64: batch['uv_mask_64'] }
uv_normals = batch['uv_normals']
batch_size = uv_static.size(0)
body = batch['mesh_body_posed'][0].cpu().numpy()
body_faces = smpl.f
i=0
frame = Frame(batch['id'][i], batch['frame_nr'][i], batch['garment'][i])
fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
mask = fix_data['mask']
faces = fix_data['faces']
os.makedirs(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}', exist_ok=True)
gt = batch['mesh_posed'][i].cpu().numpy()[mask]
save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt.stl', gt, faces)
save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt_body.stl', body, body_faces)
uv_static_tiled = uv_static[i][None,:,:,:].repeat(4, 1, 1, 1)
uv_body_posed_tiled = uv_body_posed[i][None,:,:,:].repeat(4, 1, 1, 1)
uv_mask_tiled = { 512: uv_mask[512][i][None,:,:].repeat(4, 1, 1), 256: uv_mask[256][i][None,:,:].repeat(4, 1, 1), 128: uv_mask[128][i][None,:,:].repeat(4, 1, 1), 64: uv_mask[64][i][None,:,:].repeat(4, 1, 1)}
uv_normals_tiled = uv_normals[i][None,:,:,:].repeat(4, 1, 1, 1)
pose_shape = batch['pose_shape'][i][None,:].repeat(4, 1)
body_posed = batch['mesh_body_posed'][i][None,:,:].repeat(4, 1, 1)
mask_tiled = torch.from_numpy(mask[None,:]).to(body_posed.device).repeat(4, 1)
out = torch.zeros((uv_static_tiled.shape[0], 3, 64, 32)).to(uv_static_tiled.device)
for l in range(4):
with torch.inference_mode():
for b in range(4):
pred_level = model.sample_level(uv_static_tiled, uv_body_posed_tiled, uv_normals_tiled, uv_mask_tiled, out, l)
pred = pred_level
mesh_pred = uv2mesh(pred)
shape = 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(pose_shape[:, :72], pose_shape[:, 72:], mesh_pred_shaped)[0]
mesh_posed_fixed = fix_collision(mesh_pred_posed, body_posed, mask_tiled).cpu().numpy()
for k in range(4):
_mesh_posed_fixed = mesh_posed_fixed[k, mask, :]
m = tm.Trimesh(_mesh_posed_fixed, faces, process=False, validate=False)
m = connect_unconnected_vertices(m)
tm.smoothing.filter_humphrey(m, iterations=2)
verts = m.vertices
save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_lvl_{l}_{b*4+k}.stl', verts, faces)
save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_lvl_{l}_{b*4+k}_body.stl', body, body_faces)
#replace the contentent of out with its own first element tiled
if b == 3:
out = pred_level[3,:,:,:].repeat(4, 1, 1, 1)
pass
# if count > 100:
# break
def pred_pro():
from data.utils import Frame
from utils.utils import save_stl
from pathlib import Path
import trimesh as tm
from typing import List
from models.cvae_norm_model import CVAENorm
from data.downscale import fill
from torchvision.transforms import Resize
from torchvision.transforms.functional import InterpolationMode
def create_temp_frames(folder: str, frames: List[Frame]):
with open(Path(folder) / 'frames_temp.txt', 'w') as f:
f.writelines([f'{frame}\n' for frame in frames])
model.eval()
load_model(model, f'checkpoints/iccvw/{config["name"]}')
count = 0
from configs.iccvw.pyr_64_128_256_512_skirt_e2e import config as config_skirt
model_skirt: Pyramid = Pyramid(config_skirt)
load_model(model_skirt, f'checkpoints/iccvw/{config_skirt["name"]}')
models = {'non-skirt': model, 'skirt': model_skirt}
from configs.iccvw.cvae_norm_5 import config as config_norm
model_norm = CVAENorm(config_norm)
model_norm.eval()
load_model(model_norm, f'checkpoints/iccvw/{config_norm["name"]}')
uv2mesh_skirt = UV2Mesh(res=256, is_skirt=config_skirt['is_skirt'], is_smpl_orig=config['is_cape'])
uv2meshes = {'non-skirt': uv2mesh, 'skirt': uv2mesh_skirt}
# '01320_30_Dress'
frames = ['01852_70_Tshirt', '04717_276_Trousers', '07532_152_Tshirt', '02792_282_Jumpsuit'] # 01039_118_Dress 01171_213_Dress
frames = [Frame(frame_str=frame) for frame in frames]
create_temp_frames(config['train_dir'], frames)
dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[64, 128, 256, 512], is_cape=config['is_cape'])
dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False, )
batch_orig = next(iter(dataloader))
batch_size = 4
body_faces = smpl.f
masks = []
faces_orig = []
for frame in frames:
fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
masks.append(fix_data['mask'])
faces_orig.append(fix_data['faces'])
masks = np.stack(masks, axis=0)
# faces = np.stack(faces, axis=0)
resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BICUBIC).to('cuda') for res in [64, 128, 256, 512] }
mode = 'body'
for i in range(batch_size):
batch = batch_orig.copy()
uv_static = batch['uv_static']
uv_body_posed = batch['uv_body_posed']
uv_mask = { 512: batch['uv_mask'], 256: batch['uv_mask_256'], 128: batch['uv_mask_128'], 64: batch['uv_mask_64'] }
uv_normals = batch['uv_normals']
counter = 0
frame = Frame(batch['id'][i], batch['frame_nr'][i], batch['garment'][i])
# fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
# mask = fix_data['mask']
# faces = fix_data['faces']
type = 'skirt' if frame.garment in ['Skirt', 'Dress'] else 'non-skirt'
os.makedirs(Path('temp') / 'pred_pro' / mode / f'{frame.id}_{frame.frame_nr}_{frame.garment}', exist_ok=True)
# gt = batch['mesh_posed'][i].cpu().numpy()[mask]
# save_stl(Path('temp') / 'pred_pro' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt.stl', gt, faces)
# save_stl(Path('temp') / 'pred_pro' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt_body.stl', body, body_faces)
uv_static_tiled = uv_static[i][None,:,:,:].repeat(4, 1, 1, 1)
uv_body_posed_tiled = uv_body_posed[i][None,:,:,:].repeat(4, 1, 1, 1)
uv_mask_tiled = { 512: uv_mask[512][i][None,:,:].repeat(4, 1, 1), 256: uv_mask[256][i][None,:,:].repeat(4, 1, 1), 128: uv_mask[128][i][None,:,:].repeat(4, 1, 1), 64: uv_mask[64][i][None,:,:].repeat(4, 1, 1)}
uv_normals_tiled = uv_normals[i][None,:,:,:].repeat(4, 1, 1, 1)
pose_shape = batch['pose_shape'][i][None,:].repeat(4, 1)
body_posed = batch['mesh_body_posed'][i][None,:,:].repeat(4, 1, 1)
# masks_tiled = torch.from_numpy(masks[i][None,:]).to(body_posed.device).repeat(4, 1)
if mode == 'static':
batch['uv_body_posed'] = uv_body_posed_tiled
batch['mesh_body_posed'] = body_posed
batch['pose_shape'] = pose_shape
masks_tiled = masks
for res in [64, 128, 256]:
batch[f'uv_body_posed_{res}'] = batch[f'uv_body_posed_{res}'][i][None,:,:,:].repeat(4, 1, 1, 1)
# orig_img = uv_body_posed_tiled.clone()
# downscaled = {64: [], 128: [], 256: []}
# for i in range(4):
# filled = fill(orig_img[i].cpu().numpy().transpose(1,2,0), uv_mask[512][i].cpu().numpy())
# filled = torch.from_numpy(filled.transpose(2,0,1)).to(orig_img.device)
# for res in [64, 128, 256]:
# downscaled[res].append(resizers[res](filled) * uv_mask[res][i])
# for res in [64, 128, 256]:
# batch[f'uv_body_posed_{res}'] = torch.stack(downscaled[res], dim=0)
faces = faces_orig
elif mode == 'body':
batch['uv_static'] = uv_static_tiled
for res in [64, 128, 256]:
batch[f'uv_static_{res}'] = batch[f'uv_static_{res}'][i][None,:,:,:].repeat(4, 1, 1, 1)
batch['mesh_mask'] = batch['mesh_mask'][i][None,:].repeat(4, 1)
for key in uv_mask_tiled.keys():
batch[f'uv_mask_{key}'] = uv_mask_tiled[key]
batch['uv_mask'] = uv_mask_tiled[512]
masks_tiled = torch.from_numpy(masks[i][None,:]).to(body_posed.device).repeat(4, 1).cpu().numpy()
faces = [faces_orig[i], faces_orig[i], faces_orig[i], faces_orig[i]]
else:
raise Exception('mode not supported')
body_posed = batch['mesh_body_posed']
pose_shape = batch['pose_shape']
with torch.inference_mode():
pred_norm = model_norm.sample(batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']) * 1.4
batch['uv_normals'] = pred_norm
orig_img = pred_norm.clone()
downscaled = {64: [], 128: [], 256: []}
for i in range(4):
filled = fill(orig_img[i].cpu().numpy().transpose(1,2,0), uv_mask[512][i].cpu().numpy())
filled = torch.from_numpy(filled.transpose(2,0,1)).to(orig_img.device)
for res in [64, 128, 256]:
downscaled[res].append(resizers[res](filled) * uv_mask[res][i])
for res in [64, 128, 256]:
batch[f'uv_normals_{res}'] = torch.stack(downscaled[res], dim=0)
levels_input, conds = group_inputs_for_levels(batch, config)
with torch.inference_mode():
out_levels = models[type]([levels_input, conds])
pred, _, _ = out_levels[-1]
mesh_pred = uv2meshes[type](pred)
shape = 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(pose_shape[:, :72], pose_shape[:, 72:], mesh_pred_shaped)[0]
# mesh_posed_fixed = fix_collision(mesh_pred_posed, body_posed, torch.from_numpy(masks).to(body_posed.device)).cpu().numpy()
mesh_posed_fixed = mesh_pred_posed.cpu().numpy()
for k in range(4):
_mesh_posed_fixed = mesh_posed_fixed[k, masks_tiled[k], :]
m = tm.Trimesh(_mesh_posed_fixed, faces[k], process=False, validate=False)
m = connect_unconnected_vertices(m)
tm.smoothing.filter_humphrey(m, iterations=2)
verts = m.vertices
save_stl(Path('temp') / 'pred_pro' / mode / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}.stl', verts, faces[k])
save_stl(Path('temp') / 'pred_pro' / mode / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}_body.stl', batch['mesh_body_posed'][k].cpu().numpy(), body_faces)
counter += 1
pass
def train():
global g_balancing_coeff
global step
@@ -620,9 +915,44 @@ def train():
step += 1
def speed_test():
from data.utils import Frame
from utils.utils import save_stl
from pathlib import Path
import trimesh as tm
def create_temp_frames(folder: str, frame: Frame):
with open(Path(folder) / 'frames_temp.txt', 'w') as f:
f.writelines([f'{frame}\n' for _ in range(4)])
model.eval()
load_model(model, f'checkpoints/iccvw/{config["name"]}')
frames = ['01852_70_Tshirt']
frames = [Frame(frame_str=frame) for frame in frames]
create_temp_frames(config['train_dir'], frames[0])
dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[64, 128, 256, 512], is_cape=config['is_cape'])
dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
batch = next(iter(dataloader))
uv_static = batch['uv_static']
uv_body_posed = batch['uv_body_posed']
uv_mask = { 512: batch['uv_mask'], 256: batch['uv_mask_256'], 128: batch['uv_mask_128'], 64: batch['uv_mask_64'] }
uv_normals = batch['uv_normals']
for i in tqdm(range(1000)):
with torch.inference_mode():
if i ==5:
start = time.time()
model.sample(uv_static, uv_body_posed, uv_normals, uv_mask)
print(f'per batch: {(time.time() - start) / 1000}')
if __name__ == '__main__':
if not args['debug']:
init_wandb()
# if not args['debug']:
# init_wandb()
use_amp = True
smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
@@ -640,7 +970,14 @@ if __name__ == '__main__':
model: Pyramid = Pyramid(config)
sample()
# count number of parameters
total_params = sum(p.numel() for p in model.parameters())
# speed_test()
# exit()
# sample_pro()
pred_pro()
exit()
optimizer = optim.Adam(model.parameters(), lr=config['lr'])
@@ -652,23 +989,23 @@ if __name__ == '__main__':
reg_coeff = 1
step = 0
if config['pretrained_model'] != '':
if 'ignore_saved_optimizer' in config.keys() and config['ignore_saved_optimizer']:
step = load_model(model, config['pretrained_model'])
else:
step = load_model(model, config['pretrained_model'], optimizer=optimizer, iteration=config['pretrained_model_iteration'])
# if config['pretrained_model'] != '':
# if 'ignore_saved_optimizer' in config.keys() and config['ignore_saved_optimizer']:
# step = load_model(model, config['pretrained_model'])
# else:
# step = load_model(model, config['pretrained_model'], optimizer=optimizer, iteration=config['pretrained_model_iteration'])
if len(config['pyramid']['resolutions']) == 1:
for param in model.static_enc.parameters(): param.requires_grad = False
for param in model.body_enc.parameters(): param.requires_grad = False
if config['use_normals']:
for param in model.norm_enc.parameters(): param.requires_grad = False
else:
load_model(model, config['pyramid']['pretrained_pyramid'], iteration=config['pyramid']['pretrained_pyramid_iteration'])
for param in model.static_enc.parameters(): param.requires_grad = False
for param in model.body_enc.parameters(): param.requires_grad = False
if config['use_normals']:
for param in model.norm_enc.parameters(): param.requires_grad = False
# if len(config['pyramid']['resolutions']) == 1:
# for param in model.static_enc.parameters(): param.requires_grad = False
# for param in model.body_enc.parameters(): param.requires_grad = False
# if config['use_normals']:
# for param in model.norm_enc.parameters(): param.requires_grad = False
# else:
# load_model(model, config['pyramid']['pretrained_pyramid'], iteration=config['pyramid']['pretrained_pyramid_iteration'])
# for param in model.static_enc.parameters(): param.requires_grad = False
# for param in model.body_enc.parameters(): param.requires_grad = False
# if config['use_normals']:
# for param in model.norm_enc.parameters(): param.requires_grad = False
# load_model(model.levels[0].upscale_enhancer, config['pyramid']['levels'][0]['upscale_enhancer'])
# model.levels[0].upscale_enhancer.trainable = False
@@ -691,4 +1028,5 @@ if __name__ == '__main__':
# model = torch.compile(model)
wandb.watch(model, log='all', log_freq=10000, log_graph=True)
model.eval()
train()
scripts/cvae_pyramid_e2e.py
+4 -5
View File
@@ -547,7 +547,7 @@ def train():
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 } })
to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step, 'balancing_coeff': g_balancing_coeff } })
if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid')
@@ -561,10 +561,9 @@ def train():
valid_accumulator.reset()
print(step)
# if train_loss['train/euclid_error'] < config['balancing_activation_limit']:
# g_balancing_coeff += 1 / 50000
g_balancing_coeff += 1 / 100000
if step % 5000 == 0 and not args['debug']:
if step % 10000 == 0 and not args['debug']:
ckpt_saver(model, step, optimizer, None)
step += 1
@@ -593,7 +592,7 @@ if __name__ == '__main__':
# optimizer = optim.Adam(model.parameters(), lr=config['lr'])
optimizer = optim.Adam(model.parameters(), lr=config['lr'])
_MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=50)
ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=20)
scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
resizer = Resize(size=(512, 256), interpolation=InterpolationMode.BILINEAR).to(config['device'])
scripts/generate_results.py
+1 -1
View File
@@ -14,7 +14,7 @@ parser.add_argument('-l', '--level', type=int, default=-1)
args = vars(parser.parse_args())
if args['debug']:
from configs.iccvw.pyr_64_128_cape_f_2 import config_1
from configs.iccvw.pyr_64_128_256_512_e2e_var import config_1
else:
config_module = importlib.import_module(args['config'])
config_1 = config_module.config
scripts/generate_results_all.py
+3 -1
View File
@@ -10,6 +10,7 @@ parser.add_argument('--config', type=str)
parser.add_argument('--iteration', type=str, default='')
parser.add_argument('--debug', action='store_true')
parser.add_argument('-g', '--gpu', type=str, default='1')
parser.add_argument('-fr', '--frames_file', type=str, default='')
args = vars(parser.parse_args())
if args['debug']:
@@ -117,7 +118,8 @@ if __name__ == '__main__':
# remove all values from config['signature] that start with mesh_
config_1['signature'] = set([x for x in config_1['signature'] if not x.startswith('mesh_')])
dataset_test = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file=config_1['test_frames'], use_cache=config_1['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
frames_file = args['frames_file'] if args['frames_file'] != '' else config_1['test_frames']
dataset_test = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file=frames_file, use_cache=config_1['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
# dataset_test = Dataset(config['test_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
test_dataloader = DataLoader(dataset_test, batch_size=4, drop_last=True, shuffle=True, )
scripts/get_error.py
+20 -12
View File
@@ -8,13 +8,16 @@ parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
parser.add_argument('--config', type=str)
parser.add_argument('--debug', action='store_true')
parser.add_argument('--iteration', type=str, default='190000')
parser.add_argument('--iteration', type=str, default='')
parser.add_argument('-g', '--gpu', type=str)
parser.add_argument('-n', '--normal', action='store_true')
parser.add_argument('-e', '--edge', action='store_true')
parser.add_argument('-cr', '--collision_ratio', action='store_true')
parser.add_argument('-a', '--all', action='store_true')
args = vars(parser.parse_args())
if args['debug']:
from configs.iccvw.pyr_64_128_256_512_cape_f_e2e import config as config_1
from configs.iccvw.pyr_64_128_256_512_e2e import config as config_1
else:
config_module = importlib.import_module(args['config'])
config_1 = config_module.config
@@ -32,7 +35,7 @@ from data.utils import Frame
def calculate():
def calculate(type=''):
all_levels = args['all']
if all_levels:
errors = {0: dict(), 1: dict(), 2: dict(), 3: dict()}
@@ -41,8 +44,7 @@ def calculate():
errors = dict()
errors_mean = dict()
with open(f'{results_folder}/errors/{name}.csv', 'r') as f:
with open(f'{results_folder}/errors/{name}{type}.csv', 'r') as f:
lines = f.readlines()
for line in lines:
@@ -50,11 +52,11 @@ def calculate():
category = line[2]
if all_levels:
level = int(line[3][-1])
if category not in errors:
if category not in errors[level].keys():
errors[level][category] = []
errors[level][category].append(float(line[4]))
else:
if category not in errors:
if category not in errors.keys():
errors[category] = []
errors[category].append(float(line[3]))
@@ -62,17 +64,18 @@ def calculate():
for i in range(4):
for category in errors[i]:
errors_mean[i][category] = round(np.mean(errors[i][category]), 5)
print(f'Level {i}: {errors_mean[i]}')
print(f'Level {i} error{type}: {errors_mean[i]}')
mean = np.mean(list(errors_mean[i].values()))
print(f'Final error level {i}: {mean}')
print(f'Final error{type} level {i}: {mean}')
pass
else:
for category in errors:
errors_mean[category] = round(np.mean(errors[category]), 5)
print(errors_mean)
mean = np.mean(list(errors_mean.values()))
print(f'Final error: {mean}')
mean = round(np.mean(list(errors_mean.values())), 5)
print(f'Final error{type}: {mean}')
@@ -98,5 +101,10 @@ if __name__ == '__main__':
name = f'{name}_{args["iteration"]}'
out_folder = f'{results_folder}/{name}'
if args['normal']:
calculate('_normal')
if args['edge']:
calculate('_edge')
if args['collision_ratio']:
calculate('_collision_ratio')
calculate()
scripts/get_error_hood.py
+28 -46
View File
@@ -2,30 +2,9 @@ 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')
parser.add_argument('--iteration', type=str, default='')
parser.add_argument('-g', '--gpu', type=str)
parser.add_argument('-n', '--normal', action='store_true')
parser.add_argument('-e', '--edge', action='store_true')
parser.add_argument('-cr', '--collision_ratio', action='store_true')
parser.add_argument('-a', '--all', action='store_true')
args = vars(parser.parse_args())
if args['debug']:
from configs.iccvw.pyr_64_128_256_512_cape_f_e2e import config as config_1
else:
config_module = importlib.import_module(args['config'])
config_1 = config_module.config
import time
import argparse
import tqdm
from tqdm import tqdm
import random
import numpy as np
import plotly.graph_objects as go
@@ -35,50 +14,53 @@ from data.utils import Frame
def calculate(type=''):
def calculate(type='chamfer'):
errors = dict()
errors_mean = dict()
with open(f'{folder}/{hood_csv}', 'r') as f:
lines_hood = f.readlines()
with open(f'{folder}/{ours_csv}', 'r') as f:
lines_ours = f.readlines()
files_hood = sorted(os.listdir(f'{folder}/{run}_hood/{type}'))
errors_hood = [np.load(f'{folder}/{run}_hood/{type}/{file}') for file in tqdm(files_hood)]
files_ours = sorted(os.listdir(f'{folder}/hood_3_var_ours/{type}'))
errors_ours = [np.load(f'{folder}/hood_3_var_ours/{type}/{file}') for file in tqdm(files_ours)]
for line in lines_hood:
line = line.split(',')
category = line[2]
for error, frame in zip(errors_hood, files_hood):
frame_id, frame_nr, frame_garment = frame.split('_')
category = frame_garment
if category not in errors:
errors[category] = []
errors[category].append(float(line[3]))
errors[category].append(error)
for category in errors:
errors_mean[category] = round(np.mean(errors[category]), 5)
print(errors_mean)
mean = round(np.mean(list(errors_mean.values())), 5)
print(f'Final HOOD error{type}: {mean}')
print(f'Final HOOD error_{type}: {mean}')
for line in lines_ours:
line = line.split(',')
category = line[2]
# errors = dict()
# errors_mean = dict()
if category not in errors:
errors[category] = []
errors[category].append(float(line[3]))
# for error, frame in zip(errors_ours, files_ours):
# frame_id, frame_nr, frame_garment = frame.split('_')
# category = frame_garment
for category in errors:
errors_mean[category] = round(np.mean(errors[category]), 5)
print(errors_mean)
# if category not in errors:
# errors[category] = []
# errors[category].append(error)
mean = round(np.mean(list(errors_mean.values())), 5)
print(f'Final OURS error{type}: {mean}')
# for category in errors:
# errors_mean[category] = round(np.mean(errors[category]), 5)
# print(errors_mean)
# mean = round(np.mean(list(errors_mean.values())), 5)
# print(f'Final OURS error_{type}: {mean}')
if __name__ == '__main__':
folder = '/data-local/127_2/results/hood'
hood_csv = 'hood_2_hood.csv'
ours_csv = 'hood_2_ours.csv'
run = 'hood_3'
calculate()
for type in ['chamfer', 'area', 'collision_ratio']: # ['chamfer', 'collision', 'area', 'collision_ratio']:
calculate(type)
scripts/plot_ablation.py
+40 -25
View File
@@ -25,7 +25,7 @@ import importlib
from data.utils import Frame
from data.dataset import read_frames_file
from utils.utils import save_stl
from utils.utils import save_stl, create_temp_frames
from utils.checkpoints import CheckpointSaver, load_model
from utils.postprocess import fix_collision
from losses.losses import collision_loss, normal_loss, kl_loss
@@ -70,25 +70,21 @@ def step_pyramid(model, batch, config):
return mesh_pred_posed
def create_temp_frames(folder: str, frame: Frame):
with open(Path(folder) / 'frames_temp.txt', 'w') as f:
f.writelines([f'{frame}\n' for _ in range(4)])
if __name__ == '__main__':
results_folder = Path('/data-local/127_2/results')
out_folder = Path('temp/ablation')
out_folder = Path('temp/supplementary')
os.makedirs(out_folder, exist_ok=True)
smpl = SMPL(orig_smpl=config_1['is_cape'], path='smpl', device=config_1['device'])
dataset_resolutions = []
dataset_resolutions = [64, 128, 256, 512]
experiments = {
'final': { 'nonskirt': ('pyr_64_128_256_512', 100000), 'skirt': ('pyr_64_128_256_512_skirt', 200000) },
'e2e': { 'nonskirt': ('pyr_64_128_256_512_e2e', 195000), 'skirt': ('pyr_64_128_256_512_skirt_e2e', 200000) },
'skin': { 'nonskirt': ('pyr_64_128_256_512_skin', 225000), 'skirt': ('pyr_64_128_256_512_skin_skirt', 200000) },
'norm': { 'nonskirt': ('pyr_64_128_256_512_norm', 200000), 'skirt': ('pyr_64_128_256_512_norm_skirt', 200000) },
'reg': { 'nonskirt': ('pyr_64_128_256_512_reg', 200000), 'skirt': ('pyr_64_128_256_512_reg_skirt', 200000) },
# 'inc': { 'nonskirt': ('pyr_64_128_256_512', 100000), 'skirt': ('pyr_64_128_256_512_skirt', 200000) },
'e2e': { 'nonskirt': ('pyr_64_128_256_512_e2e', 250000), 'skirt': ('pyr_64_128_256_512_skirt_e2e', 200000) },
# 'skin': { 'nonskirt': ('pyr_64_128_256_512_skin', 225000), 'skirt': ('pyr_64_128_256_512_skin_skirt', 200000) },
# 'norm': { 'nonskirt': ('pyr_64_128_256_512_norm', 200000), 'skirt': ('pyr_64_128_256_512_norm_skirt', 200000) },
# 'reg': { 'nonskirt': ('pyr_64_128_256_512_reg', 200000), 'skirt': ('pyr_64_128_256_512_reg_skirt', 200000) },
}
configs = dict()
@@ -111,8 +107,23 @@ if __name__ == '__main__':
# frames = ['00054_1_Dress', '00005_1_Jumpsuit']
frames = ['03849_27_Dress', '00054_1_Dress', '01852_70_Tshirt']
frames = [Frame(frame_str=frame) for frame in frames]
# frames = ['03849_27_Dress', '00054_1_Dress', '01852_70_Tshirt']
# frames = ['01852_70_Tshirt'] # , '03849_27_Dress', '00054_1_Dress'
# frames = [Frame(frame_str=frame) for frame in frames]
frames_file = configs['e2e']['nonskirt']['train_frames']
with open(os.path.join(config_1['train_dir'], frames_file), 'r') as f:
frames = f.readlines()
frames_1 = [Frame(frames.split()[0], int(frames.split()[1]), frames.split()[2]) for frames in frames]
frames_file = configs['e2e']['skirt']['train_frames']
with open(os.path.join(config_1['train_dir'], frames_file), 'r') as f:
frames = f.readlines()
frames_2 = [Frame(frames.split()[0], int(frames.split()[1]), frames.split()[2]) for frames in frames]
frames = frames_1 + frames_2
random.shuffle(frames)
body_faces = smpl.f
for frame in frames:
@@ -130,8 +141,8 @@ if __name__ == '__main__':
body = batch['mesh_body_posed'][0].cpu().numpy()
gt = batch['mesh_posed'][0].cpu().numpy()
save_stl(out_folder / f'{frame}_gt.stl', gt[mask, :], faces)
save_stl(out_folder / f'{frame}_gt_body.stl', body, body_faces)
# save_stl(out_folder / f'{frame}_gt.stl', gt[mask, :], faces)
# save_stl(out_folder / f'{frame}_gt_body.stl', body, body_faces)
for experiment in experiments.keys():
@@ -139,20 +150,24 @@ if __name__ == '__main__':
mesh_posed_fixed = fix_collision(pred, batch['mesh_body_posed'], batch['mesh_mask']).cpu().numpy()
mesh_posed_fixed = mesh_posed_fixed[0, mask, :]
try:
m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False)
tm.smoothing.filter_humphrey(m, iterations=2)
verts = m.vertices
except:
verts = mesh_posed_fixed
print('failed to apply smoothing')
save_stl(out_folder / f'{frame}_{experiment}_body.stl', body, body_faces)
save_stl(out_folder / f'{frame}_{experiment}_mesh.stl', verts, faces)
# save_stl(out_folder / f'{frame}_{experiment}_body.stl', body, body_faces)
# save_stl(out_folder / f'{frame}_{experiment}_mesh.stl', verts, faces)
# fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
# go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1)
# ])
# fig.update_layout(scene=dict(aspectmode='cube'))
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
# fig.update_layout(title_text=f'{frame} {experiment}')
# fig.show()
fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1)
])
fig.update_layout(scene=dict(aspectmode='cube'))
fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
fig.update_layout(title_text=f'{frame} {experiment}')
fig.show()
pass
pass
scripts/plot_ablation_pyramid.py
+15 -17
View File
@@ -24,7 +24,7 @@ from pathlib import Path
from data.utils import Frame
from data.dataset import read_frames_file
from utils.utils import save_stl
from utils.utils import save_stl, create_temp_frames
from utils.checkpoints import CheckpointSaver, load_model
from utils.postprocess import fix_collision
from losses.losses import collision_loss, normal_loss, kl_loss
@@ -38,7 +38,6 @@ from torch.utils.data import DataLoader
import trimesh as tm
from zipfile import ZipFile
from configs.iccvw.pyr_64_128_256_512_e2e import config as config_1
from configs.iccvw.pyr_64_128_256_512_skirt_e2e import config as config_2
@@ -76,15 +75,13 @@ def step_pyramid(model, batch, config):
return out
def create_temp_frames(folder: str, frame: Frame):
with open(Path(folder) / 'frames_temp.txt', 'w') as f:
f.writelines([f'{frame}\n' for _ in range(4)])
if __name__ == '__main__':
results_folder = Path('/data-local/127_2/results')
out_folder = Path('temp/pyramid_ablation')
iteration = {'nonskirt': 190000, 'skirt': 200000 }
iteration = {'nonskirt': 250000, 'skirt': 200000 }
smpl = SMPL(orig_smpl=config_1['is_cape'], path='smpl', device=config_1['device'])
dataset_resolutions = []
@@ -99,8 +96,9 @@ if __name__ == '__main__':
model['nonskirt'].eval()
model['skirt'].eval()
frames = ['00054_1_Dress', '01070_13_Jumpsuit', '03849_27_Dress', '06057_152_Dress',
frames = ['01070_13_Jumpsuit', '00054_1_Dress', '03849_27_Dress', '06057_152_Dress',
'03179_274_Dress', '04247_102_Tshirt', '01852_70_Tshirt']
frames = ['00027_1_Dress']
frames_corrupt = ['00300_69_Jumpsuit', '07608_77_Jumpsuit', '06955_117_Dress']
frames = [Frame(frame_str=frame) for frame in frames]
@@ -109,8 +107,8 @@ if __name__ == '__main__':
config = { 'nonskirt': config_1, 'skirt': config_2 }
for frame in tqdm(frames):
create_temp_frames(config_1['train_dir'], frame)
dataset = Dataset(config_1['train_dir'], config_1['signature'], device=config_1['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
create_temp_frames(config_1['test_dir'], frame)
dataset = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
picker = 'skirt' if frame.garment in ['Skirt', 'Dress'] else 'nonskirt'
@@ -136,15 +134,15 @@ if __name__ == '__main__':
tm.smoothing.filter_humphrey(m, iterations=2)
verts = m.vertices
save_stl(out_folder / f'{frame}_lvl{i}_body.stl', body, body_faces)
save_stl(out_folder / f'{frame}_lvl{i}_mesh.stl', verts, faces)
# save_stl(out_folder / f'{frame}_lvl{i}_body.stl', body, body_faces)
# save_stl(out_folder / f'{frame}_lvl{i}_mesh.stl', verts, faces)
# fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
# go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1)
# ])
# fig.update_layout(scene=dict(aspectmode='cube'))
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
# fig.show()
fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1)
])
fig.update_layout(scene=dict(aspectmode='cube'))
fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
fig.show()
pass
pass
scripts/plot_cloth3d.py
+148
View File
@@ -0,0 +1,148 @@
import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
import time
import argparse
import random
import tqdm
import numpy as np
import plotly.graph_objects as go
import plotly.express as ex
from utils.postprocess import fix_collision
import torch
from tqdm import tqdm
from data.utils import Frame
from data.cloth3d_io import readPC2Frame, readPC2, readOBJ
from utils.utils import pickle_load, save_stl
from scipy.spatial import cKDTree
from smpl.smpl_torch import SMPL
from scipy.spatial.transform import Rotation as R
import trimesh as tm
from plotly.subplots import make_subplots
smpl = SMPL(path='smpl', device='cpu')
def deorient(verts, pose, inverse=True):
assert pose.shape[1] == 72
rodr = smpl.rodrigues(torch.tensor(pose, device=smpl.device).reshape(-1, 1, 3)).reshape(verts.shape[0], -1, 3, 3).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=inverse)
return rotated
def readBIN(fname):
if '.' not in fname: fname += '.bin'
elif '.' in fname and not fname.endswith('.bin'):
print("File name extension should be '.bin'")
return
with open(fname, 'rb') as f:
F = np.frombuffer(f.read(), dtype=np.uint16).astype(np.int32)
return F.reshape((-1,3))
def calculate():
garments = sorted(os.listdir('/data-local/127_2/cloth3d_processed_test/mesh_mask'))
garments = [g[:-4] for g in garments]
# random.shuffle(garments)
# garments = [('01909_Jumpsuit', 175), ('00960_Dress', 41), ('01573_Dress', 40), ('00079_Tshirt', 190)]
body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32)
for garment in tqdm(garments):
# garment = garment_data[0]
picker = 'skirt' if garment.split('_')[1] in ['Dress', 'Skirt'] else 'nonskirt'
id, garment_type = garment.split('_')
if not os.path.exists(f'temp/cloth3d/{id}/outfit.pc16'):
continue
if not os.path.exists(os.path.join(cloth3d_folder, id, garment_type + '.pc16')):
continue
pred_cloth3d = readPC2(f'temp/cloth3d/{id}/outfit.pc16', True)['V']
pred_cloth3d_faces = readBIN(f'temp/cloth3d/{id}/outfit.bin')
gt = readPC2(os.path.join(cloth3d_folder, id, garment_type + '.pc16'), True)['V']
mask = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')['mask']
pose = np.load(f'/data-local/127_2/cloth3d_processed_test/pose/{id}.npy')
gt = deorient(gt, pose, True)
# gt = deorient(gt, pose, True)
body = readPC2(f'/data-local/127_2/cloth3d_processed_test/mesh_body_posed_no_orient/{id}.pc16', True)['V']
fix_data = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')
faces = fix_data['faces']
i = np.random.randint(0, len(pred_cloth3d))
# i = garment_data[1]
try:
if calculate_ours:
pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'].item()['lvl3']
except:
print('missed one')
continue
mesh_posed_fixed = fix_collision(torch.from_numpy(pred_ours[None,:,:]).to('cuda'), torch.from_numpy(body[i][None,:,:]).to('cuda'), torch.from_numpy(mask[None,:]).to('cuda')).cpu().numpy()
mesh_posed_fixed = mesh_posed_fixed[0, mask, :]
try:
m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False)
tm.smoothing.filter_humphrey(m, iterations=2)
pred_ours = m.vertices
except:
pred_ours = mesh_posed_fixed
print('failed to apply smoothing')
# fig = go.Figure(data=[go.Mesh3d(x=pred_ours[:,0], y=pred_ours[:,2], z=pred_ours[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='blue')])
# fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
# 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', scene_aspectratio=dict(x=1, y=1, z=1))
# fig.show()
fig = go.Figure(data=[go.Mesh3d(x=pred_cloth3d[i][:,0], y=pred_cloth3d[i][:,1], z=pred_cloth3d[i][:,2], i=pred_cloth3d_faces[:,0], j=pred_cloth3d_faces[:,1], k=pred_cloth3d_faces[:,2], color='blue')])
fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
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', scene_aspectratio=dict(x=1, y=1, z=1))
fig.show()
# fig = go.Figure(data=[go.Mesh3d(x=gt[i][:,0], y=gt[i][:,2], z=gt[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')])
# fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
# 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', scene_aspectratio=dict(x=1, y=1, z=1))
# fig.show()
pass
# save_stl(f'temp/hood/{id}_{i}_{garment_type}_gt.stl', gt[i], pred_hood_faces)
# save_stl(f'temp/hood/{id}_{i}_{garment_type}_gt_body.stl', body[i], body_faces)
# save_stl(f'temp/hood/new/{id}_{i}_{garment_type}_hood.stl', pred_hood[i], pred_hood_faces)
# save_stl(f'temp/hood/new/{id}_{i}_{garment_type}_hood_body.stl', body[i], body_faces)
# save_stl(f'temp/hood/{id}_{i}_{garment_type}_ours.stl', pred_ours, faces)
# save_stl(f'temp/hood/{id}_{i}_{garment_type}_ours_body.stl', body[i], body_faces)
if __name__ == '__main__':
hood_folder = '/data-local/127_2/results/hood/all_2_105k'
ours_folder = dict()
ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_200000'
ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_2_250000'
cloth3d_folder = '/data-net/dataset/cloth3d/test'
results_folder = '/data-local/127_2/results'
calculate_ours = True
calculate()
scripts/plot_hood.py
+131
View File
@@ -0,0 +1,131 @@
import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
import time
import argparse
import random
import tqdm
import numpy as np
import plotly.graph_objects as go
import plotly.express as ex
from utils.postprocess import fix_collision
import torch
from tqdm import tqdm
from data.utils import Frame
from data.cloth3d_io import readPC2Frame, readPC2, readOBJ
from utils.utils import pickle_load, save_stl
from scipy.spatial import cKDTree
from smpl.smpl_torch import SMPL
from scipy.spatial.transform import Rotation as R
import trimesh as tm
from plotly.subplots import make_subplots
smpl = SMPL(path='smpl', device='cpu')
def deorient(verts, pose, inverse=True):
assert pose.shape[1] == 72
rodr = smpl.rodrigues(torch.tensor(pose, device=smpl.device).reshape(-1, 1, 3)).reshape(verts.shape[0], -1, 3, 3).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=inverse)
return rotated
def calculate():
# garments = sorted(os.listdir(hood_folder))
# garments = [g[:-4] for g in garments]
# random.shuffle(garments)
garments = [('01909_Jumpsuit', 175), ('00960_Dress', 41), ('01573_Dress', 40), ('00079_Tshirt', 190)]
body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32)
for garment_data in tqdm(garments):
garment = garment_data[0]
picker = 'skirt' if garment.split('_')[1] in ['Dress', 'Skirt'] else 'nonskirt'
id, garment_type = garment.split('_')
pred_hood_data = pickle_load(f'{hood_folder}/{garment}.pkl')
pred_hood, pred_hood_faces = pred_hood_data['verts'], pred_hood_data['faces']
gt = readPC2(os.path.join(cloth3d_folder, id, garment_type + '.pc16'), True)['V'][2:]
mask = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')['mask']
pose = np.load(f'/data-local/127_2/cloth3d_processed_test/pose/{id}.npy')[2:]
gt = deorient(gt, pose, True)
body = readPC2(f'/data-local/127_2/cloth3d_processed_test/mesh_body_posed_no_orient/{id}.pc16', True)['V'][2:]
fix_data = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')
faces = fix_data['faces']
# i = np.random.randint(0, len(pred_hood))
i = garment_data[1]
try:
if calculate_ours:
pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'].item()['lvl3']
except:
print('missed one')
continue
mesh_posed_fixed = fix_collision(torch.from_numpy(pred_ours[None,:,:]).to('cuda'), torch.from_numpy(body[i][None,:,:]).to('cuda'), torch.from_numpy(mask[None,:]).to('cuda')).cpu().numpy()
mesh_posed_fixed = mesh_posed_fixed[0, mask, :]
try:
m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False)
tm.smoothing.filter_humphrey(m, iterations=2)
pred_ours = m.vertices
except:
pred_ours = mesh_posed_fixed
print('failed to apply smoothing')
# fig = go.Figure(data=[go.Mesh3d(x=pred_ours[:,0], y=pred_ours[:,2], z=pred_ours[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='blue')])
# fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
# 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', scene_aspectratio=dict(x=1, y=1, z=1))
# fig.show()
fig = go.Figure(data=[go.Mesh3d(x=pred_hood[i][:,0], y=pred_hood[i][:,2], z=pred_hood[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')])
fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
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', scene_aspectratio=dict(x=1, y=1, z=1))
fig.show()
# fig = go.Figure(data=[go.Mesh3d(x=gt[i][:,0], y=gt[i][:,2], z=gt[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')])
# fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
# 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', scene_aspectratio=dict(x=1, y=1, z=1))
# fig.show()
pass
# save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_gt.stl', gt[i], pred_hood_faces)
# save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_gt_body.stl', body[i], body_faces)
# save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_hood.stl', pred_hood[i], pred_hood_faces)
# save_stl(f'temp/hood/ne2w/{id}_{i}_{garment_type}_hood_body.stl', body[i], body_faces)
# save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_ours.stl', pred_ours, faces)
# save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_ours_body.stl', body[i], body_faces)
if __name__ == '__main__':
hood_folder = '/data-local/127_2/results/hood/all_2_105k'
ours_folder = dict()
ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_200000'
ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_2_250000'
cloth3d_folder = '/data-net/dataset/cloth3d/test'
results_folder = '/data-local/127_2/results'
calculate_ours = True
calculate()
scripts/upscaler_train.py
+8 -5
View File
@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
args = vars(parser.parse_args())
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = '2'
os.environ["CUDA_VISIBLE_DEVICES"] = '3'
import time
@@ -39,10 +39,10 @@ from utils.utils import InfiniteDataLoader as DataLoader
from pytorch_msssim import MS_SSIM
low_res = 64
is_skirt = False
is_cape = True
is_male = True
low_res = 256
is_skirt = True
is_cape = False
is_male = False
skirt_modifier = '_skirt' if is_skirt else ''
cape_modifier = '_cape' if is_cape else ''
gender_modifier = '' if not is_cape else ('_m' if is_male else '_f')
@@ -244,6 +244,9 @@ def train():
step += 1
if step > 100000:
break
if __name__ == '__main__':
if not args['debug']:
utils/postprocess.py
+2
View File
@@ -35,6 +35,8 @@ def fix_collision(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr=.004
F_fix_coll_local = F_fix_coll_orig
else:
raise ValueError("V must have 14475 or 3638 vertices")
if V.shape[0] == 1:
F_fix_coll_local = F_fix_coll_local[:1]
mesh = Meshes(B, F_fix_coll_local)
B_vn = torch.stack(mesh.verts_normals_list())
utils/utils.py
+28 -1
View File
@@ -1,9 +1,11 @@
import os
import sys
import pickle
import numpy as np
from torch.utils.data import DataLoader
from stl import mesh as mesh_stl
from pathlib import Path
from data.utils import Frame
class Accumulator():
@@ -60,3 +62,28 @@ def save_stl(file, mesh, faces):
obj.vectors[i][j] = mesh[f[j],:]
obj.save(file)
def create_temp_frames(folder: str, frame: Frame):
with open(Path(folder) / 'frames_temp.txt', 'w') as f:
f.writelines([f'{frame}\n' for _ in range(4)])
def pickle_load(file):
"""
Load a pickle file.
"""
with open(file, 'rb') as f:
loadout = pickle.load(f)
return loadout
def pickle_dump(loadout, file):
"""
Dump a pickle file. Create the directory if it does not exist.
"""
os.makedirs(os.path.dirname(str(file)), exist_ok=True)
with open(file, 'wb') as f:
pickle.dump(loadout, f)