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

Compare commits

base: luni/cvpr_2023:dbbcd5f14ec3d622d60e77640c2c67928f48dff8
Branches Tags
luni/cvpr_2023:main luni/cvpr_2023:iccv_review luni/cvpr_2023:vto luni/cvpr_2023:extrapolation luni/cvpr_2023:skip_conn
..
compare: luni/cvpr_2023:skip_conn
Branches Tags
luni/cvpr_2023:main luni/cvpr_2023:iccv_review luni/cvpr_2023:vto luni/cvpr_2023:extrapolation luni/cvpr_2023:skip_conn
2 Commits
dbbcd5f14e .. skip_conn
Author SHA1 Message Date
luni 5a1276f045 trying skip connection 2023-05-24 17:46:13 +02:00
luni 53217c6e5d direct training 2023-05-23 22:57:08 +02:00
34 changed files with 865 additions and 4142 deletions
Expand all files Collapse all files
configs/abl_norm.py
-192
View File
@@ -1,192 +0,0 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '0'
config['name'] = 'abl_torch_norm'
config['notes'] = ''
config['tags'] = ['iccv', 'abl_norm', 'torch']
config['script'] = 'scripts/cvae.py'
config['pretrained_model'] = ''
config['use_skinning'] = True # and deshaping
config['remove_global_orientation'] = True
config['use_normals'] = True
config['use_extra_losses'] = True # edge, collision, 3d mesh
config['use_offsets'] = True
config['use_norm_loss'] = True
ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_normals'], config['use_extra_losses'], config['use_offsets']].count(False)
assert ablation_count == 1 or ablation_count == 0
config['avg_garment_path'] = 'custom_data/avg/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv train_no_shape_one_hot_avg_uv
if config['use_skinning']:
assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
config['batch_size'] = 4
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = False
config['is_cape'] = False
config['device'] = 'cuda'
config['train_dir'] = '/data-local/127/train'
config['valid_dir'] = '/data-local/127/valid'
config['test_dir'] = '/data-net/dataset/test_ready'
config['train_frames'] = None
config['valid_frames'] = None
config['static_model'] = 'checkpoints/iccv/static_torch_orig_2'
config['body_model'] = 'checkpoints/iccv/body_torch_1'
config['normals_model'] = 'checkpoints/iccv/torch/norm_torch_1'
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['latent_dim'] = 128
###############
config['lr'] = 0.0001# 0.0001
###############
config['starting_balancing_coeff'] = 0.0001
config['signature']: set = { 'id',
'frame_nr',
'pose_shape',
'uv_body_posed',
'mesh_body_posed',
'mesh_posed',
'uv_normals',
'uv_static',
'mesh_mask',
'uv_mask',
}
if config['use_skinning']:
config['signature'].add('uv_unposed')
config['signature'].add('mesh_unposed')
else:
config['signature'].add('uv_posed')
config['upscale_enhancer'] = ''
config['freeze_encoders'] = True
config['use_one_hot'] = False
config['bottleneck_scale'] = 2
config['cvae'] = dict()
config['cvae']['device'] = config['device']
config['cvae']['is_label'] = False
config['cvae']['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
config['cvae']['kernel_size'] = 7
config['cvae']['bottleneck_size'] = 64 * config['bottleneck_scale']
config['cvae']['latent_dim'] = config['cvae']['bottleneck_size']
block_setting = [
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 3),
CNBlockConfig(128, 256, 3),
CNBlockConfig(256, config['cvae']['bottleneck_size'], 9),
CNBlockConfig(config['cvae']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cvae']['bottleneck_size'] * 2, 256, 3),
CNBlockConfig(256, 128, 9),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 3),
CNBlockConfig(32, 3, 3),
]
config['cvae']['encoder_block_settings'] = block_setting
config['cvae']['decoder_block_settings'] = block_setting_reverse
####################### cond_ae_norm ############################
config['cond_ae_norm'] = dict()
config['cond_ae_norm']['device'] = config['device']
config['cond_ae_norm']['is_label'] = True
config['cond_ae_norm']['latent_dim'] = config['latent_dim']
config['cond_ae_norm']['kernel_size'] = 7
config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 9),
CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
]
config['cond_ae_norm']['encoder_block_settings'] = block_setting
config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_norm']['activation'] = 'tanh'
####################### cond_ae_body ############################
config['cond_ae_body'] = dict()
config['cond_ae_body']['device'] = config['device']
config['cond_ae_body']['is_label'] = False
config['cond_ae_body']['latent_dim'] = config['latent_dim']
config['cond_ae_body']['kernel_size'] = 7
config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 9),
CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
]
config['cond_ae_body']['encoder_block_settings'] = block_setting
config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_body']['activation'] = 'tanh'
####################### cond_ae_static ############################
config['cond_ae_static'] = dict()
config['cond_ae_static']['device'] = config['device']
config['cond_ae_static']['is_label'] = False
config['cond_ae_static']['latent_dim'] = config['latent_dim']
config['cond_ae_static']['kernel_size'] = 7
config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(8, 16, 3),
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 8, 3),
CNBlockConfig(8, 3, 3),
]
config['cond_ae_static']['encoder_block_settings'] = block_setting
config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_static']['activation'] = 'tanh'
configs/ae_norm_1.py
-73
View File
@@ -1,73 +0,0 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '0'
config['name'] = 'norm_torch_2'
config['notes'] = ' '
config['tags'] = ['torch', 'norm']
config['script'] = 'scripts/ae_norm_train.py'
config['pretrained_model'] = ''
config['batch_size'] = 4
config['train_dir'] = '/data-local/127/train'
config['valid_dir'] = '/data-local/127/valid'
config['test_dir'] = '/data-net/dataset/test_ready'
config['train_frames'] = None
config['valid_frames'] = None
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = False
config['is_cape'] = False
config['remove_global_orientation'] = True
config['device'] = 'cuda'
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['latent_dim'] = 128
config['bottleneck_scale'] = 2
config['cond_ae'] = dict()
config['cond_ae']['device'] = config['device']
config['cond_ae']['is_label'] = False
config['cond_ae']['latent_dim'] = config['latent_dim']
config['cond_ae']['kernel_size'] = 7
config['cond_ae']['bottleneck_size'] = 40 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 9),
CNBlockConfig(128, config['cond_ae']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
# CNBlockConfig(3, 3, 3),
]
config['cond_ae']['encoder_block_settings'] = block_setting
config['cond_ae']['decoder_block_settings'] = block_setting_reverse
config['cond_ae']['activation'] = 'tanh'
config['lr'] = 0.0001
config['cond_ae']['upscale_output'] = False
config['signature'] = { 'id',
'frame_nr',
'uv_mask',
'uv_normals',
'mesh_mask',
'mesh_normals',
}
configs/cvae_norm_1.py
-172
View File
@@ -1,172 +0,0 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '0'
config['name'] = 'cvae_norm_1'
config['notes'] = ''
config['tags'] = ['iccv', 'cvae_norm', 'torch']
config['script'] = 'scripts/cvae_norm.py'
config['pretrained_model'] = ''
config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
config['batch_size'] = 4
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = False
config['is_cape'] = False
config['device'] = 'cuda'
config['train_dir'] = '/data-local/127/train'
config['valid_dir'] = '/data-local/127/valid'
config['test_dir'] = '/data-net/dataset/test_ready'
config['train_frames'] = None
config['valid_frames'] = None
config['static_model'] = 'checkpoints/iccv/static_torch_orig_2'
config['body_model'] = 'checkpoints/iccv/body_torch_1'
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['latent_dim'] = 128
###############
config['lr'] = 0.0001# 0.0001
###############
config['starting_balancing_coeff'] = 0.0001
config['signature']: set = { 'id',
'frame_nr',
'pose_shape',
'uv_body_posed',
'uv_normals',
'uv_static',
'mesh_mask',
'uv_mask',
'mesh_normals',
}
config['upscale_enhancer'] = ''
config['freeze_encoders'] = True
config['use_one_hot'] = False
config['bottleneck_scale'] = 2
config['cvae'] = dict()
config['cvae']['device'] = config['device']
config['cvae']['is_label'] = False
config['cvae']['input_features'] = 3+3+3
config['cvae']['kernel_size'] = 7
config['cvae']['bottleneck_size'] = 64 * config['bottleneck_scale']
config['cvae']['latent_dim'] = config['cvae']['bottleneck_size']
block_setting = [
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 3),
CNBlockConfig(128, 256, 3),
CNBlockConfig(256, config['cvae']['bottleneck_size'], 9),
CNBlockConfig(config['cvae']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(176, 256, 3),
CNBlockConfig(256, 128, 9),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 3),
CNBlockConfig(32, 3, 3),
]
config['cvae']['encoder_block_settings'] = block_setting
config['cvae']['decoder_block_settings'] = block_setting_reverse
####################### cond_ae_curv ############################
config['cond_ae_curv'] = dict()
config['cond_ae_curv']['device'] = config['device']
config['cond_ae_curv']['is_label'] = True
config['cond_ae_curv']['latent_dim'] = config['latent_dim']
config['cond_ae_curv']['kernel_size'] = 7
config['cond_ae_curv']['bottleneck_size'] = 40 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 9),
CNBlockConfig(128, config['cond_ae_curv']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], 128, 3),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
]
config['cond_ae_curv']['encoder_block_settings'] = block_setting
config['cond_ae_curv']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_curv']['activation'] = 'tanh'
####################### cond_ae_body ############################
config['cond_ae_body'] = dict()
config['cond_ae_body']['device'] = config['device']
config['cond_ae_body']['is_label'] = False
config['cond_ae_body']['latent_dim'] = config['latent_dim']
config['cond_ae_body']['kernel_size'] = 7
config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 9),
CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
]
config['cond_ae_body']['encoder_block_settings'] = block_setting
config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_body']['activation'] = 'tanh'
####################### cond_ae_static ############################
config['cond_ae_static'] = dict()
config['cond_ae_static']['device'] = config['device']
config['cond_ae_static']['is_label'] = False
config['cond_ae_static']['latent_dim'] = config['latent_dim']
config['cond_ae_static']['kernel_size'] = 7
config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(8, 16, 3),
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 8, 3),
CNBlockConfig(8, 3, 3),
]
config['cond_ae_static']['encoder_block_settings'] = block_setting
config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_static']['activation'] = 'tanh'
configs/direct_1.py
+77
View File
@@ -0,0 +1,77 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '5'
config['name'] = 'direct_1'
config['notes'] = ' '
config['tags'] = ['torch', 'direct']
config['script'] = 'scripts/direct_train.py'
config['pretrained_model'] = ''
config['use_curvature'] = False
config['batch_size'] = 4
config['train_dir'] = '/data-local/127/train'
config['valid_dir'] = '/data-local/127/valid'
config['test_dir'] = '/data-net/dataset/test_ready'
config['train_frames'] = None
config['valid_frames'] = None
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = False
config['is_cape'] = False
config['remove_global_orientation'] = True
config['device'] = 'cuda'
config['input_features'] = 6
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['latent_dim'] = 128
config['bottleneck_scale'] = 2
config['is_label'] = False
config['kernel_size'] = 7
config['bottleneck_size'] = 64 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 128, 3),
CNBlockConfig(128, 256, 9),
CNBlockConfig(256, config['bottleneck_size'], 9),
CNBlockConfig(config['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['bottleneck_size'], 256, 3),
CNBlockConfig(256, 128, 9),
CNBlockConfig(128, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
# CNBlockConfig(3, 3, 3),
]
config['encoder_block_settings'] = block_setting
config['decoder_block_settings'] = block_setting_reverse
config['activation'] = 'tanh'
config['lr'] = 0.0001
config['reg_loss_coeff'] = 0.1
config['signature'] = { 'id',
'frame_nr',
'pose_shape',
'uv_body_posed',
'mesh_mask',
'mesh_body_posed',
'mesh_posed',
'uv_unposed',
'mesh_unposed',
'uv_static',
'uv_mask',
}
configs/direct_2.py
+78
View File
@@ -0,0 +1,78 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '0'
config['name'] = 'direct_2'
config['notes'] = 'with curvature input'
config['tags'] = ['torch', 'direct']
config['script'] = 'scripts/direct_train.py'
config['pretrained_model'] = ''
config['use_curvature'] = True
config['batch_size'] = 4
config['train_dir'] = '/data-local/127/train'
config['valid_dir'] = '/data-local/127/valid'
config['test_dir'] = '/data-net/dataset/test_ready'
config['train_frames'] = None
config['valid_frames'] = None
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = False
config['is_cape'] = False
config['remove_global_orientation'] = True
config['device'] = 'cuda'
config['input_features'] = 10
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['latent_dim'] = 128
config['bottleneck_scale'] = 2
config['is_label'] = False
config['kernel_size'] = 7
config['bottleneck_size'] = 64 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 128, 3),
CNBlockConfig(128, 256, 9),
CNBlockConfig(256, config['bottleneck_size'], 9),
CNBlockConfig(config['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['bottleneck_size'], 256, 3),
CNBlockConfig(256, 128, 9),
CNBlockConfig(128, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
# CNBlockConfig(3, 3, 3),
]
config['encoder_block_settings'] = block_setting
config['decoder_block_settings'] = block_setting_reverse
config['activation'] = 'tanh'
config['lr'] = 0.0001
config['reg_loss_coeff'] = 0.1
config['signature'] = { 'id',
'frame_nr',
'pose_shape',
'uv_curvature',
'mesh_mask',
'mesh_body_posed',
'mesh_posed',
'uv_unposed',
'mesh_unposed',
'uv_static',
'uv_mask',
'uv_body_posed',
}
configs/direct_3.py
+78
View File
@@ -0,0 +1,78 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '4'
config['name'] = 'direct_3'
config['notes'] = 'with curvature input and skip connection'
config['tags'] = ['torch', 'direct']
config['script'] = 'scripts/direct_train.py'
config['pretrained_model'] = ''
config['use_curvature'] = True
config['batch_size'] = 4
config['train_dir'] = '/data-local/127/train'
config['valid_dir'] = '/data-local/127/valid'
config['test_dir'] = '/data-net/dataset/test_ready'
config['train_frames'] = None
config['valid_frames'] = None
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = False
config['is_cape'] = False
config['remove_global_orientation'] = True
config['device'] = 'cuda'
config['input_features'] = 10
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['latent_dim'] = 128
config['bottleneck_scale'] = 2
config['is_label'] = False
config['kernel_size'] = 7
config['bottleneck_size'] = 64 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 128, 3),
CNBlockConfig(128, 256, 9),
CNBlockConfig(256, config['bottleneck_size'], 9),
CNBlockConfig(config['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['bottleneck_size'], 256, 3),
CNBlockConfig(256*2, 128, 9),
CNBlockConfig(128*2, 32, 9),
CNBlockConfig(32*2, 16, 3),
CNBlockConfig(16*2, 8, 3),
CNBlockConfig(8, 3, 3),
]
config['encoder_block_settings'] = block_setting
config['decoder_block_settings'] = block_setting_reverse
config['activation'] = 'tanh'
config['lr'] = 0.0001
config['reg_loss_coeff'] = 0.1
config['signature'] = { 'id',
'frame_nr',
'pose_shape',
'uv_curvature',
'mesh_mask',
'mesh_body_posed',
'mesh_posed',
'uv_unposed',
'mesh_unposed',
'uv_static',
'uv_mask',
'uv_body_posed',
}
custom_data/zero_body.npy
BIN
View File
Binary file not shown.
custom_utils/custom_utils.py
+1
View File
@@ -43,6 +43,7 @@ class UV2Mesh(nn.Module):
if uv.shape[1] == 3 or uv.shape[1] == 7: if uv.shape[1] == 3 or uv.shape[1] == 7:
uv = torch.permute(uv, [0,2,3,1]) uv = torch.permute(uv, [0,2,3,1])
if uv.shape[1] != 512: if uv.shape[1] != 512:
raise NotImplementedError
uv = Resize(uv[:,:,None], size=(512, 256), interpolation=InterpolationMode.BILINEAR) uv = Resize(uv[:,:,None], size=(512, 256), interpolation=InterpolationMode.BILINEAR)
# batch_size = uv.shape[0] # batch_size = uv.shape[0]
data/calculate_normals.py
-112
View File
@@ -1,112 +0,0 @@
import os
import sys
import numpy as np
import trimesh as tm
import scipy.io as sio
import plotly.graph_objects as go
import plotly.express as px
from tqdm import tqdm
import torch
from scipy.spatial.transform import Rotation as R
import argparse
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from data.cloth3d_io import readOBJ, readPC2
from data.utils import Frame
from smpl.smpl_torch import SMPL
def calculate_normals(vertices, mesh_file, poses):
_, faces, _, _ = readOBJ(mesh_file)
rodr = smpl.rodrigues(torch.tensor(poses.T.reshape(-1, 1, 3), dtype=torch.float32)).numpy().reshape(vertices.shape[0], -1, 3, 3)
normals = []
for i in range(vertices.shape[0]):
r = R.from_matrix(rodr[i, 0, :, :])
# assert r.as_matrix() == rodr[frame, 0, :, :]
out = r.apply(vertices[i], inverse=True)
# create a plotly mesh plot of out and vertices [i]
# fig = go.Figure(data=[
# go.Mesh3d(
# x=out[:, 0], y=out[:, 2], z=out[:, 1],
# # i, j and k give the vertices of triangles
# i=faces[:, 0], j=faces[:, 1], k=faces[:, 2],
# name='y',
# showscale=False
# ),
# go.Mesh3d(
# x=vertices[i][:, 0], y=vertices[i][:, 1], z=vertices[i][:, 2],
# # i, j and k give the vertices of triangles
# i=faces[:, 0], j=faces[:, 1], k=faces[:, 2],
# name='x',
# showscale=False
# )
# ])
# fig.show()
mesh = tm.Trimesh(vertices=out, faces=faces, process=False, validate=False)
normals.append(mesh.vertex_normals)
return np.stack(normals)
def get_frames(folder):
uv_static_folder = os.path.join(folder, 'pose_shape')
frames = sorted(os.listdir(uv_static_folder))
frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames]
return frames
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("-f", "--folder", help = "end of range", default='/data-net/dataset/00761')
parser.add_argument("-f_obj", "--folder_object", help = "end of range", default='/data-net/dataset/cloth3d/train')
parser.add_argument("-s", "--start", help = "start of range", default=0)
parser.add_argument("-e", "--end", help = "end of range", default=10)
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
args = parser.parse_args()
obj_folder = args.folder_object # '/data-net/dataset/cloth3d/train'
folder = args.folder #'/data-local/127/train'
os.makedirs(os.path.join(folder, 'mesh_normals'), exist_ok=True)
frames = get_frames(folder)
smpl = SMPL('cpu', 'smpl')
ids = {frame.id for frame in frames}
ids = sorted(list(ids))[int(args.start):int(args.end)]
for id in tqdm(ids):
# find obj files in the folder
garment_obj_folder = os.path.join(obj_folder, id[:5])
garments = [file.split('.')[0] for file in os.listdir(garment_obj_folder) if file.endswith('.obj')]
if len(id) == 5:
garment = garments[0]
elif len(id) == 6:
garment = garments[1]
else:
raise ValueError('Invalid frame id')
mesh_file = os.path.join(garment_obj_folder, f'{garment}.obj')
vertices =readPC2(os.path.join(garment_obj_folder, f'{garment}.pc16'), True)['V']
poses = sio.loadmat(os.path.join(garment_obj_folder, f'info.mat'))['poses']
normals = calculate_normals(vertices, mesh_file, poses)
map = sio.loadmat(os.path.join(garment_obj_folder, f'{garment}_map.mat'))['map'][0]
mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy'))
mesh_normals = normals[:,map,:]
mesh_normals = mesh_normals * mask[None,:, None]
for i in range(vertices.shape[0]):
# np.savez(os.path.join(folder, 'mesh_normals', f'{id}_{i}.npz'), mesh_normals=mesh_normals)
np.savez_compressed(os.path.join(folder, 'mesh_normals', f'{id}_{i}.npz'), mesh_normals=mesh_normals[i])
# pass
# print(frame.id, frame.frame_nr)
data/cloth3d_io.py
+1 -12
View File
@@ -342,6 +342,7 @@ import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
import plotly.graph_objects as go import plotly.graph_objects as go
from custom_utils.custom_utils import _maskFaces
def plot_3d(mesh, faces): def plot_3d(mesh, faces):
fig = go.Figure() fig = go.Figure()
@@ -369,19 +370,7 @@ def plot_3d(mesh, faces):
return fig return fig
def quads2tris(F):
out = []
for f in F:
if len(f) == 3: out += [f]
elif len(f) == 4: out += [[f[0],f[1],f[2]],
[f[0],f[2],f[3]]]
else: print("This should not happen...")
return np.array(out, np.int32)
if __name__ == '__main__': if __name__ == '__main__':
from custom_utils.custom_utils import _maskFaces
folder = '/home/workspace/data/dataset/cloth3d/train' folder = '/home/workspace/data/dataset/cloth3d/train'
F = np.load("custom_utils/smpl_faces.npy")#.astype(np.int32) F = np.load("custom_utils/smpl_faces.npy")#.astype(np.int32)
data/convert_to_uv.py
+148 -125
View File
@@ -12,7 +12,7 @@ import pickle
import numpy as np import numpy as np
import scipy as sp import scipy as sp
import trimesh as tm import trimesh as tm
# import cmapy as cm import cmapy as cm
import tables as tb import tables as tb
import scipy.io as sio import scipy.io as sio
import scipy.interpolate as spi import scipy.interpolate as spi
@@ -24,7 +24,7 @@ import plotly.graph_objects as go
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from data.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo, read_posed_body_frame, read_posed_body, read_posed_garment, read_curvatures from data.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo, read_posed_body_frame, read_posed_body, read_posed_garment, read_curvatures
from data.cloth3d_io import readPC2 from data.utils import save_to_h5file
os.environ["CUDA_VISIBLE_DEVICES"] = "6" os.environ["CUDA_VISIBLE_DEVICES"] = "6"
@@ -34,26 +34,26 @@ def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rb
uv_values = np.zeros((uv_map.shape[0], 3)) uv_values = np.zeros((uv_map.shape[0], 3))
uv_values[m2uv[:, 0]] = verts uv_values[m2uv[:, 0]] = verts
uv_values[m2uv[:, 1]] += verts uv_values[m2uv[:, 1]] += verts
uv_values = uv_values[uv_map_mask[type], :] # add garment mask uv_values = uv_values[uv_map_mask, :] # add garment mask
uv_values /= uv_counter[type][:, np.newaxis] uv_values /= uv_counter[:, np.newaxis]
return_garment_uv_mask = False return_garment_uv_mask = False
if mask is not None: if mask is not None:
garment_mask = np.zeros((uv_map.shape[0]), dtype=bool) garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
garment_mask[m2uv[:, 0]] = mask garment_mask[m2uv[:, 0]] = mask
garment_mask[m2uv[:, 1]] += mask garment_mask[m2uv[:, 1]] += mask
garment_mask = garment_mask[uv_map_mask[type]] # add garment mask garment_mask = garment_mask[uv_map_mask] # add garment mask
else: else:
garment_mask = np.sum(uv_values, axis=-1) != 0 garment_mask = np.sum(uv_values, axis=-1) != 0
# test # test
verts_uv[uv_map_discreet_unique[type][:, 0], uv_map_discreet_unique[type][:, 1], :] = uv_values verts_uv[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
# verts_uv_save = verts_uv.copy() # verts_uv_save = verts_uv.copy()
# garment_mask = np.sum(uv_values, axis=-1) != 0 # garment_mask = np.sum(uv_values, axis=-1) != 0
unique_uv_values = uv_values[garment_mask, :] unique_uv_values = uv_values[garment_mask, :]
uv_map_discreet_unique_unique = uv_map_discreet_unique[type][garment_mask, :] uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
rbf_dest_current = rbf_dest[type] rbf_dest_current = rbf_dest
if garment_uv_mask is None and mask is not None: if garment_uv_mask is None and mask is not None:
garment_uv_mask = np.zeros((res * 2, res)) garment_uv_mask = np.zeros((res * 2, res))
@@ -62,8 +62,8 @@ def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rb
garment_uv_mask_border = garment_uv_mask.astype(bool) garment_uv_mask_border = garment_uv_mask.astype(bool)
garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool) garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)
garment_uv_mask_border = ~garment_uv_mask_border & garment_uv_mask garment_uv_mask_border = ~garment_uv_mask_border & garment_uv_mask
garment_uv_mask = garment_uv_mask * uv_mask[type][:,:,0] garment_uv_mask = garment_uv_mask * uv_mask[:,:,0]
rbf_dest_current = rbf_dest[type] * garment_uv_mask rbf_dest_current = rbf_dest * garment_uv_mask
return_garment_uv_mask = True return_garment_uv_mask = True
@@ -86,7 +86,7 @@ def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rb
# rbf_dest[:, [0,1]] = rbf_dest[:, [1, 0]] # rbf_dest[:, [0,1]] = rbf_dest[:, [1, 0]]
result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = final result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = final
result[uv_map_discreet_unique[type][:, 0], uv_map_discreet_unique[type][:, 1], :] = uv_values result[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
if garment_uv_mask is not None: if garment_uv_mask is not None:
result = result * garment_uv_mask[:,:,np.newaxis] result = result * garment_uv_mask[:,:,np.newaxis]
@@ -102,26 +102,26 @@ def convert_to_uv_batch_body(res, m2uv, uv_map, verts, neighbours=15, mask=None,
uv_values = np.zeros((batch, uv_map.shape[0], 3)) uv_values = np.zeros((batch, uv_map.shape[0], 3))
uv_values[:, m2uv[:, 0]] = verts uv_values[:, m2uv[:, 0]] = verts
uv_values[:, m2uv[:, 1]] += verts uv_values[:, m2uv[:, 1]] += verts
uv_values = uv_values[:, uv_map_mask['non-skirt'], :] uv_values = uv_values[:, uv_map_mask, :]
uv_values /= uv_counter['non-skirt'][np.newaxis, :, np.newaxis] uv_values /= uv_counter[np.newaxis, :, np.newaxis]
return_garment_uv_mask = False return_garment_uv_mask = False
# test # test
verts_uv[:, uv_map_discreet_unique['non-skirt'][:, 0], uv_map_discreet_unique['non-skirt'][:, 1], :] = uv_values verts_uv[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
if mask is not None: if mask is not None:
garment_mask = np.zeros((uv_map.shape[0]), dtype=bool) garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
garment_mask[m2uv[:, 0]] = mask garment_mask[m2uv[:, 0]] = mask
garment_mask[m2uv[:, 1]] += mask garment_mask[m2uv[:, 1]] += mask
garment_mask = garment_mask[uv_map_mask['non-skirt']] # add garment mask garment_mask = garment_mask[uv_map_mask] # add garment mask
else: else:
garment_mask = np.sum(uv_values[0,:,:], axis=-1) != 0 garment_mask = np.sum(uv_values[0,:,:], axis=-1) != 0
unique_uv_values = uv_values[:, garment_mask, :] unique_uv_values = uv_values[:, garment_mask, :]
uv_map_discreet_unique_unique = uv_map_discreet_unique['non-skirt'][garment_mask, :] uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
rbf_dest_current = rbf_dest['non-skirt'] rbf_dest_current = rbf_dest
# add back mesh vertices that are not part of the garment # add back mesh vertices that are not part of the garment
uv_map_discreet_unique_unique_inv = uv_map_discreet_unique['non-skirt'][np.invert(garment_mask), :] uv_map_discreet_unique_unique_inv = uv_map_discreet_unique[np.invert(garment_mask), :]
rbf_dest_fill = np.zeros((res * 2, res), dtype=bool) rbf_dest_fill = np.zeros((res * 2, res), dtype=bool)
rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True
@@ -133,7 +133,7 @@ def convert_to_uv_batch_body(res, m2uv, uv_map, verts, neighbours=15, mask=None,
return_garment_uv_mask = True return_garment_uv_mask = True
if garment_uv_mask is not None: if garment_uv_mask is not None:
rbf_dest_current = (rbf_dest['non-skirt'] + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border rbf_dest_current = (rbf_dest + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border
rbf_dest_current = np.nonzero(rbf_dest_current) rbf_dest_current = np.nonzero(rbf_dest_current)
rbf_dest_current = np.array(rbf_dest_current).T rbf_dest_current = np.array(rbf_dest_current).T
@@ -172,18 +172,18 @@ def convert_to_uv_batch(res, m2uv, uv_map, verts, neighbours=15, mask=None, meth
garment_mask = np.zeros((uv_map.shape[0]), dtype=bool) garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
garment_mask[m2uv[:, 0]] = mask garment_mask[m2uv[:, 0]] = mask
garment_mask[m2uv[:, 1]] += mask garment_mask[m2uv[:, 1]] += mask
garment_mask = garment_mask[uv_map_mask[type]] # add garment mask garment_mask = garment_mask[uv_map_mask] # add garment mask
rbf_dest_current = rbf_dest[type] rbf_dest_current = rbf_dest
# add back mesh vertices that are not part of the garment # add back mesh vertices that are not part of the garment
uv_map_discreet_unique_unique_inv = uv_map_discreet_unique[type][np.invert(garment_mask), :] uv_map_discreet_unique_unique_inv = uv_map_discreet_unique[np.invert(garment_mask), :]
rbf_dest_fill = np.zeros((res * 2, res), dtype=bool) rbf_dest_fill = np.zeros((res * 2, res), dtype=bool)
rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True
assert not (garment_uv_mask is None and mask is not None) assert not (garment_uv_mask is None and mask is not None)
assert garment_uv_mask is not None assert garment_uv_mask is not None
rbf_dest_current = (rbf_dest[type] + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border rbf_dest_current = (rbf_dest + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border
rbf_dest_current = np.nonzero(rbf_dest_current) rbf_dest_current = np.nonzero(rbf_dest_current)
rbf_dest_current = np.array(rbf_dest_current).T rbf_dest_current = np.array(rbf_dest_current).T
@@ -196,11 +196,11 @@ def convert_to_uv_batch(res, m2uv, uv_map, verts, neighbours=15, mask=None, meth
uv_values = np.zeros((uv_map.shape[0], 3)) uv_values = np.zeros((uv_map.shape[0], 3))
uv_values[m2uv[:, 0]] = verts[i] uv_values[m2uv[:, 0]] = verts[i]
uv_values[m2uv[:, 1]] += verts[i] uv_values[m2uv[:, 1]] += verts[i]
uv_values = uv_values[uv_map_mask[type], :] uv_values = uv_values[uv_map_mask, :]
uv_values /= uv_counter[type][:, np.newaxis] uv_values /= uv_counter[:, np.newaxis]
unique_uv_values = uv_values[garment_mask, :] unique_uv_values = uv_values[garment_mask, :]
uv_map_discreet_unique_unique = uv_map_discreet_unique[type][garment_mask, :] uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
interpolator = spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values[:,:], neighbors=neighbours, kernel='linear') interpolator = spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values[:,:], neighbors=neighbours, kernel='linear')
result = np.zeros((res * 2, res, 3), np.float32) result = np.zeros((res * 2, res, 3), np.float32)
@@ -386,12 +386,12 @@ def test(verts, faces, values):
def get_shapedirs(): def get_shapedirs():
to_keep = sio.loadmat(os.path.join('smpl/keep_remove.mat'))['keep'].reshape(-1) to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1)
shapedirs = [] shapedirs = []
with open(os.path.join('smpl/super_model_f.pkl'), 'rb') as f: with open(os.path.join('custom_utils/super_model_f.pkl'), 'rb') as f:
female = pickle.load(f, encoding="latin1") female = pickle.load(f, encoding="latin1")
shapedirs.append(female['shapedirs'][to_keep, :]) shapedirs.append(female['shapedirs'][to_keep, :])
with open(os.path.join('smpl/super_model_m.pkl'), 'rb') as f: with open(os.path.join('custom_utils/super_model_m.pkl'), 'rb') as f:
male = pickle.load(f, encoding="latin1") male = pickle.load(f, encoding="latin1")
shapedirs.append(male['shapedirs'][to_keep, :]) shapedirs.append(male['shapedirs'][to_keep, :])
return shapedirs return shapedirs
@@ -399,139 +399,162 @@ def get_shapedirs():
if __name__ == '__main__': if __name__ == '__main__':
parser = argparse.ArgumentParser() 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='/home/workspace/data/dataset/cloth3d/train')
parser.add_argument("-s", "--start", help = "start of range", default=0, type=int) parser.add_argument("-s", "--start", help = "start of range", default=70500)
parser.add_argument("-e", "--end", help = "end of range", default=100, type=int) parser.add_argument("-e", "--end", help = "end of range", default=90000)
parser.add_argument("-n", "--name", help = "end of range", default='lhunor') parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
parser.add_argument("-sk", "--skirt", help = "end of range", default=0)
args = parser.parse_args() args = parser.parse_args()
F = {'non-skirt': np.load("smpl/smpl_faces.npy"), 'skirt': np.load("smpl/smpl_faces_skirt.npy")} F = np.load("custom_utils/smpl_faces" + ("_skirt" if args.skirt == 1 else "") + ".npy")
uv_map = {'non-skirt': np.load('custom_utils/stretched_uv_map_3.npy'), 'skirt': np.load('custom_utils/stretched_uv_map_skirt.npy')} uv_map = np.load('custom_utils/stretched_uv_map_' + ("skirt" if args.skirt == 1 else "3") + '.npy')
m2uv = {'non-skirt': np.load('custom_utils/m2uv.npy', allow_pickle=True).item(), 'skirt': np.load('custom_utils/m2uv_skirt.npy', allow_pickle=True).item()} m2uv = np.load('custom_utils/m2uv' + ("_skirt" if args.skirt == 1 else "") + '.npy', allow_pickle=True).item()
m2uv = {k: m2uv_array(v) for k, v in m2uv.items()} m2uv = m2uv_array(m2uv)
uv_mask = {'non-skirt': np.load('custom_utils/mask_256x512.npy'), 'skirt': np.load('custom_utils/mask_256x512_skirt.npy')} uv_mask = np.load('custom_utils/mask_256x512' + ("_skirt" if args.skirt == 1 else "") + '.npy')
curvature_label_type = '6_25 '
uv_res = 256 uv_res = 256
uv_map_discreet = dict() uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
uv_map_discreet_unique = dict() # uv_mask_dilate = cv2.erode(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_CROSS, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=1).astype(bool)[:, :, np.newaxis]
uv_counter = dict() # uv_mask_border = uv_mask * np.invert(uv_mask_dilate)
rbf_dest = dict() uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
uv_map_mask = dict()
for k,v in uv_mask.items():
uv_mask[k] = (cv2.resize(v * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
# uv_mask_dilate = cv2.erode(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_CROSS, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=1).astype(bool)[:, :, np.newaxis]
# uv_mask_border = uv_mask * np.invert(uv_mask_dilate)
uv_mask[k] = cv2.dilate(v.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
uv_map_discreet[k] = np.floor(uv_map[k] * uv_res).astype(np.int32) uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
uv_map_mask[k] = unique_uv_map_mask(uv_map_discreet[k]) uv_map_mask = unique_uv_map_mask(uv_map_discreet)
uv_counter[k] = np.zeros((uv_map[k].shape[0])) uv_counter = np.zeros((uv_map.shape[0]))
uv_counter[k][m2uv[k][:, 0]] = 1 uv_counter[m2uv[:, 0]] = 1
uv_counter[k][m2uv[k][:, 1]] += 1 uv_counter[m2uv[:, 1]] += 1
uv_counter[k] = uv_counter[k][uv_map_mask[k]] uv_counter = uv_counter[uv_map_mask]
uv_map_discreet_unique[k] = uv_map_discreet[k][uv_map_mask[k], :] uv_map_discreet_unique = uv_map_discreet[uv_map_mask, :]
rbf_dest[k] = np.ones((uv_res * 2, uv_res)) rbf_dest = np.ones((uv_res * 2, uv_res))
rbf_dest[k][uv_map_discreet_unique[k][:, 0], uv_map_discreet_unique[k][:, 1]] = 0 rbf_dest[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = 0
rbf_dest[k][np.invert(uv_mask[k][:,:,0])] = 0 rbf_dest[np.invert(uv_mask[:,:,0])] = 0
######################################################################################### #########################################################################################
use_no_orient = False
no_shape = False
is_skirt = True if args.skirt == 1 else False
folder = args.folder folder = args.folder
data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_normals', 'body_posed', 'mesh_unposed'] # 'mesh_unposed', neighbourhood = 25
frames_file = 'frames_all.txt'
data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_curvature', 'body_posed'] # 'mesh_unposed',
ids = sorted(os.listdir(folder + '/mesh_body_posed_no_orient'))[args.start:args.end] print("skirt" if args.skirt == 1 else "not skirt")
ids = [id.split('.')[0].split('_')[0] for id in ids] print(f'params: -f {args.folder} -s {args.start} -e {args.end} -n {args.name} -skirt {args.skirt}')
skirt_modifier = '_skirt' if is_skirt else ''
orient_mod = '_no_orient' if use_no_orient else ''
shape_mod = '_no_shape' if no_shape else ''
frames = []
with open(os.path.join(folder, frames_file), 'r') as f:
lines = f.readlines()
lines = [line.strip() for line in lines]
frames = lines[int(args.start):int(args.end)]
ids = [frame.split()[0] for frame in frames]
ids = sorted(list(set(ids)))
shapedirs = get_shapedirs() shapedirs = get_shapedirs()
root = '/'.join(folder.split('/')[:-1]) root = '/'.join(folder.split('/')[:-1])
type = folder.split('/')[-1] type = folder.split('/')[-1]
# h5file_curv = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='r')
# h5file = tb.open_file(os.path.join(root, f'{type}_extras_uv.h5'), mode='a')
os.makedirs(folder, exist_ok=True) extras_folder = os.path.join(root, f'{type}_extras')
os.makedirs(os.path.join(folder, 'uv_posed_no_orient'), exist_ok=True) os.makedirs(extras_folder, exist_ok=True)
os.makedirs(os.path.join(folder, 'uv_static'), exist_ok=True) os.makedirs(os.path.join(extras_folder, 'uv_posed' + orient_mod), exist_ok=True)
os.makedirs(os.path.join(folder, 'uv_normals'), exist_ok=True) os.makedirs(os.path.join(extras_folder, 'uv_static'), exist_ok=True)
os.makedirs(os.path.join(folder, 'uv_unposed_no_shape'), exist_ok=True) os.makedirs(os.path.join(extras_folder, 'uv_curvature'), exist_ok=True)
os.makedirs(os.path.join(folder, 'uv_body_posed_no_orient'), exist_ok=True) os.makedirs(os.path.join(extras_folder, 'uv_unposed'), exist_ok=True)
os.makedirs(os.path.join(folder, 'uv_mask'), exist_ok=True) os.makedirs(os.path.join(extras_folder, 'uv_body_posed' + orient_mod), exist_ok=True)
os.makedirs(os.path.join(extras_folder, 'uv_mask'), exist_ok=True)
for id in tqdm.tqdm(ids): for id in tqdm.tqdm(ids):
garments = glob.glob(os.path.join(folder, 'mesh_mask', f'{id}_*'))
garments = [garment.split('/')[-1].split('.')[0].split('_')[-1] for garment in garments]
# if os.path.exists(os.path.join(folder, 'uv_static', f'{id}_{garment}.npz')):
# continue
for k, garment in enumerate(garments): info = loadInfo(os.path.join(folder, id, 'info.mat'))
# if not ('Dress' in garments or 'Skirt' in garments): garments = list(info['outfit'].keys())
# if k != 1: assert info['poses'].shape[1] <= 300
# files = glob.glob(os.path.join(folder, 'uv_body_posed_no_orient', f'{id}_*_{garment}.npz'))
# for file in files: # put all this in a function and call for all garments
# shutil.move(file, '_'.join(file.split('_')[:-1]) + '.npz')
# also masking is not the same as just filling the empty values with neighbouring
# values, at least the plots are different
for garment in garments:
# garment = garments[0]
# if all the uv maps have been generated, skip this id
# if all([os.path.exists(os.path.join(extras_folder, 'uv_body_posed' + orient_mod, f'{id}_{i}_{garment}.npz')) for i in range(info['poses'].shape[1])]):
# continue # continue
mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}_{garment}.npy')) mask = read_garment_mask(folder, id, garment)
x = readPC2(os.path.join(folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V']
mesh_posed = np.zeros((x.shape[0], 14475, 3), dtype=np.float32)
mesh_posed[:,mask, :] = x
if 'mesh_unposed' in data_to_convert:
mesh_unposed = np.zeros((x.shape[0], 14475, 3), dtype=np.float32)
mesh_unposed[:,mask, :] = readPC2(os.path.join(folder, 'mesh_unposed_no_shape' ,f'{id}_{garment}.pc16'), True)['V']
if 'mesh_static' in data_to_convert:
mesh_static = np.zeros((14475, 3), dtype=np.float32)
mesh_static[mask, :] = np.load(os.path.join(folder, 'mesh_static', f'{id}_{garment}.npy'))
if 'mesh_normals' in data_to_convert:
mesh_normals = np.zeros((x.shape[0], 14475, 3), dtype=np.float32)
mesh_normals[:,mask, :] = readPC2(os.path.join(folder, 'mesh_normals', f'{id}_{garment}.pc16'), True)['V']
if 'body_posed' in data_to_convert and k != 1:
body_posed = readPC2(os.path.join(folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V']
start = time.perf_counter()
type = 'skirt' if garment == 'Dress' or garment == 'Skirt' else 'non-skirt'
uv_static, uv_mask_local, garment_uv_mask_border = convert_to_uv(uv_res, m2uv[type], uv_map[type], mesh_static, neighbours=3, mask=mask, method='rbf', garment_uv_mask=None)
uv_static = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_static[None, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
uv_static = uv_static[0].astype(np.float32)
if 'mesh_posed' in data_to_convert: if 'mesh_posed' in data_to_convert:
uv_posed = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border) # mesh_posed_frame = read_posed_garment_frame(folder, id, garment, 0, mask=mask, shape=info['shape'], pose=info['poses'][:,0], deorient=use_no_orient, unshape=no_shape)
mesh_posed = read_posed_garment(folder, id, garment, mask=mask, shape=info['shape'], poses=info['poses'], deorient=use_no_orient, unshape=no_shape)
if 'mesh_unposed' in data_to_convert: 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) raise NotImplementedError()
if 'mesh_normals' in data_to_convert: mesh_unposed = read_unposed_garment(folder, id, garment, mask=mask, shape=info['shape'], unshape=no_shape)
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 'mesh_static' in data_to_convert:
if 'body_posed' in data_to_convert and k != 1: mesh_static = read_garment_template(folder, id, garment, mask)
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) if 'mesh_curvature' in data_to_convert:
mesh_curvature = read_curvatures(folder + '_extras', id, garment, count=info['poses'].shape[1])
if 'body_posed' in data_to_convert:
body_posed = read_posed_body(folder, id, info['poses'], use_no_orient)
start = time.perf_counter()
uv_static, uv_mask_local, garment_uv_mask_border = convert_to_uv(uv_res, m2uv, uv_map, mesh_static, neighbours=3, mask=mask, method='rbf', garment_uv_mask=None)
uv_static = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_static[None, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
uv_static = uv_static.astype(np.float32)
if 'mesh_posed' in data_to_convert:
uv_posed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
if 'mesh_unposed' in data_to_convert:
raise NotImplementedError()
uv_unposed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
if 'mesh_curvature' in data_to_convert:
uv_curvature = convert_to_uv_batch_label(uv_res, m2uv, uv_map, mesh_curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
if 'body_posed' in data_to_convert:
uv_body_posed = convert_to_uv_batch_body(uv_res, m2uv, uv_map, body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border)
print(time.perf_counter() - start) print(time.perf_counter() - start)
# ex.imshow(uv_mask_local).show() # mesh_static
# ex.imshow((uv_static+1)/2).show() # save_to_h5file(h5file, id, 0, garment, uv_static, 'uv_static', 'float32')
# ex.imshow((uv_posed[0]+1)/2).show() # save_to_h5file(h5file, id, 0, garment, uv_mask, 'uv_mask', 'float32')
# ex.imshow((uv_unposed[0]+1)/2).show() np.savez_compressed(os.path.join(extras_folder, 'uv_static', f'{id}_{garment}.npz'), uv_static=uv_static)
# ex.imshow((uv_normals[0]+1.5)/3).show() np.savez_compressed(os.path.join(extras_folder, 'uv_mask', f'{id}_{garment}.npz'), uv_mask=uv_mask_local)
# ex.imshow((uv_body_posed[0]+1)/2).show()
# pass
np.savez_compressed(os.path.join(folder, 'uv_static', f'{id}_{garment}.npz'), uv_static=uv_static)
np.savez_compressed(os.path.join(folder, 'uv_mask', f'{id}_{garment}.npz'), uv_mask=uv_mask_local)
if 'mesh_posed' in data_to_convert: if 'mesh_posed' in data_to_convert:
for i in range(len(uv_posed)): for i in range(len(uv_posed)):
np.savez_compressed(os.path.join(folder, 'uv_posed_no_orient', f'{id}_{i}_{garment}.npz'), uv_posed=uv_posed[i]) np.savez_compressed(os.path.join(extras_folder, 'uv_posed' + orient_mod, f'{id}_{i}_{garment}.npz'), uv_posed=uv_posed[i])
# save_to_h5file(h5file, id, i, garment, uv_posed[i, :, :, :], 'uv_posed' + orient_mod + shape_mod, 'float32')
if 'mesh_unposed' in data_to_convert: if 'mesh_unposed' in data_to_convert:
for i in range(len(uv_unposed)): raise NotImplementedError()
np.savez_compressed(os.path.join(folder, 'uv_unposed_no_shape', f'{id}_{i}_{garment}.npz'), uv_unposed=uv_unposed[i]) for i in range(uv_unposed.shape[0]):
if 'mesh_normals' in data_to_convert: save_to_h5file(h5file, id, i, garment, uv_unposed[i, :, :, :], 'uv_unposed' + orient_mod + shape_mod, 'float32')
for i in range(len(uv_normals)): if 'mesh_curvature' in data_to_convert:
np.savez_compressed(os.path.join(folder, 'uv_normals', f'{id}_{i}_{garment}.npz'), uv_normals=uv_normals[i]) for i in range(len(uv_curvature)):
if 'body_posed' in data_to_convert and k != 1: np.savez_compressed(os.path.join(extras_folder, 'uv_curvature', f'{id}_{i}_{garment}.npz'), uv_curvature=uv_curvature[i].astype(np.uint8))
# save_to_h5file(h5file, id, i, garment, uv_curvature[i, :, :, :], 'uv_curvature', 'uint8')
if 'body_posed' in data_to_convert:
for i in range(len(uv_body_posed)): for i in range(len(uv_body_posed)):
np.savez_compressed(os.path.join(folder, 'uv_body_posed_no_orient', f'{id}_{i}.npz'), uv_body_posed=uv_body_posed[i]) np.savez_compressed(os.path.join(extras_folder, 'uv_body_posed' + orient_mod, f'{id}_{i}_{garment}.npz'), uv_body_posed=uv_body_posed[i])
# save_to_h5file(h5file, id, i, garment, uv_body_posed[i, :, :, :], 'uv_body_posed', 'float32')
# pass pass
h5file_curv.close()
data/dataset.py
+65 -132
View File
@@ -21,7 +21,6 @@ from torchvision.transforms.functional import InterpolationMode
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from data.utils import Frame from data.utils import Frame
from data.cloth3d_io import readPC2, readPC2Frame
def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=None): def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=None):
@@ -61,6 +60,7 @@ class StaticCache():
self.folder = folder self.folder = folder
assert set(resolutions).issubset(set([32, 64, 128, 256, 512])) assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
self.resolutions = resolutions self.resolutions = resolutions
self.filled = ''
self.is_cape = is_cape self.is_cape = is_cape
self.device = device self.device = device
@@ -74,65 +74,50 @@ class StaticCache():
self.uv_masks[res] = down.cpu().numpy() self.uv_masks[res] = down.cpu().numpy()
def read(self, id, garment): def read(self, id):
data = dict() data = dict()
mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', f'{id}_{garment}.npy')); mesh_mask = mesh_mask.astype(bool)
data['mesh_mask'] = mesh_mask
for res in self.resolutions: for res in self.resolutions:
res_string = '' res_string = ''
if res != 512: if res != 512:
res_string = '_' + str(res) res_string = '_' + str(res)
if 'uv_static' + res_string in self.signature: if 'uv_static' + res_string in self.signature:
uv_static = np.load(os.path.join(self.folder, f'uv_static{res_string}', f'{id}_{garment}.npz')); uv_static = uv_static['uv_static']; data['uv_static' + res_string] = uv_static uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; data['uv_static' + res_string] = uv_static
if 'uv_mask' + res_string in self.signature: if 'uv_mask' + res_string in self.signature:
if self.is_cape: if self.is_cape:
data['uv_mask' + res_string] = self.uv_masks[res] data['uv_mask' + res_string] = self.uv_masks[res]
else: else:
uv_mask = np.load(os.path.join(self.folder, f'uv_mask{res_string}', f'{id}_{garment}.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask uv_mask = np.load(os.path.join(self.folder, 'uv_mask' + res_string, id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask
if 'mesh_static' in self.signature: if 'mesh_static' in self.signature:
mesh_static = np.load(os.path.join(self.folder, 'mesh_static', f'{id}_{garment}.npy')) mesh_static = np.load(os.path.join(self.folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static.astype(np.float32)
data['mesh_static'] = np.zeros((14475, 3), dtype=np.float32) if 'mesh_mask' in self.signature:
data['mesh_static'][mesh_mask] = mesh_static.astype(np.float32) if self.is_cape:
data['mesh_mask'] = np.ones((3638), np.float32)
else:
mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask
# if 'uv_shaped_body' in self.signature: # if 'uv_shaped_body' in self.signature:
# uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; data['uv_shaped_body'] = uv_shaped_body # uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; data['uv_shaped_body'] = uv_shaped_body
return data return data
def __getitem__(self, frame) -> dict: def __getitem__(self, id) -> dict:
if frame not in self.store: if id not in self.store:
item = self.read(frame.id, frame.garment) item = self.read(id)
stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()} stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()}
for key, value in stored_data.items(): for key, value in stored_data.items():
if key.startswith('uv_') and value.shape[-1] == 3: if key.startswith('uv_') and value.shape[-1] == 3:
stored_data[key] = torch.permute(value, [2,0,1]) stored_data[key] = torch.permute(value, [2,0,1])
self.store[frame] = stored_data self.store[id] = stored_data
data = { key: value.to(self.device) for key, value in self.store[frame].items()} data = { key: value.to(self.device) for key, value in self.store[id].items()}
return data return data
def read_frames_file(frames_file, shuffle=False):
frames = []
with open(frames_file, 'r') as f:
lines = f.readlines()
lines = [line.strip() for line in lines]
frames = [Frame(line.split()[0], int(line.split()[1]), line.split()[2]) for line in lines]
if shuffle:
random.shuffle(frames)
return frames
class Dataset(Dataset): class Dataset(Dataset):
def __init__(self, folder, signature, resolutions=[512], frames_file='', device='cuda', use_cache=False, is_cape=False): def __init__(self, folder, signature, resolutions=[512], frames_file='frames_all.txt', device='cuda', use_cache=False, is_cape=False):
self.folder = folder self.folder = folder
self.device = device self.device = device
self.is_cape = is_cape self.is_cape = is_cape
@@ -140,20 +125,13 @@ class Dataset(Dataset):
self.resolutions = resolutions self.resolutions = resolutions
self.transforms = tv.transforms.ToTensor() self.transforms = tv.transforms.ToTensor()
self.frames = read_frames_file(os.path.join(folder, frames_file), shuffle=True) self.frames = read_frames(folder, frames_file=frames_file, is_cape=is_cape, shuffle=True)
self.body_faces = np.load('smpl/smpl_faces.npy')
self.static_cache = StaticCache(folder, signature, is_cape=is_cape, resolutions=resolutions, device=device) self.static_cache = StaticCache(folder, signature, is_cape=is_cape, resolutions=resolutions, device=device)
self.cache = dict() if use_cache else None self.cache = dict() if use_cache else None
ids = set([frame.id for frame in self.frames])
self.pose_cache = dict()
for id in ids:
self.pose_cache[id] = np.load(os.path.join(self.folder, 'pose', f'{id}.npy'))
self.shape_cache = dict()
for id in ids:
self.shape_cache[id] = np.load(os.path.join(self.folder, 'shape', f'{id}.npy'))
def __len__(self): def __len__(self):
return len(self.frames) return len(self.frames)
@@ -167,56 +145,66 @@ class Dataset(Dataset):
data = { **data, **self.static_cache[frame.id] } data = { **data, **self.static_cache[frame.id] }
return data return data
separator = '_'
if self.is_cape:
separator = '.'
uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient_2'
if self.is_cape:
uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient'
item = dict() item = dict()
mesh_mask = self.static_cache[frame]['mesh_mask'].cpu().numpy().astype(bool)
item['id'] = frame.id
item['frame_nr'] = frame.frame_nr
item['garment'] = frame.garment
if 'id' in self.signature:
item['id'] = frame.id
if 'frame_nr' in self.signature:
item['frame_nr'] = frame.frame_nr
if 'pose_shape' in self.signature: if 'pose_shape' in self.signature:
pose = self.pose_cache[frame.id][frame.frame_nr] pose_shape = np.load(os.path.join(self.folder, 'pose_shape', frame.id + separator + frame.frame_nr + '.npy')).astype(np.float32)
pose[:3] = 0 pose_shape[:3] = 0
item['pose_shape'] = np.concatenate((pose, self.shape_cache[frame.id]), axis=0) item['pose_shape'] = pose_shape
if 'mesh_body_posed' in self.signature: if 'mesh_body_posed' in self.signature:
mesh_body_posed_no_orient = readPC2Frame(os.path.join(self.folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), frame.frame_nr, True) mesh_posed_body_no_orient = np.load(os.path.join(self.folder, 'mesh_posed_body_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
item['mesh_body_posed'] = mesh_body_posed_no_orient item['mesh_body_posed'] = mesh_posed_body_no_orient
if 'mesh_posed' in self.signature: if 'mesh_posed' in self.signature:
mesh_posed = readPC2Frame(os.path.join(self.folder, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) mesh_posed = np.load(os.path.join(self.folder, 'mesh_posed_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed_no_orient'].astype(np.float32)
item['mesh_posed'] = np.zeros((14475, 3), dtype=np.float32) item['mesh_posed'] = mesh_posed
item['mesh_posed'][mesh_mask] = mesh_posed
if 'mesh_unposed' in self.signature: if 'mesh_unposed' in self.signature:
mesh_unposed_no_shape = readPC2Frame(os.path.join(self.folder, 'mesh_unposed_no_shape', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) mesh_unposed_no_shape = np.load(os.path.join(self.folder, 'mesh_unposed_no_shape', frame.id + separator + frame.frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
item['mesh_unposed'] = np.zeros((14475, 3), dtype=np.float32) item['mesh_unposed'] = mesh_unposed_no_shape
item['mesh_unposed'][mesh_mask] = mesh_unposed_no_shape if 'mesh_curvature' in self.signature:
if 'mesh_normals' in self.signature: mesh_curvature = np.load(os.path.join(self.folder, 'mesh_curvature', frame.id + separator + frame.frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
mesh_normals = readPC2Frame(os.path.join(self.folder, 'mesh_normals', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) item['mesh_curvature'] = mesh_curvature.astype(np.int32)
item['mesh_normals'] = np.zeros((14475, 3), dtype=np.float32)
item['mesh_normals'][mesh_mask] = mesh_normals
for res in self.resolutions: for res in self.resolutions:
res_string = '' res_string = ''
if res != 512: if res != 512:
res_string = '_' + str(res) res_string = '_' + str(res)
if 'uv_normals' + res_string in self.signature: if 'uv_curvature' + res_string in self.signature:
uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_normals = uv_normals['uv_normals'] uv_curvature = np.load(os.path.join(self.folder, 'uv_curvature' + res_string, frame.id + separator + frame.frame_nr + '.npz'));
item['uv_normals' + res_string] = uv_normals if res_string != '':
uv_curvature = uv_curvature['uv_curvature']
else:
uv_curvature = uv_curvature['curvature_uv']
uv_curvature = torch.nn.functional.one_hot(torch.tensor(uv_curvature, device=self.device, dtype=torch.long), 7).float().to(self.device)
item['uv_curvature' + res_string] = uv_curvature
if 'uv_unposed' + res_string in self.signature: if 'uv_unposed' + res_string in self.signature:
uv_unposed_no_shape = np.load(os.path.join(self.folder, 'uv_unposed_no_shape' + res_string, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed'] uv_unposed_no_shape = np.load(os.path.join(self.folder, 'uv_unposed_no_shape' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']
item['uv_unposed' + res_string] = uv_unposed_no_shape item['uv_unposed' + res_string] = uv_unposed_no_shape
if 'uv_body_posed' + res_string in self.signature: if 'uv_body_posed' + res_string in self.signature:
uv_body_posed_no_orient = np.load(os.path.join(self.folder, 'uv_body_posed_no_orient' + res_string, f'{frame.id}_{frame.frame_nr}.npz')); uv_body_posed_no_orient = uv_body_posed_no_orient['uv_body_posed'] uv_posed_body_no_orient = np.load(os.path.join(self.folder, uv_posed_body_no_orient_folder_name + res_string, frame.id + separator + frame.frame_nr + '.npz'))
item['uv_body_posed' + res_string] = uv_body_posed_no_orient uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']
item['uv_body_posed' + res_string] = uv_posed_body_no_orient
if 'uv_posed' + res_string in self.signature: if 'uv_posed' + res_string in self.signature:
uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_posed = uv_posed['uv_posed'] uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient']
item['uv_posed' + res_string] = uv_posed item['uv_posed' + res_string] = uv_posed
data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor or type(value) == int) else value) for key, value in item.items()} data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in item.items()}
for key, value in data.items(): for key, value in data.items():
if key.startswith('uv_') and (value.shape[-1] == 3 or value.shape[-1] == 7): if key.startswith('uv_') and (value.shape[-1] == 3 or value.shape[-1] == 7):
@@ -226,93 +214,38 @@ class Dataset(Dataset):
self.cache[frame] = data self.cache[frame] = data
data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in data.items()} data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in data.items()}
data = { **data, **self.static_cache[frame] } data = { **data, **self.static_cache[frame.id] }
return data return data
if __name__ == '__main__': if __name__ == '__main__':
from torch.utils.data import DataLoader from torch.utils.data import DataLoader
import plotly.graph_objects as go
import plotly.express as ex
from plotly.subplots import make_subplots
from custom_utils.custom_utils import _maskFaces
F = np.load('smpl/smpl_faces.npy')
F_skirt = np.load('smpl/smpl_faces_skirt.npy')
device = "cuda" if torch.cuda.is_available() else "cpu" device = "cuda" if torch.cuda.is_available() else "cpu"
signature = { 'pose_shape', signature = { 'id',
'frame_nr',
'pose_shape',
'uv_posed', 'uv_posed',
'uv_body_posed', 'uv_body_posed',
'mesh_body_posed', 'mesh_body_posed',
'mesh_posed', 'mesh_posed',
'mesh_unposed', 'mesh_unposed',
'uv_unposed', 'uv_unposed',
'uv_normals', 'uv_curvature',
'mesh_normals', 'mesh_curvature',
'uv_static', 'uv_static',
'mesh_static', 'mesh_static',
'mesh_mask', 'mesh_mask',
'uv_mask', 'uv_mask',
'uv_shaped_body',
} }
dataset = Dataset('/data-local/127/cloth3d_processed', signature, resolutions=[512], device=device, frames_file='frames_train.txt') dataset = Dataset('/data-local/127/train', signature, resolutions=[512], device=device, frames_file='frames_jumpsuit.txt')
train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False) train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
for batch in train_dataloader: for batch in train_dataloader:
batch batch
b = 0
mesh_mask = batch['mesh_mask'][b].cpu().numpy()
faces = F_skirt if batch['garment'][b] in ['Skirt', 'Dress'] else F
_, faces = _maskFaces(np.zeros((14475, 3)), faces, mesh_mask)
mesh_static = batch['mesh_static'][b].cpu().numpy()[mesh_mask]
mesh_posed = batch['mesh_posed'][b].cpu().numpy()[mesh_mask]
mesh_unposed = batch['mesh_unposed'][b].cpu().numpy()[mesh_mask]
mesh_body_posed = batch['mesh_body_posed'][b].cpu().numpy()
fig = make_subplots(rows=2, cols=2,
specs=[[{'type': 'mesh3d'}, {'type': 'mesh3d'}],
[{'type': 'mesh3d'}, {'type': 'mesh3d'}]])
fig.add_trace(go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,2], z=mesh_static[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=1)
fig.add_trace(go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=2)
fig.add_trace(go.Mesh3d(x=mesh_unposed[:,0], y=mesh_unposed[:,2], z=mesh_unposed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=2, col=1)
fig.add_trace(go.Mesh3d(x=mesh_body_posed[:,0], y=mesh_body_posed[:,2], z=mesh_body_posed[:,1], i=F[:,0], j=F[:,1], k=F[:,2], opacity=1), row=2, col=2)
fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], autorange=False),
yaxis=dict(range=[-1, 1], autorange=False),
zaxis=dict(range=[-1, 1], autorange=False),
aspectmode='cube'),
scene2=dict(xaxis=dict(range=[-1, 1], autorange=False),
yaxis=dict(range=[-1, 1], autorange=False),
zaxis=dict(range=[-1, 1], autorange=False),
aspectmode='cube'),
scene3=dict(xaxis=dict(range=[-1, 1], autorange=False),
yaxis=dict(range=[-1, 1], autorange=False),
zaxis=dict(range=[-1, 1], autorange=False),
aspectmode='cube'),
scene4=dict(xaxis=dict(range=[-1, 1], autorange=False),
yaxis=dict(range=[-1, 1], autorange=False),
zaxis=dict(range=[-1, 1], autorange=False),
aspectmode='cube'))
fig.show()
uv_mask = batch['uv_mask'][b].cpu().numpy()
uv_mask = np.repeat(uv_mask[np.newaxis, :, :], 3, axis=0)
uv_static = (batch['uv_static'][b].cpu().numpy() + 1) / 2
uv_posed = (batch['uv_posed'][b].cpu().numpy() + 1) / 2
uv_unposed = (batch['uv_unposed'][b].cpu().numpy() + 1) / 2
uv_body_posed = (batch['uv_body_posed'][b].cpu().numpy() + 1) / 2
uv_normals = (batch['uv_normals'][b].cpu().numpy() + 1.5) / 3
img = np.concatenate([uv_mask, uv_static, uv_posed, uv_unposed, uv_body_posed, uv_normals], axis=2)
img = np.transpose(img, (1, 2, 0))
ex.imshow(img).show()
pass pass
pass pass
@@ -341,7 +274,7 @@ if __name__ == '__main__':
# for i, batch in enumerate(dataset): # for i, batch in enumerate(dataset):
# print(i) # print(i)
# batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic)) # batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
# ex.imshow(batch['uv_normals'][0].numpy().argmax(-1) / 6.0).show() # ex.imshow(batch['uv_curvature'][0].numpy().argmax(-1) / 6.0).show()
# sum_uv += batch['uv_unposed'].numpy().sum(axis=0) # sum_uv += batch['uv_unposed'].numpy().sum(axis=0)
# sum_mask += batch['uv_mask'].numpy().sum(axis=0) # sum_mask += batch['uv_mask'].numpy().sum(axis=0)
data/investigate.py
+14 -63
View File
@@ -4,26 +4,11 @@ import random
import numpy as np import numpy as np
import plotly.graph_objects as go import plotly.graph_objects as go
from pymeshfix import MeshFix from pymeshfix import MeshFix
import trimesh
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from data.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo, read_posed_body_frame, read_posed_body, read_posed_garment, read_curvatures, read_curvature_frame, writeOBJ from data.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo, read_posed_body_frame, read_posed_body, read_posed_garment, read_curvatures, read_curvature_frame, writeOBJ
# from custom_utils.custom_utils import _maskFaces from custom_utils.custom_utils import _maskFaces
def _maskFaces(V, F, mask, return_unfixed=False):
V_mask = mask
F_mask = np.take(V_mask, F.reshape(-1), axis=0)
F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
V_ind = np.zeros(V_mask.shape)
V_ind[V_mask] = np.arange(V_mask.sum())
F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
V_new = V[V_mask, :]
if return_unfixed:
return V_new, F_new, F[F_mask, :]
else:
return V_new, F_new
def get_frames(folder, frames_file): def get_frames(folder, frames_file):
frames = [] frames = []
@@ -148,7 +133,7 @@ def check(folder, frame):
v = mesh_posed[mask,:] v = mesh_posed[mask,:]
is_skirt = (garment == 'Dress' or garment == 'Skirt') is_skirt = (garment == 'Dress' or garment == 'Skirt')
skirt_mod = '_skirt' if is_skirt else '' skirt_mod = '_skirt' if is_skirt else ''
F = np.load(f'smpl/smpl_faces{skirt_mod}.npy').astype(np.int32) F = np.load(f'custom_utils/smpl_faces{skirt_mod}.npy').astype(np.int32)
_, faces, unfixed = _maskFaces(np.zeros((14475,3)), F, mask, True) _, faces, unfixed = _maskFaces(np.zeros((14475,3)), F, mask, True)
meshfix = MeshFix(v, faces) meshfix = MeshFix(v, faces)
@@ -195,22 +180,7 @@ def check(folder, frame):
faces_3 = mask_faces(F, mask_2, mask_2_raw, repeated, mesh_posed) faces_3 = mask_faces(F, mask_2, mask_2_raw, repeated, mesh_posed)
# writeOBJ('temp/investigate_3.obj', v_2, faces_3) # writeOBJ('temp/investigate_3.obj', v_2, faces_3)
# plot(v_2, faces_3) plot(v_2, faces_3)
# save stl
# mesh = trimesh.Trimesh(vertices=v_2, faces=faces_3, process=False, validate=False)
# mesh.export('temp/investigate_00761.stl')
# # read stl file
# mesh = trimesh.load('temp/fixed_00761.stl', process=False, validate=False)
# mesh = trimesh.load('temp/fixed_00761.ply', process=False, validate=False)
# plotly plot of mesh with v_2 vertices and faces from trimesh mesh
# faces_4 = faces_3 #np.asarray(mesh.faces)
# go.Figure(data = [go.Mesh3d(x=v_2[:, 0], y=v_2[:, 1], z=v_2[:, 2], i=faces_4[:, 0], j=faces_4[:, 1], k=faces_4[:, 2],)
# ]).show()
# plot repeated # plot repeated
@@ -219,9 +189,9 @@ def check(folder, frame):
r.extend(rr) r.extend(rr)
x = mesh_posed x = mesh_posed
y = mesh_posed[r,:] y = mesh_posed[r,:]
# go.Figure(data = [go.Scatter3d(x=x[:, 0], y=x[:, 1], z=x[:, 2], marker=dict(size=3, line=dict(width=0),), mode='markers'), go.Figure(data = [go.Scatter3d(x=x[:, 0], y=x[:, 1], z=x[:, 2], marker=dict(size=3, line=dict(width=0),), mode='markers'),
# go.Scatter3d(x=y[:, 0], y=y[:, 1], z=y[:, 2], marker=dict(size=5, line=dict(width=0),), mode='markers'), go.Scatter3d(x=y[:, 0], y=y[:, 1], z=y[:, 2], marker=dict(size=5, line=dict(width=0),), mode='markers'),
# ]).show() ]).show()
# unq, idx, count = np.unique(v, axis=0, return_counts=True, return_index=True) # unq, idx, count = np.unique(v, axis=0, return_counts=True, return_index=True)
@@ -237,49 +207,30 @@ def check(folder, frame):
# faces_4 = mask_faces(faces.astype(np.int64), mask_3, mask_2_raw, repeated, v) # faces_4 = mask_faces(faces.astype(np.int64), mask_3, mask_2_raw, repeated, v)
# plot(v[mask_3,:], faces_4) # plot(v[mask_3,:], faces_4)
return mask_2_raw, faces_3
pass pass
def remask(folder, frame, mask, out_folder):
id, frame_nr = frame.split('_')
frame_nr = int(frame_nr)
if frame_nr == 0:
mesh_static = np.load(os.path.join(folder, 'mesh_static', f'{id}.npz'))['mesh_static']
mesh_static_new = np.zeros((14475,3), dtype=np.float32)
mesh_static_new[mask,:] = mesh_static[mask,:]
np.savez_compressed(os.path.join(out_folder, 'mesh_static', f'{id}.npz'), mesh_static=mesh_static_new)
# mesh_posed = read_posed_garment_frame(folder, id, garment, frame_nr, mask=None, shape=None, pose=info['poses'][:,frame_nr], deorient=True, unshape=None)
mesh_posed = np.load(os.path.join(folder, 'mesh_posed_no_orient', f'{id}_{frame_nr}.npz'))['mesh_posed_no_orient']
mesh_posed_new = np.zeros((14475,3), dtype=np.float32)
mesh_posed_new[mask,:] = mesh_posed[mask,:]
np.savez_compressed(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{frame_nr}.npz'), mesh_posed_no_orient=mesh_posed)
if __name__ == "__main__": if __name__ == "__main__":
folder = '/data-local/127/train' folder = '/home/workspace/data/dataset/cloth3d/train'
frames = [f'00761_{i}' for i in range(300)] frames = get_frames(folder, 'frames_all.txt')
root = '/'.join(folder.split('/')[:-1]) root = '/'.join(folder.split('/')[:-1])
type = folder.split('/')[-1] type = folder.split('/')[-1]
use_no_orient = False use_no_orient = False
no_shape = False no_shape = False
is_skirt = False is_skirt = True
skirt_modifier = '_skirt' if is_skirt else '' skirt_modifier = '_skirt' if is_skirt else ''
orient_mod = '_no_orient' if use_no_orient else '' orient_mod = '_no_orient' if use_no_orient else ''
shape_mod = '_no_shape' if no_shape else '' shape_mod = '_no_shape' if no_shape else ''
mask = np.load('/data-net/dataset/00761/mesh_mask/00761.npy')
faces = np.load('/data-net/dataset/00761/faces/00761.npy')
for frame in frames: while True:
# frame = random.choice(frames) frame = random.choice(frames)
# frame = '00761 0 Tshirt' frame = '05936 0 Dress'
remask(folder, frame, mask, '/data-net/dataset/00761') check(folder, frame)
pass pass
data/preprocessing.py
+5 -453
View File
@@ -1,39 +1,22 @@
import os import os
import sys import sys
import argparse
parser = argparse.ArgumentParser()
# parser.add_argument("-f", "--folder", help = "end of range")
parser.add_argument("-s", "--start", help = "start of range", type=int, default=0)
parser.add_argument("-e", "--end", help = "end of range", type=int, default=1)
parser.add_argument("-g", "--gpu", help = "gpu id")
# parser.add_argument("-n", "--name", help = "end of range")
args = parser.parse_args()
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
import math import math
import shutil import shutil
import warnings import warnings
import tables as tb import tables as tb
import numpy as np import numpy as np
import scipy.io as sio import scipy.io as sio
import trimesh as tm
from tqdm import tqdm from tqdm import tqdm
from tables import NaturalNameWarning from tables import NaturalNameWarning
warnings.filterwarnings('ignore', category=NaturalNameWarning) warnings.filterwarnings('ignore', category=NaturalNameWarning)
import plotly.graph_objects as go
import torch
from scipy.spatial.transform import Rotation as R
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from data.cloth3d_io import readPC2, loadInfo, writePC2, readOBJ, quads2tris from data.cloth3d_io import readPC2, loadInfo
from data.utils import Frame from data.utils import Frame
from smpl.smpl_torch import SMPL
from custom_utils.custom_utils import _maskFaces
@@ -149,443 +132,12 @@ def create_extras_file(folder):
h5file.close() h5file.close()
def deorient(verts, pose):
rodr = smpl.rodrigues(torch.tensor(pose.T, device=smpl.device).reshape(-1, 1, 3)).reshape(verts.shape[0], -1, 3, 3).cpu().numpy()
rotated = np.zeros_like(verts)
for frame in range(verts.shape[0]):
r = R.from_matrix(rodr[frame, 0, :, :])
rotated[frame] = r.apply(verts[frame], inverse=True)
return rotated
def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder):
ids = os.listdir(folder)
os.makedirs(os.path.join(out_folder, 'mesh_posed_no_orient'), exist_ok=True)
os.makedirs(os.path.join(out_folder, 'mesh_mask'), exist_ok=True)
for id in tqdm(ids):
# id = '01467'
garment_dir = os.path.join(folder, id)
if not os.path.isdir(garment_dir):
continue
garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
info = loadInfo(os.path.join(folder, id, 'info.mat'))
for garment in garments:
if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
continue
verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
reg_data = sio.loadmat(os.path.join(folder_reg, id, garment + '_weights_fixed.mat'))
verts = (verts[:, reg_data['F'], :] * reg_data['W'][None, :, :, None]).sum(axis=2)
mask_idx = sio.loadmat(os.path.join(folder_reg, id, garment + '_mask.mat'))['mask_idx'].reshape((-1)) - 1
mask = np.zeros((14475), dtype=bool)
mask[mask_idx] = True
verts = deorient(verts, info['poses'])
# # plot as scatter plot with plotly
# import plotly.graph_objects as go
# frame = 5
# fig = go.Figure(data=[go.Scatter3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], mode='markers', marker=dict(size=1))])
# fig.show()
writePC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), verts, True)
np.save(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy'), mask)
def create_pose_shape(folder, out_folder):
ids = os.listdir(folder)
os.makedirs(os.path.join(out_folder, 'pose'), exist_ok=True)
os.makedirs(os.path.join(out_folder, 'shape'), exist_ok=True)
for id in tqdm(ids):
garment_dir = os.path.join(folder, id)
if not os.path.isdir(garment_dir):
continue
info = loadInfo(os.path.join(folder, id, 'info.mat'))
shape = np.concatenate([info['shape'], np.asarray([info['gender']])]).astype(np.float32)
np.save(os.path.join(out_folder, 'pose', f'{id}.npy'), info['poses'].T.astype(np.float32))
np.save(os.path.join(out_folder, 'shape', f'{id}.npy'), shape)
def create_mesh_static(folder, out_folder):
ids = os.listdir(folder)
os.makedirs(os.path.join(out_folder, 'mesh_static'), exist_ok=True)
for id in tqdm(ids):
garment_dir = os.path.join(folder, id)
if not os.path.isdir(garment_dir):
continue
garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
for garment in garments:
if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
continue
verts_static, _, _, _ = readOBJ(os.path.join(garment_dir, f'{garment}.obj'))
reg_data = sio.loadmat(os.path.join(folder_reg, id, garment + '_weights_fixed.mat'))
verts_static = (verts_static[reg_data['F'], :] * reg_data['W'][:, :, None]).sum(axis=1).astype(np.float32)
# rotate 90 degrees around x axis
verts_static = np.matmul(verts_static, np.array([[1,0,0],[0,0,-1],[0,1,0]]))
# fig = go.Figure(data=[go.Scatter3d(x=verts_static[:,0], y=verts_static[:,2], z=verts_static[:,1], mode='markers', marker=dict(size=1))])
# fig.update_layout(scene_aspectmode='cube')
# fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
# fig.show()
np.save(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.npy'), verts_static.astype(np.float32))
def create_mesh_body_posed_no_orient(folder, out_folder):
ids = os.listdir(folder)
os.makedirs(os.path.join(out_folder, 'mesh_body_posed_no_orient'), exist_ok=True)
for id in tqdm(ids):
garment_dir = os.path.join(folder, id)
if not os.path.isdir(garment_dir):
continue
info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
pose, shape, gender = info['poses'], info['shape'], info['gender']
pose = np.transpose(pose, [1, 0]).astype(np.float32)
pose[:, :3] = 0
beta = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32)
T = smpl(torch.tensor(pose, device=smpl.device), torch.tensor(np.tile(beta[np.newaxis, :], [pose.shape[0], 1]), device=smpl.device))[0]
T = T.cpu().numpy()
writePC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), T, True)
# verts = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V']
# # plot as scatter plot with plotly
# import plotly.graph_objects as go
# frame = 5
# fig = go.Figure(data=[go.Scatter3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], mode='markers', marker=dict(size=1))])
# fig.show()
pass
def calculate_normals(vertices, faces):
normals = []
for i in range(vertices.shape[0]):
mesh = tm.Trimesh(vertices=vertices[i], faces=faces, process=False, validate=False)
normals.append(mesh.vertex_normals)
return np.stack(normals)
def create_mesh_normals(folder, out_folder):
ids = os.listdir(folder)
os.makedirs(os.path.join(out_folder, 'mesh_normals'), exist_ok=True)
# F = np.load('smpl/smpl_faces.npy')
# F_skirt = np.load('smpl/smpl_faces_skirt.npy')
for id in tqdm(ids):
garment_dir = os.path.join(folder, id)
if not os.path.isdir(garment_dir):
continue
garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
for garment in garments:
if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
continue
verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
verts = deorient(verts, info['poses'])
_, faces, _, _ = readOBJ(os.path.join(garment_dir, garment + '.obj'))
faces = quads2tris(faces)
normals = calculate_normals(verts, faces.astype(np.int32))
reg_data = sio.loadmat(os.path.join(folder_reg, id, garment + '_weights_fixed.mat'))
normals = (normals[:, reg_data['F'], :] * reg_data['W'][None, :, :, None]).sum(axis=2)
writePC2(os.path.join(out_folder, 'mesh_normals', f'{id}_{garment}.pc16'), normals, True)
# normals = readPC2(os.path.join(out_folder, 'mesh_normals', f'{id}_{garment}.pc16'), True)['V']
# colors = np.clip(np.abs(normals), 0, 1)
# verts = (verts[:, reg_data['F'], :] * reg_data['W'][None, :, :, None]).sum(axis=2)
# # plot as scatter plot with plotly
# import plotly.graph_objects as go
# frame = 5
# fig = go.Figure(data=[go.Scatter3d(
# x=verts[frame,:,0],
# y=verts[frame,:,2],
# z=verts[frame,:,1],
# mode='markers',
# marker=dict(
# size=1,
# color=normals[frame,:,:],
# colorscale='Viridis',
# opacity=0.8
# )
# )])
# moved_verts = verts + 0.05 * normals
# # Create the scatter plot for the moved vertices
# fig.add_trace(go.Scatter3d(
# x=moved_verts[frame,:,0],
# y=moved_verts[frame,:,2],
# z=moved_verts[frame,:,1],
# mode='markers',
# marker=dict(
# size=1,
# color='red',
# opacity=0.8
# )
# ))
# fig.show()
# # create a plotly mesh plot for verts and faces
# fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)])
# fig.show()
pass
def verify_registration_ans_mask(folder, out_folder):
ids = os.listdir(folder)
for id in ids:
garment_dir = os.path.join(folder, id)
if not os.path.isdir(garment_dir):
continue
garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
for garment in garments:
vertices = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V']
mask = np.load(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy'))
if mask.sum() == vertices.shape[1]:
print(id, garment)
pass
def create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=0, end=1):
ids = sorted(os.listdir(folder))[start:end]
os.makedirs(os.path.join(out_folder, 'mesh_unposed_no_shape'), exist_ok=True)
F = np.load('smpl/smpl_faces.npy')
# F_skirt = np.load('smpl/smpl_faces_skirt.npy')
for id in tqdm(ids):
garment_dir = os.path.join(folder, id)
if not os.path.isdir(garment_dir):
continue
garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
# info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
pose = np.load(os.path.join(out_folder, 'pose', f'{id}.npy'))
shape = np.tile(np.load(os.path.join(out_folder, 'shape', f'{id}.npy'))[None,:], [pose.shape[0],1])
pose = torch.tensor(pose, dtype=torch.float32, device='cuda')
shape = torch.tensor(shape, dtype=torch.float32, device='cuda')
for garment in garments:
if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
continue
mesh_static = np.load(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.npy'))
mesh_posed = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V']
mask = np.load(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy'))
# mask_idx = np.where(mask)[0]
pose[:,:3] = 0
# pose = pose[:10]
# shape = shape[:10]
# mesh_posed = mesh_posed[:10]
unposed = np.zeros((14475, 3), dtype=np.float32)
unposed[mask] = mesh_static
unposed = torch.tensor(np.tile(unposed[None,:,:], [mesh_posed.shape[0],1,1]), dtype=torch.float32, device='cuda')
unposed.requires_grad = True
mesh_posed = torch.tensor(mesh_posed, dtype=torch.float32, device='cuda')
optimizer = torch.optim.Adam([unposed], lr=0.01)
# verts_1 = unposed[0].cpu().detach().numpy()[mask]
# _, faces = _maskFaces(np.zeros((14475, 3)), F, mask)
# fig = go.Figure(data=[go.Mesh3d(x=verts_1[:,0], y=verts_1[:,2], z=verts_1[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)])
# fig.update_layout(scene = dict(
# xaxis = dict(range=[-1,1],),
# yaxis = dict(range=[-1,1],),
# zaxis = dict(range=[-1,1],),),)
# fig.update_layout(scene_aspectmode='cube')
# fig.show()
# verts_2 = mesh_posed[0].cpu().detach().numpy()
# fig = go.Figure(data=[go.Mesh3d(x=verts_2[:,0], y=verts_2[:,2], z=verts_2[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)])
# fig.update_layout(scene = dict(
# xaxis = dict(range=[-1,1],),
# yaxis = dict(range=[-1,1],),
# zaxis = dict(range=[-1,1],),),)
# fig.update_layout(scene_aspectmode='cube')
# fig.show()
# zero_body = smpl(torch.zeros_like(pose), torch.zeros_like(shape))[0]
# verts_3 = smpl(pose, shape, zero_body)[0][0].cpu().detach().numpy()
# # verts_3 = zero_body[0].cpu().detach().numpy()
# faces_3 = F
# fig = go.Figure(data=[go.Mesh3d(x=verts_3[:,0], y=verts_3[:,2], z=verts_3[:,1], i=faces_3[:,0], j=faces_3[:,1], k=faces_3[:,2], color='lightpink', opacity=1)])
# fig.update_layout(scene = dict(
# xaxis = dict(range=[-1,1],),
# yaxis = dict(range=[-1,1],),
# zaxis = dict(range=[-1,1],),),)
# fig.update_layout(scene_aspectmode='cube')
# fig.show()
for i in range(5000):
optimizer.zero_grad()
pred = smpl(pose, shape, unposed)[0]
loss : torch.Tensor = ((pred[:,mask,:] - mesh_posed)**2).sum()
# Backward pass
loss.backward()
optimizer.step()
with torch.inference_mode():
euclidean_error = torch.sqrt(((pred[:,mask,:] - mesh_posed)**2).sum(-1)).mean()
# Print the loss
if i % 50 == 0:
print(f"it: {i} loss: {loss:.8f} euclidean_error: {euclidean_error:.8f}")
if euclidean_error < 0.0000001:
break
if euclidean_error < 0.0000005 and i > 2000:
break
if euclidean_error < 0.000001 and i > 4000:
break
unposed = unposed.detach().cpu().numpy()
gender = int(shape[0].cpu().numpy()[-1] > 0)
shape_offsets = torch.tensordot(shape[:,:10], smpl.shapedirs[gender], dims=[[1], [2]])
shape_offsets = shape_offsets.cpu().numpy()[:,smpl.to_keep,:]
unshaped = (unposed - shape_offsets)[:,mask,:]
writePC2(os.path.join(out_folder, 'mesh_unposed_no_shape', f'{id}_{garment}.pc16'), unshaped, True)
# verts = unposed[:,mask,:]
# verts = unshaped
# _, faces = _maskFaces(np.zeros((14475, 3)), F, mask)
# frame = 9
# fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)])
# fig.update_layout(scene = dict(
# xaxis = dict(range=[-1,1],),
# yaxis = dict(range=[-1,1],),
# zaxis = dict(range=[-1,1],),),)
# fig.update_layout(scene_aspectmode='cube')
# fig.show()
# verts = pred.detach().cpu().numpy()[:,mask,:]
# frame = 9
# fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)])
# fig.update_layout(scene = dict(
# xaxis = dict(range=[-1,1],),
# yaxis = dict(range=[-1,1],),
# zaxis = dict(range=[-1,1],),),)
# fig.update_layout(scene_aspectmode='cube')
# fig.show()
pass
def create_frames_file(folder, out_folder, split=0.9):
ids = os.listdir(folder)
frames_non_skirt = []
frames_skirt = []
ids_non_skirt = set()
ids_skirt = set()
for id in tqdm(ids):
garment_dir = os.path.join(folder, id)
if not os.path.isdir(garment_dir):
continue
info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
frames = info['poses'].shape[1]
assert info['poses'].shape[0] == 72
for garment in garments:
for i in range(frames):
if garment == 'Skirt' or garment == 'Dress':
frames_skirt.append(Frame(id, i, garment))
ids_skirt.add(id)
else:
frames_non_skirt.append(Frame(id, i, garment))
ids_non_skirt.add(id)
nr_train_non_skirt = int(len(ids_non_skirt) * split)
nr_train_skirt = int(len(ids_skirt) * split)
ids_skirt = list(ids_skirt)
ids_non_skirt = list(ids_non_skirt)
with open(os.path.join(out_folder, 'frames_train.txt'), 'w') as f:
for frame in frames_non_skirt:
if frame.id in ids_non_skirt[:nr_train_non_skirt]:
f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
with open(os.path.join(out_folder, 'frames_train_skirt.txt'), 'w') as f:
for frame in frames_skirt:
if frame.id in ids_skirt[:nr_train_skirt]:
f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
with open(os.path.join(out_folder, 'frames_valid.txt'), 'w') as f:
for frame in frames_non_skirt:
if frame.id in ids_non_skirt[nr_train_non_skirt:]:
f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
with open(os.path.join(out_folder, 'frames_valid_skirt.txt'), 'w') as f:
for frame in frames_skirt:
if frame.id in ids_skirt[nr_train_skirt:]:
f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
if __name__ == '__main__': if __name__ == '__main__':
folder = '/home/workspace/data/dataset/cloth3d/train'
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')
# clean_dataset(folder) # clean_dataset(folder)
# generate_frames_file(folder, 'frames_all.txt') # generate_frames_file(folder, 'frames_all.txt')
# split_data(folder, 0.9) # split_data(folder, 0.9)
# create_extras_file(folder) create_extras_file(folder)
# create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder)
# verify_registration_ans_mask(folder, out_folder)
# create_pose_shape(folder, out_folder)
# create_mesh_static(folder, out_folder)
# create_mesh_body_posed_no_orient(folder, out_folder)
# create_mesh_normals(folder, out_folder)
# create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end)
create_frames_file(folder, out_folder)
data/registration.py
-461
View File
@@ -1,461 +0,0 @@
import os
import sys
import random
import numpy as np
import plotly.graph_objects as go
import torch
from torch.nn.functional import mse_loss
from pytorch3d.structures import Meshes
from pytorch3d.loss import mesh_laplacian_smoothing, mesh_edge_loss, mesh_normal_consistency, chamfer_distance
from pymeshfix import MeshFix
# os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from data.cloth3d_io import (
read_garment_mask,
read_garment_template,
read_posed_garment_frame,
loadInfo,
read_posed_body_frame,
read_posed_body,
read_posed_garment,
read_curvatures,
read_curvature_frame,
writeOBJ,
readOBJ,
quads2tris
)
from smpl.smpl_torch import SMPL
# from custom_utils.custom_utils import _maskFaces
device = 'cuda'
smpl = SMPL(device, 'smpl')
import torch.nn.functional as F
from pytorch3d.structures import Meshes
def find_closest_vertices(mesh1: Meshes, mesh2: Meshes, vertex_indices: torch.Tensor) -> torch.Tensor:
"""
Given two PyTorch3D meshes and a set of vertex indices for mesh1,
find the closest vertices from mesh2 to these vertices.
Args:
mesh1: A PyTorch3D Meshes object representing the first mesh.
mesh2: A PyTorch3D Meshes object representing the second mesh.
vertex_indices: A tensor of shape (N,) representing the indices of the vertices in mesh1.
Returns:
A tensor of shape (N,) representing the indices of the closest vertices in mesh2.
"""
# Extract the vertices of mesh1 and mesh2.
vertices1 = mesh1.verts_packed()
vertices2 = mesh2.verts_packed()
# Extract the vertices of mesh1 that correspond to the given indices.
vertices1_selected = vertices1[vertex_indices]
# Compute the pairwise distances between the selected vertices of mesh1 and all vertices of mesh2.
distances = F.pairwise_distance(vertices1_selected.unsqueeze(1), vertices2.unsqueeze(0), p=2)
# Find the indices of the closest vertices in mesh2.
closest_indices = torch.argmin(distances, dim=1)
return closest_indices
def _maskFaces(V, F, mask, return_unfixed=False):
V_mask = mask
F_mask = np.take(V_mask, F.reshape(-1), axis=0)
F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
V_ind = np.zeros(V_mask.shape)
V_ind[V_mask] = np.arange(V_mask.sum())
F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
V_new = V[V_mask, :]
if return_unfixed:
return V_new, F_new, F[F_mask, :]
else:
return V_new, F_new
def get_frames(folder, frames_file):
frames = []
with open(os.path.join(folder, frames_file), 'r') as f:
for line in f:
frames.append(line.strip())
return frames
def plot(mesh, faces):
fig = go.Figure()
fig.add_trace(go.Mesh3d(
x=mesh[:,0],
y=mesh[:,1],
z=mesh[:,2],
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),)
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
)
helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,1],),
zaxis = dict(nticks=4, range=[-1,1],),)
fig.update_layout(scene = helper,
)
fig.show()
def build_lookup_table(repeated):
# iterate through the list and create a dict from each sub array in the list where each dict has one element of the sub array as key and the other elements as value as a list and repeat for each element
lookup = {}
for i in range(len(repeated)):
for j in range(len(repeated[i])):
if len(repeated[i]) == 1:
continue
lookup[repeated[i][j]] = [repeated[i][k] for k in range(len(repeated[i]))] # if k != j
return lookup
from data.cloth3d_io import readPC2Frame
import scipy.io as sio
import scipy
import scipy.spatial as spatial
def local_read_posed_garment_frame(folder, id, garment, frame, pose=None, deorient=False):
pc16_path = os.path.join(folder, id, garment + '.pc16')
verts = readPC2Frame(pc16_path, frame, True)
map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1))
if deorient:
R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
# R = spatial.transform.Rotation.from_euler([0,0,-zrot])
verts = R.apply(verts)
return verts, map
def local_read_posed_body_frame(folder, id, frame, pose=None, deorient=False):
pc16_path = os.path.join(folder, id, 'Body_posed.pc16')
verts = readPC2Frame(pc16_path, frame, True)
if deorient:
R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
# R = spatial.transform.Rotation.from_euler([0,0,-zrot])
verts = R.apply(verts)
return verts
def edge_loss_custom(x, y, edges):
x_e = x[:,edges[:,0],:] - x[:,edges[:,1],:]
x_e = torch.sum(x_e ** 2, -1)
x_e = torch.sqrt(x_e + 1e-9)
y_e = y[:,edges[:,0],:] - y[:,edges[:,1],:]
y_e = torch.sum(y_e ** 2, -1)
y_e = torch.sqrt(y_e + 1e-9)
d_e = y_e - x_e
loss = torch.mean(d_e ** 2, dim=-1)[0]
return loss
def crotch(mesh, body, crotch_idx, nr=75):
crotch_middle_on_body = crotch_idx
# create kdtree for body and mesh
body_tree = spatial.cKDTree(body)
mesh_tree = spatial.cKDTree(mesh)
# find closest point on mesh from crotch middle on body along the z axis
crotch_middle_on_mesh = find_closest_point_along_z(body[crotch_middle_on_body], mesh)
crotch_body = body_tree.query(body[crotch_middle_on_body], k=nr)[1]
crotch_mesh = mesh_tree.query(mesh[crotch_middle_on_mesh], k=nr)[1]
body_verts = body[crotch_body]
mesh_verts = mesh[crotch_mesh] - (mesh[crotch_middle_on_mesh] - body[crotch_middle_on_body])
mesh_verts = spatial.cKDTree(mesh_verts)
closest_mesh_verts = mesh_verts.query(body_verts, k=1)[1]
crotch_mesh = crotch_mesh[closest_mesh_verts]
return crotch_body, crotch_mesh
def find_closest_point_along_z(point, verts, eps=0.01):
min_dist = np.inf
closest_point = None
for i, v in enumerate(verts):
dist = np.linalg.norm(point - v)
if dist < min_dist and np.abs(point[0] - v[0]) < eps and np.abs(point[1] - v[1]) < eps:
min_dist = dist
closest_point = i
return closest_point
def mesh_edge_lengths(mesh: Meshes) -> torch.Tensor:
"""
Given a PyTorch3D Meshes object, compute the lengths of all edges in the mesh.
Args:
mesh: A PyTorch3D Meshes object representing the mesh.
Returns:
A tensor of shape (E,) representing the lengths of all edges in the mesh,
where E is the total number of edges in the mesh.
"""
# Extract the vertices and faces of the mesh.
vertices = mesh.verts_packed()
faces = mesh.faces_packed()
# Compute the pairwise distances between all pairs of vertices that form an edge.
edge_indices = faces.view(-1, 3)
edge_vertices = vertices[edge_indices]
edge_lengths = torch.norm(edge_vertices[:, 1:] - edge_vertices[:, :-1], dim=2)
return edge_lengths
def register(folder, frame):
id, frame_nr, garment = frame.split()
frame_nr = int(frame_nr)
info = loadInfo(os.path.join(folder, id, 'info.mat'))
raw_mask = read_garment_mask(folder, id, garment, return_raw=True)
# convert it to indices
mask = raw_mask > 0
mesh_posed, map = local_read_posed_garment_frame(folder, id, garment, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
body_posed = local_read_posed_body_frame(folder, id, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
template, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj'))
faces = quads2tris(faces)
# boundaries = raw_mask == 2
# # invert map and use it to invert boudaries
# map_inv = np.zeros(map.max()+1, dtype=np.int32)
# map_inv[map] = np.arange(len(map))
# boundaries = boundaries[map_inv]
# # get the boundary indices
# mesh1_boundaries = torch.tensor(np.where(boundaries)[0]).to(device)
# pose = torch.zeros(1, 72).to('cuda')
shape = torch.zeros(1, 11).to(device)
shape[:, :10] = torch.tensor(info['shape']).to(device)
shape[:, -1] = torch.tensor(info['gender']).to(device)
# for star pose
# pose[:, 5] = 3.1415/4
# pose[:,8] = -3.1415/4
STAR_POSE = np.zeros((24, 3), np.float32)
STAR_POSE[1,2] = .15
STAR_POSE[2,2] = -.15
pose = torch.tensor(STAR_POSE.flatten()[None,:]).to(device)
star_body = smpl(pose, shape, remove=True)[0][0].cpu().numpy() - smpl(pose, shape, remove=True)[4][0].cpu().numpy()
# rotate star body with 90 degrees around x axis
R = spatial.transform.Rotation.from_euler('x', np.pi/2)
star_body = R.apply(star_body).astype(np.float32)
is_skirt = (garment == 'Dress' or garment == 'Skirt')
skirt_mod = '_skirt' if is_skirt else ''
F = np.load(f'smpl/smpl_faces{skirt_mod}.npy').astype(np.int32)
##################### mask body #####################
new_crotch = 4729
# star_body = star_body[mask]
# new_crotch = np.zeros((14475), bool)
# new_crotch[4729] = True
# new_crotch = new_crotch[mask]
# new_crotch = np.where(new_crotch)[0][0]
# _, F = _maskFaces(np.zeros((14475, 3)), F, mask)
# plot(template, faces)
# v = mesh_posed[mask,:]
##################### boundaries #####################
# meshfix = MeshFix(template, faces)
# holes = meshfix.extract_holes()
# new_border = []
# for border in np.array(holes.points):
# # find the index of the vertex in the original mesh
# ind = np.argwhere(np.all(template == border, axis=1))
# new_border.extend(ind.ravel())
# mesh1_boundaries = torch.tensor(new_border).to(device)
# np.save('temp/body_boundaries.npy', new_border)
mesh1_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
# meshfix = MeshFix(template, faces)
# holes = meshfix.extract_holes()
# new_border = []
# for border in np.array(holes.points):
# # find the index of the vertex in the original mesh
# ind = np.argwhere(np.all(template == border, axis=1))
# new_border.extend(ind.ravel())
# mesh1_boundaries = torch.tensor(new_border).to(device)
# np.save('temp/body_boundaries.npy', new_border)
# mesh1_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
##################### crotch #####################
crotch_body, crotch_mesh = crotch(template, star_body, crotch_idx=new_crotch, nr=50)
# export star_body as obj
# writeOBJ('temp/star_body.obj', star_body, F)
template = torch.tensor(template).to(device)
template_orig = template.clone()
faces = torch.tensor(faces).to(device)
star_body = torch.tensor(star_body).to(device)
F = torch.tensor(F).to(device)
# verts = template.cpu().numpy()
# star_body = star_body.cpu().numpy()
# faces = faces.detach().cpu().numpy()
# F = F.detach().cpu().numpy()
# fig = go.Figure(data=[
# go.Mesh3d(
# x=verts[:,0], y=verts[:,1], z=verts[:,2],
# i=faces[:,0], j=faces[:,1], k=faces[:,2],
# color='lightpink', opacity=0.50),
# go.Mesh3d(
# x=star_body[:,0], y=star_body[:,1], z=star_body[:,2],
# i=F[:,0], j=F[:,1], k=F[:,2],
# color='lightblue', opacity=0.50),
# go.Scatter3d(
# x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
# mode='markers', marker=dict(size=2, color='red')),
# go.Scatter3d(
# x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
# mode='markers', marker=dict(size=2, color='blue')),
# ])
# # make it cuve and limit to -1 and -1
# fig.update_layout(scene = dict(
# xaxis = dict(range=[-1,1],),
# yaxis = dict(range=[-1,1],),
# zaxis = dict(range=[-1,1],),),)
# fig.show()
# Define the meshes to optimize
mesh1 = Meshes(verts=template[None,:,:], faces=faces[None,:,:])
mesh2 = Meshes(verts=star_body[None,:,:], faces=F[None,:,:])
mesh2_boundaries = find_closest_vertices(mesh1, mesh2, mesh1_boundaries)
avg_edge_length = mesh_edge_lengths(mesh1).mean()
# Define the optimizer and learning rate
mesh1.verts_packed().requires_grad = True
optimizer = torch.optim.Adam([mesh1.verts_packed()], lr=0.0001)
# optimizer = torch.optim.SGD([mesh1.verts_packed()], lr=0.001)
# Define the loss functions
def laplacian_smoothing_loss(mesh):
return mesh_laplacian_smoothing(mesh, method="uniform") # cotcurv uniform cot
def edge_loss(mesh):
return mesh_edge_loss(mesh, target_length=avg_edge_length)
def normal_consistency_loss(mesh):
return mesh_normal_consistency(mesh)
# Optimization loop
for i in range(10000):
optimizer.zero_grad()
# find boundaries with the find holes function?
# Forward pass
mesh1_out = Meshes(verts=[mesh1.verts_packed()], faces=[mesh1.faces_packed()])
ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list()), torch.stack(mesh2.verts_list()))
bound_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,mesh1_boundaries,:], torch.stack(mesh2.verts_list())[:,mesh2_boundaries,:])
# crotch_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:], torch.stack(mesh2.verts_list())[:,crotch_body,:],)
crotch_ch_loss = torch.mean((torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:] - torch.stack(mesh2.verts_list())[:,crotch_body,:])**2, dim=2).mean()
loss: torch.Tensor = ch_loss * 10 + (bound_ch_loss + crotch_ch_loss) * 5
# loss: torch.Tensor = ch_loss + ch_norm_loss + (bound_ch_loss + bound_ch_norm_loss) * 1
loss += laplacian_smoothing_loss(mesh1_out) * 0.05
loss += edge_loss(mesh1_out) * 100
# loss += edge_loss_custom(torch.stack(mesh1_out.verts_list()), template_orig[None,:,:], mesh1.edges_packed()) * 100
loss += normal_consistency_loss(mesh1_out) * 0.1
# Backward pass
loss.backward()
optimizer.step()
# Print the loss
if i % 10 == 0:
print(f"Loss: {loss.item()} chamfer: {ch_loss.item()} ")
# create a plotly plot with the body and the mesh1_out
verts = mesh1_out.verts_packed().detach().cpu().numpy()
star_body = star_body.detach().cpu().numpy()
faces = faces.detach().cpu().numpy()
F = F.detach().cpu().numpy()
mesh1_boundaries = mesh1_boundaries.detach().cpu().numpy()
mesh2_boundaries = mesh2_boundaries.detach().cpu().numpy()
fig = go.Figure(data=[
go.Mesh3d(
x=verts[:,0], y=verts[:,1], z=verts[:,2],
i=faces[:,0], j=faces[:,1], k=faces[:,2],
color='lightpink', opacity=0.80),
go.Mesh3d(
x=star_body[:,0], y=star_body[:,1], z=star_body[:,2],
i=F[:,0], j=F[:,1], k=F[:,2],
color='lightblue', opacity=0.50),
go.Scatter3d(
x=verts[mesh1_boundaries,0], y=verts[mesh1_boundaries,1], z=verts[mesh1_boundaries,2],
mode='markers', marker=dict(size=2, color='red')),
go.Scatter3d(
x=star_body[mesh2_boundaries,0], y=star_body[mesh2_boundaries,1], z=star_body[mesh2_boundaries,2],
mode='markers', marker=dict(size=2, color='blue')),
go.Scatter3d(
x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
mode='markers', marker=dict(size=2, color='red')),
go.Scatter3d(
x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
mode='markers', marker=dict(size=2, color='blue')),
])
# make it cuve and limit to -1 and -1
fig.update_layout(scene = dict(
xaxis = dict(range=[-1,1],),
yaxis = dict(range=[-1,1],),
zaxis = dict(range=[-1,1],),),)
fig.update_layout(scene_aspectmode='cube')
fig.show()
pass
if __name__ == "__main__":
folder = '/data-net/code/cloth3d/train'
# folder = r'D:\Work\code\cloth3d\train'
# frames = [f'00761_{i}' for i in range(300)]
# folder = r'Z:\Work\Huni\Barcelona\Dataset\cloth3d\train'
frames = get_frames(folder, 'frames_all.txt')
root = '/'.join(folder.split('/')[:-1])
type = folder.split('/')[-1]
use_no_orient = False
no_shape = False
is_skirt = False
skirt_modifier = '_skirt' if is_skirt else ''
orient_mod = '_no_orient' if use_no_orient else ''
shape_mod = '_no_shape' if no_shape else ''
# mask = np.load('/data-net/dataset/00761/mesh_mask/00761.npy')
# faces = np.load('/data-net/dataset/00761/faces/00761.npy')
for frame in frames:
frame = random.choice(frames)
# frame = '00761 0 Tshirt'
frame = '01893 256 Jumpsuit'
register(folder, frame)
break
pass
data/registration_post.py
-482
View File
@@ -1,482 +0,0 @@
import os
import sys
import random
import numpy as np
import plotly.graph_objects as go
import torch
from torch.nn.functional import mse_loss
from pytorch3d.structures import Meshes
from pytorch3d.loss import mesh_laplacian_smoothing, mesh_edge_loss, mesh_normal_consistency, chamfer_distance
from pymeshfix import MeshFix
# os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from data.cloth3d_io import (
read_garment_mask,
read_garment_template,
read_posed_garment_frame,
loadInfo,
read_posed_body_frame,
read_posed_body,
read_posed_garment,
read_curvatures,
read_curvature_frame,
writeOBJ,
readOBJ,
quads2tris
)
from smpl.smpl_torch import SMPL
# from custom_utils.custom_utils import _maskFaces
device = 'cuda'
smpl = SMPL(device, 'smpl')
import torch.nn.functional as F
from pytorch3d.structures import Meshes
def find_closest_vertices(mesh1: Meshes, mesh2: Meshes, vertex_indices: torch.Tensor) -> torch.Tensor:
"""
Given two PyTorch3D meshes and a set of vertex indices for mesh1,
find the closest vertices from mesh2 to these vertices.
Args:
mesh1: A PyTorch3D Meshes object representing the first mesh.
mesh2: A PyTorch3D Meshes object representing the second mesh.
vertex_indices: A tensor of shape (N,) representing the indices of the vertices in mesh1.
Returns:
A tensor of shape (N,) representing the indices of the closest vertices in mesh2.
"""
# Extract the vertices of mesh1 and mesh2.
vertices1 = mesh1.verts_packed()
vertices2 = mesh2.verts_packed()
# Extract the vertices of mesh1 that correspond to the given indices.
vertices1_selected = vertices1[vertex_indices]
# Compute the pairwise distances between the selected vertices of mesh1 and all vertices of mesh2.
distances = F.pairwise_distance(vertices1_selected.unsqueeze(1), vertices2.unsqueeze(0), p=2)
# Find the indices of the closest vertices in mesh2.
closest_indices = torch.argmin(distances, dim=1)
return closest_indices
def _maskFaces(V, F, mask, return_unfixed=False):
V_mask = mask
F_mask = np.take(V_mask, F.reshape(-1), axis=0)
F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
V_ind = np.zeros(V_mask.shape)
V_ind[V_mask] = np.arange(V_mask.sum())
F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
V_new = V[V_mask, :]
if return_unfixed:
return V_new, F_new, F[F_mask, :]
else:
return V_new, F_new
def get_frames(folder, frames_file):
frames = []
with open(os.path.join(folder, frames_file), 'r') as f:
for line in f:
frames.append(line.strip())
return frames
def plot(mesh, faces):
fig = go.Figure()
fig.add_trace(go.Mesh3d(
x=mesh[:,0],
y=mesh[:,1],
z=mesh[:,2],
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),)
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
)
helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,1],),
zaxis = dict(nticks=4, range=[-1,1],),)
fig.update_layout(scene = helper,
)
fig.show()
def build_lookup_table(repeated):
# iterate through the list and create a dict from each sub array in the list where each dict has one element of the sub array as key and the other elements as value as a list and repeat for each element
lookup = {}
for i in range(len(repeated)):
for j in range(len(repeated[i])):
if len(repeated[i]) == 1:
continue
lookup[repeated[i][j]] = [repeated[i][k] for k in range(len(repeated[i]))] # if k != j
return lookup
from data.cloth3d_io import readPC2Frame
import scipy.io as sio
import scipy
import scipy.spatial as spatial
def local_read_posed_garment_frame(folder, id, garment, frame, pose=None, deorient=False):
pc16_path = os.path.join(folder, id, garment + '.pc16')
verts = readPC2Frame(pc16_path, frame, True)
map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1))
if deorient:
R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
# R = spatial.transform.Rotation.from_euler([0,0,-zrot])
verts = R.apply(verts)
return verts, map
def local_read_posed_body_frame(folder, id, frame, pose=None, deorient=False):
pc16_path = os.path.join(folder, id, 'Body_posed.pc16')
verts = readPC2Frame(pc16_path, frame, True)
if deorient:
R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
# R = spatial.transform.Rotation.from_euler([0,0,-zrot])
verts = R.apply(verts)
return verts
def edge_loss_custom(x, y, edges):
x_e = x[:,edges[:,0],:] - x[:,edges[:,1],:]
x_e = torch.sum(x_e ** 2, -1)
x_e = torch.sqrt(x_e + 1e-9)
y_e = y[:,edges[:,0],:] - y[:,edges[:,1],:]
y_e = torch.sum(y_e ** 2, -1)
y_e = torch.sqrt(y_e + 1e-9)
d_e = y_e - x_e
loss = torch.mean(d_e ** 2, dim=-1)[0]
return loss
def crotch(mesh, body, crotch_idx, nr=75):
crotch_middle_on_body = crotch_idx
# create kdtree for body and mesh
body_tree = spatial.cKDTree(body)
mesh_tree = spatial.cKDTree(mesh)
# find closest point on mesh from crotch middle on body along the z axis
crotch_middle_on_mesh = find_closest_point_along_z(body[crotch_middle_on_body], mesh)
crotch_body = body_tree.query(body[crotch_middle_on_body], k=nr)[1]
crotch_mesh = mesh_tree.query(mesh[crotch_middle_on_mesh], k=nr)[1]
body_verts = body[crotch_body]
mesh_verts = mesh[crotch_mesh] - (mesh[crotch_middle_on_mesh] - body[crotch_middle_on_body])
mesh_verts = spatial.cKDTree(mesh_verts)
closest_mesh_verts = mesh_verts.query(body_verts, k=1)[1]
crotch_mesh = crotch_mesh[closest_mesh_verts]
return crotch_body, crotch_mesh
def find_closest_point_along_z(point, verts, eps=0.01):
min_dist = np.inf
closest_point = None
for i, v in enumerate(verts):
dist = np.linalg.norm(point - v)
if dist < min_dist and np.abs(point[0] - v[0]) < eps and np.abs(point[1] - v[1]) < eps:
min_dist = dist
closest_point = i
return closest_point
def mesh_edge_lengths(mesh: Meshes) -> torch.Tensor:
"""
Given a PyTorch3D Meshes object, compute the lengths of all edges in the mesh.
Args:
mesh: A PyTorch3D Meshes object representing the mesh.
Returns:
A tensor of shape (E,) representing the lengths of all edges in the mesh,
where E is the total number of edges in the mesh.
"""
# Extract the vertices and faces of the mesh.
vertices = mesh.verts_packed()
faces = mesh.faces_packed()
# Compute the pairwise distances between all pairs of vertices that form an edge.
edge_indices = faces.view(-1, 3)
edge_vertices = vertices[edge_indices]
edge_lengths = torch.norm(edge_vertices[:, 1:] - edge_vertices[:, :-1], dim=2)
return edge_lengths
# try to optimize the already registered garment ??????????
def register(folder, frame):
id, frame_nr, garment = frame.split()
frame_nr = int(frame_nr)
info = loadInfo(os.path.join(folder, id, 'info.mat'))
raw_mask = read_garment_mask(folder, id, garment, return_raw=True)
# convert it to indices
mask = raw_mask > 0
mesh_posed, map = local_read_posed_garment_frame(folder, id, garment, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
body_posed = local_read_posed_body_frame(folder, id, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
template, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj'))
faces = quads2tris(faces)
boundaries = raw_mask == 2
registered_boundaries = torch.tensor(np.where(boundaries[mask])[0]).to(device)
# invert map and use it to invert boudaries
# map_inv = np.zeros(map.max()+1, dtype=np.int32)
# map_inv[map] = np.arange(len(map))
# boundaries = boundaries[map_inv]
# # get the boundary indices
# mesh1_boundaries = torch.tensor(np.where(boundaries)[0]).to(device)
# # pose = torch.zeros(1, 72).to('cuda')
# shape = torch.zeros(1, 11).to(device)
# shape[:, :10] = torch.tensor(info['shape']).to(device)
# shape[:, -1] = torch.tensor(info['gender']).to(device)
# # for star pose
# # pose[:, 5] = 3.1415/4
# # pose[:,8] = -3.1415/4
# STAR_POSE = np.zeros((24, 3), np.float32)
# STAR_POSE[1,2] = .15
# STAR_POSE[2,2] = -.15
# pose = torch.tensor(STAR_POSE.flatten()[None,:]).to(device)
# star_body = smpl(pose, shape, remove=True)[0][0].cpu().numpy() - smpl(pose, shape, remove=True)[4][0].cpu().numpy()
# # rotate star body with 90 degrees around x axis
# R = spatial.transform.Rotation.from_euler('x', np.pi/2)
# star_body = R.apply(star_body).astype(np.float32)
is_skirt = (garment == 'Dress' or garment == 'Skirt')
skirt_mod = '_skirt' if is_skirt else ''
F = np.load(f'smpl/smpl_faces{skirt_mod}.npy').astype(np.int32)
##################### mask body #####################
new_crotch = 4729
# star_body = star_body[mask]
new_crotch = np.zeros((14475), bool)
new_crotch[4729] = True
new_crotch = new_crotch[mask]
new_crotch = np.where(new_crotch)[0][0]
_, F = _maskFaces(np.zeros((14475, 3)), F, mask)
registered = template[map][mask]
# plot(template, faces)
# v = mesh_posed[mask,:]
##################### boundaries #####################
# meshfix = MeshFix(template, faces)
# holes = meshfix.extract_holes()
# new_border = []
# for border in np.array(holes.points):
# # find the index of the vertex in the original mesh
# ind = np.argwhere(np.all(template == border, axis=1))
# new_border.extend(ind.ravel())
# mesh1_boundaries = torch.tensor(new_border).to(device)
# np.save('temp/body_boundaries.npy', new_border)
template_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
# meshfix = MeshFix(template, faces)
# holes = meshfix.extract_holes()
# new_border = []
# for border in np.array(holes.points):
# # find the index of the vertex in the original mesh
# ind = np.argwhere(np.all(template == border, axis=1))
# new_border.extend(ind.ravel())
# mesh1_boundaries = torch.tensor(new_border).to(device)
# np.save('temp/body_boundaries.npy', new_border)
# mesh1_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
##################### crotch #####################
# crotch_body, crotch_mesh = crotch(template, star_body, crotch_idx=new_crotch, nr=50)
# export star_body as obj
# writeOBJ('temp/star_body.obj', star_body, F)
template = torch.tensor(template).to(device)
template_orig = template.clone()
faces = torch.tensor(faces).to(device)
registered = torch.tensor(registered).to(device)
F = torch.tensor(F).to(device)
# verts = template.cpu().numpy()
# star_body = star_body.cpu().numpy()
# faces = faces.detach().cpu().numpy()
# F = F.detach().cpu().numpy()
# fig = go.Figure(data=[
# go.Mesh3d(
# x=verts[:,0], y=verts[:,1], z=verts[:,2],
# i=faces[:,0], j=faces[:,1], k=faces[:,2],
# color='lightpink', opacity=0.50),
# go.Mesh3d(
# x=star_body[:,0], y=star_body[:,1], z=star_body[:,2],
# i=F[:,0], j=F[:,1], k=F[:,2],
# color='lightblue', opacity=0.50),
# go.Scatter3d(
# x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
# mode='markers', marker=dict(size=2, color='red')),
# go.Scatter3d(
# x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
# mode='markers', marker=dict(size=2, color='blue')),
# ])
# # make it cuve and limit to -1 and -1
# fig.update_layout(scene = dict(
# xaxis = dict(range=[-1,1],),
# yaxis = dict(range=[-1,1],),
# zaxis = dict(range=[-1,1],),),)
# fig.show()
# Define the meshes to optimize
mesh1 = Meshes(verts=registered[None,:,:], faces=F[None,:,:])
mesh2 = Meshes(verts=template[None,:,:], faces=faces[None,:,:])
mesh1_boundaries = registered_boundaries
mesh2_boundaries = template_boundaries
# mesh2_boundaries = find_closest_vertices(mesh1, mesh2, mesh1_boundaries)
avg_edge_length = mesh_edge_lengths(mesh1).mean()
# Define the optimizer and learning rate
mesh1.verts_packed().requires_grad = True
optimizer = torch.optim.Adam([mesh1.verts_packed()], lr=0.0001)
# optimizer = torch.optim.SGD([mesh1.verts_packed()], lr=0.001)
# Define the loss functions
def laplacian_smoothing_loss(mesh):
return mesh_laplacian_smoothing(mesh, method="uniform") # cotcurv uniform cot
def edge_loss(mesh):
return mesh_edge_loss(mesh, target_length=avg_edge_length)
def normal_consistency_loss(mesh):
return mesh_normal_consistency(mesh)
# Optimization loop
for i in range(10000):
optimizer.zero_grad()
# find boundaries with the find holes function?
# Forward pass
mesh1_out = Meshes(verts=[mesh1.verts_packed()], faces=[mesh1.faces_packed()])
ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list()), torch.stack(mesh2.verts_list()))
bound_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,mesh1_boundaries,:], torch.stack(mesh2.verts_list())[:,mesh2_boundaries,:])
# crotch_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:], torch.stack(mesh2.verts_list())[:,crotch_body,:],)
# crotch_ch_loss = torch.mean((torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:] - torch.stack(mesh2.verts_list())[:,crotch_body,:])**2, dim=2).mean()
loss: torch.Tensor = ch_loss * 3 + (bound_ch_loss) * 20
# loss: torch.Tensor = ch_loss + ch_norm_loss + (bound_ch_loss + bound_ch_norm_loss) * 1
lap_loss = laplacian_smoothing_loss(mesh1_out)
e_loss = edge_loss(mesh1_out)
# loss += edge_loss_custom(torch.stack(mesh1_out.verts_list()), template_orig[None,:,:], mesh1.edges_packed()) * 100
norm_loss = normal_consistency_loss(mesh1_out)
loss += lap_loss * 0.05 + e_loss * 100 + norm_loss * 0.1
# Backward pass
loss.backward()
optimizer.step()
# Print the loss
if i % 10 == 0:
print(f"it: {i} : loss: {loss:.5f} chamfer: {ch_loss:.6f} bound_chamfer: {bound_ch_loss:.6f} laplacian: {lap_loss:.6f} edge: {e_loss:.6f} normal: {norm_loss:.5f}")
# create a plotly plot with the body and the mesh1_out
registered = mesh1_out.verts_packed().detach().cpu().numpy()
template = mesh2.verts_packed().detach().cpu().numpy()
faces = faces.detach().cpu().numpy()
F = F.detach().cpu().numpy()
# faces, F = F, faces
# verts, star_body = star_body, verts
mesh1_boundaries, mesh2_boundaries = mesh1_boundaries.detach().cpu().numpy(), mesh2_boundaries.detach().cpu().numpy()
# mesh2_boundaries = mesh1_boundaries.detach().cpu().numpy()
fig = go.Figure(data=[
go.Mesh3d(
x=template[:,0], y=template[:,1], z=template[:,2],
i=faces[:,0], j=faces[:,1], k=faces[:,2],
color='lightblue', opacity=0.80),
go.Mesh3d(
x=registered[:,0], y=registered[:,1], z=registered[:,2],
i=F[:,0], j=F[:,1], k=F[:,2],
color='lightpink', opacity=0.50),
go.Scatter3d(
x=registered[mesh1_boundaries,0], y=registered[mesh1_boundaries,1], z=registered[mesh1_boundaries,2],
mode='markers', marker=dict(size=2, color='red')),
go.Scatter3d(
x=template[mesh2_boundaries,0], y=template[mesh2_boundaries,1], z=template[mesh2_boundaries,2],
mode='markers', marker=dict(size=2, color='blue')),
# go.Scatter3d(
# x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
# mode='markers', marker=dict(size=2, color='red')),
# go.Scatter3d(
# x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
# mode='markers', marker=dict(size=2, color='blue')),
])
# make it cuve and limit to -1 and -1
fig.update_layout(scene = dict(
xaxis = dict(range=[-1,1],),
yaxis = dict(range=[-1,1],),
zaxis = dict(range=[-1,1],),),)
fig.update_layout(scene_aspectmode='cube')
fig.show()
# create a plotly wireframe plot of registered
fig = go.Figure()
for edge in mesh1_out.edges_packed().detach().cpu().numpy():
fig.add_trace(go.Scatter3d(
x=registered[edge,0], y=registered[edge,1], z=registered[edge,2],
mode='lines', line=dict(color='red', width=1)))
fig.update_layout(scene = dict(
xaxis = dict(range=[-1,1],),
yaxis = dict(range=[-1,1],),
zaxis = dict(range=[-1,1],),),)
fig.update_layout(scene_aspectmode='cube')
fig.show()
pass
if __name__ == "__main__":
folder = '/data-net/code/cloth3d/train'
# folder = r'D:\Work\code\cloth3d\train'
# frames = [f'00761_{i}' for i in range(300)]
# folder = r'Z:\Work\Huni\Barcelona\Dataset\cloth3d\train'
frames = get_frames(folder, 'frames_all.txt')
root = '/'.join(folder.split('/')[:-1])
type = folder.split('/')[-1]
use_no_orient = False
no_shape = False
is_skirt = False
skirt_modifier = '_skirt' if is_skirt else ''
orient_mod = '_no_orient' if use_no_orient else ''
shape_mod = '_no_shape' if no_shape else ''
# mask = np.load('/data-net/dataset/00761/mesh_mask/00761.npy')
# faces = np.load('/data-net/dataset/00761/faces/00761.npy')
for frame in frames:
frame = random.choice(frames)
# frame = '00761 0 Tshirt'
frame = '01893 256 Jumpsuit'
register(folder, frame)
break
pass
data/registration_weights.py
-1071
View File
File diff suppressed because it is too large Load Diff
data/utils.py
+1 -1
View File
@@ -1,7 +1,7 @@
import tables as tb import tables as tb
class Frame: class Frame:
def __init__(self, id: str, frame_nr: int, garment: str=''): def __init__(self, id, frame_nr, garment=''):
self.id = id self.id = id
self.frame_nr = frame_nr self.frame_nr = frame_nr
self.garment = garment self.garment = garment
docker/docker-compose.yml
-1
View File
@@ -34,7 +34,6 @@ services:
- '/data-net/guest/lhunor/data:/data-net' - '/data-net/guest/lhunor/data:/data-net'
- '/data-fast/126-data1/lhunor/data:/data-local/126' - '/data-fast/126-data1/lhunor/data:/data-local/126'
- '/data-fast/127-data1/lhunor/data:/data-local/127' - '/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' - '/data-fast/128-data1/lhunor/data:/data-local/128'
deploy: deploy:
resources: resources:
losses/losses.py
+5 -5
View File
@@ -18,12 +18,12 @@ from torch_cluster import knn
smpl_to_keep_orig = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1) smpl_to_keep_orig = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
smpl_to_keep_orig = smpl_to_keep_orig[smpl_to_keep_orig < 6890] smpl_to_keep_orig = smpl_to_keep_orig[smpl_to_keep_orig < 6890]
# to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1) to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
F_coll = torch.tensor(np.load("smpl/smpl_faces.npy"), device='cuda') F_coll = torch.tensor(np.load("smpl/smpl_faces.npy")[to_keep,:], device='cuda')
F_coll_skirt = torch.tensor(np.load("smpl/smpl_faces_skirt.npy"), device='cuda') F_coll_skirt = torch.tensor(np.load("smpl/smpl_faces_skirt.npy")[to_keep,:], device='cuda')
F_coll_orig = torch.tensor(np.load("smpl/smpl_faces_orig.npy")[:6890,:], device='cuda') F_coll_orig = torch.tensor(np.load("smpl/smpl_faces_orig.npy")[:6890,:], device='cuda')
F_norm = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[None,:,:], [4,1,1]), device='cuda') F_norm = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[None,:,:], [4,1,1]), device='cuda') F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
F_norm_orig = torch.tensor(np.tile(np.load("smpl/smpl_faces_orig.npy")[:6890,:][None,:,:], [4,1,1]), device='cuda') F_norm_orig = torch.tensor(np.tile(np.load("smpl/smpl_faces_orig.npy")[:6890,:][None,:,:], [4,1,1]), device='cuda')
F_norm_orig = F_norm_orig[:,smpl_to_keep_orig,:] F_norm_orig = F_norm_orig[:,smpl_to_keep_orig,:]
models/base.py
+9 -9
View File
@@ -35,10 +35,10 @@ class VariationalEncoder(nn.Module):
class AEEncoder(nn.Module): class AEEncoder(nn.Module):
def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7, activation='tanh', is_label=False, nr_of_labels=None, device='cuda'): def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7, activation='tanh', is_label=False, nr_of_labels=None, device='cuda', input_features=3):
super(AEEncoder, self).__init__() super(AEEncoder, self).__init__()
self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_labels, device=device) self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_labels, device=device, input_features=input_features)
if activation == 'tanh': if activation == 'tanh':
self.act = nn.Tanh() self.act = nn.Tanh()
elif activation == 'sigmoid': elif activation == 'sigmoid':
@@ -53,13 +53,13 @@ class AEEncoder(nn.Module):
else: else:
x = inputs x = inputs
x = self.encoder(x) x, features = self.encoder(x)
x = self.act(x) x = self.act(x)
if len(inputs) == 2: if len(inputs) == 2:
return x, inputs[1] return x, features, inputs[1]
else: else:
return x return x, features
class Decoder(nn.Module): class Decoder(nn.Module):
@@ -75,12 +75,12 @@ class Decoder(nn.Module):
self.softmax = nn.Softmax(dim=1) self.softmax = nn.Softmax(dim=1)
def forward(self, inputs): def forward(self, inputs):
if len(inputs) == 2: if len(inputs) == 3:
x, mask = inputs x, features, mask = inputs
else: else:
x = inputs x, features = inputs
x = self.decoder(x) x = self.decoder([x, features])
if self.is_label: if self.is_label:
x = x * mask[:,None,:,:] x = x * mask[:,None,:,:]
models/cond_encoder.py
+1 -1
View File
@@ -8,7 +8,7 @@ class AutoEncoder(nn.Module):
def __init__(self, config): def __init__(self, config):
super(AutoEncoder, self).__init__() super(AutoEncoder, self).__init__()
self.encoder : AEEncoder = AEEncoder(config['encoder_block_settings'], config['kernel_size'], config['activation'], is_label=config['is_label'], device=config['device'], nr_of_labels=7) self.encoder : AEEncoder = AEEncoder(config['encoder_block_settings'], config['kernel_size'], config['activation'], is_label=config['is_label'], device=config['device'], nr_of_labels=7, input_features=config['input_features'])
self.decoder : Decoder = Decoder(config['decoder_block_settings'], kernel_size=config['kernel_size'], is_label=config['is_label'], device=config['device'], nr_of_classes=6) self.decoder : Decoder = Decoder(config['decoder_block_settings'], kernel_size=config['kernel_size'], is_label=config['is_label'], device=config['device'], nr_of_classes=6)
models/convnext_orig.py
+35 -16
View File
@@ -81,8 +81,6 @@ class ConvNeXtBase(nn.Module):
if norm_layer is None: if norm_layer is None:
norm_layer = partial(LayerNorm2d, eps=1e-6) norm_layer = partial(LayerNorm2d, eps=1e-6)
layers_list: List[nn.Module] = []
# Stem # Stem
firstconv_output_channels = block_setting[0].input_channels firstconv_output_channels = block_setting[0].input_channels
self.stem = Conv2dNormActivation( self.stem = Conv2dNormActivation(
@@ -95,7 +93,8 @@ class ConvNeXtBase(nn.Module):
activation_layer=None, activation_layer=None,
bias=True, bias=True,
).to(device) ).to(device)
layers_list.append(self.stem) self.stages = []
self.downsamples = []
total_stage_blocks = sum(cnf.num_layers for cnf in block_setting) total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
stage_block_id = 0 stage_block_id = 0
@@ -107,17 +106,20 @@ class ConvNeXtBase(nn.Module):
sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0) sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
stage.append(block(cnf.input_channels, layer_scale, sd_prob, kernel_size)) stage.append(block(cnf.input_channels, layer_scale, sd_prob, kernel_size))
stage_block_id += 1 stage_block_id += 1
layers_list.append(nn.Sequential(*stage).to(device)) self.stages.append(nn.Sequential(*stage).to(device))
if cnf.out_channels is not None: if cnf.out_channels is not None:
# Downsampling # Downsampling
layers_list.append( self.downsamples.append(
nn.Sequential( nn.Sequential(
norm_layer(cnf.input_channels), norm_layer(cnf.input_channels),
nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2), nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
).to(device) ).to(device)
) )
self.stages = nn.ModuleList(self.stages)
self.downsamples = nn.ModuleList(self.downsamples)
self.layers = nn.Sequential(*layers_list).to(device) # self.layers = nn.Sequential(*layers_list).to(device)
# self.layers = nn.ModuleList(layers_list)
self.act = nn.Tanh().to(device) self.act = nn.Tanh().to(device)
# for m in self.modules(): # for m in self.modules():
@@ -127,9 +129,16 @@ class ConvNeXtBase(nn.Module):
# nn.init.zeros_(m.bias) # nn.init.zeros_(m.bias)
def _forward_impl(self, x: Tensor) -> Tensor: def _forward_impl(self, x: Tensor) -> Tensor:
x = self.layers(x) x = self.stem(x)
features = []
for i in range(len(self.stages)):
x = self.stages[i](x)
if i == len(self.stages) - 1:
continue
features.append(x)
x = self.downsamples[i](x)
x = self.act(x) x = self.act(x)
return x return x, features
def forward(self, x: Tensor) -> Tensor: def forward(self, x: Tensor) -> Tensor:
return self._forward_impl(x) return self._forward_impl(x)
@@ -165,13 +174,14 @@ class ConvNeXtReverseBase(nn.Module):
norm_layer = partial(LayerNorm2d, eps=1e-6) norm_layer = partial(LayerNorm2d, eps=1e-6)
self.stem = nn.Sequential( self.stem = nn.Sequential(
nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=4, stride=4), nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2),
norm_layer(block_setting[0].out_channels), norm_layer(block_setting[0].out_channels),
).to(device) ).to(device)
total_stage_blocks = sum(cnf.num_layers for cnf in block_setting) total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
stage_block_id = 0 stage_block_id = 0
layers_list: List[nn.Module] = [] self.stages = []
self.upsamples = []
for i, cnf in enumerate(block_setting): for i, cnf in enumerate(block_setting):
@@ -182,17 +192,17 @@ class ConvNeXtReverseBase(nn.Module):
sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0) sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
stage.append(block(cnf.input_channels, layer_scale, sd_prob, kernel_size)) stage.append(block(cnf.input_channels, layer_scale, sd_prob, kernel_size))
stage_block_id += 1 stage_block_id += 1
layers_list.append(nn.Sequential(*stage)) self.stages.append(nn.Sequential(*stage).to(device))
if i == 0: if i == 0:
layers_list.append(self.stem) self.upsamples.append(self.stem)
else: else:
# Upsampling # Upsampling
layers_list.append( self.upsamples.append(
nn.Sequential( nn.Sequential(
norm_layer(cnf.input_channels), norm_layer(cnf.input_channels),
nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2), nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
) ).to(device)
) )
self.last = nn.Sequential( self.last = nn.Sequential(
@@ -201,7 +211,8 @@ class ConvNeXtReverseBase(nn.Module):
Permute([0, 3, 1, 2]), Permute([0, 3, 1, 2]),
).to(device) ).to(device)
self.layers = nn.Sequential(*layers_list).to(device) self.stages = nn.ModuleList(self.stages)
self.upsamples = nn.ModuleList(self.upsamples)
self.activation = nn.Tanh().to(device) if not is_label else nn.Identity().to(device) self.activation = nn.Tanh().to(device) if not is_label else nn.Identity().to(device)
@@ -213,7 +224,15 @@ class ConvNeXtReverseBase(nn.Module):
def _forward_impl(self, x) -> Tensor: def _forward_impl(self, x) -> Tensor:
x = self.layers(x) x, features = x
for i in range(len(self.stages)-1):
if i == 0:
x = self.stages[i](x)
else:
x = self.stages[i](torch.cat([x, features[-i]], dim=1))
x = self.upsamples[i](x)
x = self.stages[-1](x)
x = self.upsamples[-1](x)
x = self.last(x) x = self.last(x)
x = self.activation(x) x = self.activation(x)
return x return x
models/cvae_model.py
+13 -13
View File
@@ -27,16 +27,16 @@ class CVAE(nn.Module):
if config['body_model'] != '': if config['body_model'] != '':
load_model(self.body_enc, config['body_model'] + '_enc') load_model(self.body_enc, config['body_model'] + '_enc')
if config['use_normals']: if config['use_curvature']:
self.norm_enc : AEEncoder = AEEncoder(config['cond_ae_norm']['encoder_block_settings'], config['cond_ae_norm']['kernel_size'], config['cond_ae_norm']['activation'], device=config['device']) self.curv_enc : AEEncoder = AEEncoder(config['cond_ae_curv']['encoder_block_settings'], config['cond_ae_curv']['kernel_size'], config['cond_ae_curv']['activation'], device=config['device'], is_label=True, nr_of_labels=7)
if config['normals_model'] != '': if config['curvature_model'] != '':
load_model(self.norm_enc, config['normals_model'] + '_enc') load_model(self.curv_enc, config['curvature_model'] + '_enc')
if config['freeze_encoders']: if config['freeze_encoders']:
for param in self.static_enc.parameters(): param.requires_grad = False for param in self.static_enc.parameters(): param.requires_grad = False
for param in self.body_enc.parameters(): param.requires_grad = False for param in self.body_enc.parameters(): param.requires_grad = False
if config['use_normals']: if config['use_curvature']:
for param in self.norm_enc.parameters(): param.requires_grad = False for param in self.curv_enc.parameters(): param.requires_grad = False
self.encoder = VariationalEncoder(config['cvae']['encoder_block_settings'], config['cvae']['kernel_size'], config['cvae']['input_features'], device=config['device']) self.encoder = VariationalEncoder(config['cvae']['encoder_block_settings'], config['cvae']['kernel_size'], config['cvae']['input_features'], device=config['device'])
self.decoder : Decoder = Decoder(config['cvae']['decoder_block_settings'], kernel_size=config['cvae']['kernel_size'], device=config['device']) self.decoder : Decoder = Decoder(config['cvae']['decoder_block_settings'], kernel_size=config['cvae']['kernel_size'], device=config['device'])
@@ -48,7 +48,7 @@ class CVAE(nn.Module):
else: else:
self.avg_garment = None self.avg_garment = None
self.use_normals = config['use_normals'] self.use_curvature = config['use_curvature']
self.use_offsets_over_body = not config['use_skinning'] self.use_offsets_over_body = not config['use_skinning']
if self.use_offsets_over_body: if self.use_offsets_over_body:
@@ -98,9 +98,9 @@ class CVAE(nn.Module):
def encode_conds(self, uv_template, uv_body_posed, uv_curv): def encode_conds(self, uv_template, uv_body_posed, uv_curv):
z_body = self.body_enc(uv_body_posed) z_body = self.body_enc(uv_body_posed)
z_static = self.static_enc(uv_template) z_static = self.static_enc(uv_template)
if self.use_normals: if self.use_curvature:
z_norm = self.norm_enc(uv_curv) z_curv = self.curv_enc(uv_curv)
z_cond = torch.concat([z_static, z_body, z_norm], dim=1) z_cond = torch.concat([z_static, z_body, z_curv], dim=1)
else: else:
z_cond = torch.concat([z_static, z_body], dim=1) z_cond = torch.concat([z_static, z_body], dim=1)
@@ -108,13 +108,13 @@ class CVAE(nn.Module):
def forward(self, inputs): def forward(self, inputs):
uv_unposed, uv_template, uv_body_posed, uv_norm, uv_mask = inputs uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_norm], dim=1) x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], dim=1)
mean, logvar = self.encoder(x) mean, logvar = self.encoder(x)
z: torch.Tensor = self.reparameterize(mean, logvar) z: torch.Tensor = self.reparameterize(mean, logvar)
z = z.view((mean.shape[0], *self.latent_shape)) z = z.view((mean.shape[0], *self.latent_shape))
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_norm) z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
z = torch.concat([z, z_cond], dim=1) z = torch.concat([z, z_cond], dim=1)
x = self.decoder(z) x = self.decoder(z)
models/cvae_norm_model.py
-97
View File
@@ -1,97 +0,0 @@
import os
import time
import random
import torch
from torch import nn
import numpy as np
import scipy as sp
# import plotly.graph_objects as go
# import plotly.express as ex
from utils.checkpoints import load_model
from models.base import AEEncoder, VariationalEncoder, Decoder
class CVAENorm(nn.Module):
def __init__(self, config):
super(CVAENorm, self).__init__()
self.size = config['resolution']
self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'], device=config['device'])
if config['static_model'] != '':
load_model(self.static_enc, config['static_model'] + '_enc')
self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], config['cond_ae_body']['kernel_size'], config['cond_ae_body']['activation'], device=config['device'])
if config['body_model'] != '':
load_model(self.body_enc, config['body_model'] + '_enc')
if config['freeze_encoders']:
for param in self.static_enc.parameters(): param.requires_grad = False
for param in self.body_enc.parameters(): param.requires_grad = False
self.encoder = VariationalEncoder(config['cvae']['encoder_block_settings'], config['cvae']['kernel_size'], config['cvae']['input_features'], device=config['device'])
self.decoder : Decoder = Decoder(config['cvae']['decoder_block_settings'], kernel_size=config['cvae']['kernel_size'], device=config['device'])
self.latent_shape = (config['cvae']['latent_dim'], 8, 4)
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None, add_noise=False):
if noise is None:
eps = tf.random.normal(shape=(uv_template.shape[0], 8, 4, self.latent_dim))
else:
eps = noise
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=False)
if add_noise:
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
z = tf.concat([eps, z_cond], axis=-1)
x = self.decoder(z, training=False)
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
if self.use_offsets_over_body:
x = x + uv_body_posed
x = x * uv_mask[:,:,:,tf.newaxis]
return x
def reparameterize(self, mean: torch.Tensor, logvar: torch.Tensor, training: bool = True, fix_noise: bool = False):
if training == True:
if not fix_noise:
eps = torch.randn(size=mean.shape).to(mean.device)
else:
eps = torch.randn(size=(1, mean.shape[1]))
eps = torch.tile(eps, [mean.shape[0],1])
out = eps * torch.exp(logvar * 0.5) + mean
else:
out = torch.exp(logvar * 0.5) + mean
return out
def encode_conds(self, uv_template, uv_body_posed):
z_body = self.body_enc(uv_body_posed)
z_static = self.static_enc(uv_template)
z_cond = torch.concat([z_static, z_body], dim=1)
return z_cond
def forward(self, inputs):
uv_unposed, uv_template, uv_body_posed, uv_mask = inputs
x = torch.concat([uv_unposed, uv_template, uv_body_posed], dim=1)
mean, logvar = self.encoder(x)
z: torch.Tensor = self.reparameterize(mean, logvar)
z = z.view((mean.shape[0], *self.latent_shape))
z_cond = self.encode_conds(uv_template, uv_body_posed)
z = torch.concat([z, z_cond], dim=1)
x = self.decoder(z)
x = x * uv_mask[:,None,:,:]
return x, mean, logvar
models/direct.py
+29
View File
@@ -0,0 +1,29 @@
import torch
from torch import nn
from models.convnext_orig import ConvNeXt, CNBlock, CNBlockConfig
from typing import Any, Callable, List, Optional, Sequence
class DirectModel(nn.Module):
def __init__(self, block_setting: List[CNBlockConfig], kernel_size: int = 7, input_features: int = 16, device: str = 'cuda'):
super(DirectModel, self).__init__()
self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, input_features=input_features, device=device)
self.decoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, device=device, reversed=True)
layers = [CNBlock(block_setting[-1].input_channels, kernel_size=kernel_size) for _ in range(5)]
self.mean_branch = nn.Sequential(*layers).to(device)
layers = [CNBlock(block_setting[-1].input_channels, kernel_size=kernel_size) for _ in range(5)]
self.var_branch = nn.Sequential(*layers).to(device)
self.flatten_mean = nn.Flatten()
self.flatten_var = nn.Flatten()
def forward(self, inputs):
body, template = inputs
x = self.encoder(x)
return mean
run.py
+1 -1
View File
@@ -1,6 +1,6 @@
import os import os
import sys import sys
from configs.cvae_norm_1 import config from configs.direct_1 import config
import argparse import argparse
# parser = argparse.ArgumentParser(description='') # parser = argparse.ArgumentParser(description='')
run_parallel.py
+5 -8
View File
@@ -4,27 +4,24 @@ import os
if __name__ == '__main__': if __name__ == '__main__':
# folder = '/home/workspace/data/dataset/cloth3d/train' folder = '/home/workspace/data/dataset/cloth3d/train'
# folder = '/home/workspace/local-data/cape_processed/train/male' # folder = '/home/workspace/local-data/cape_processed/train/male'
# folder = '/home/workspace/local-data/valid_skirt' # folder = '/home/workspace/local-data/valid_skirt'
# folder = '/home/workspace/local-data/valid' # folder = '/home/workspace/local-data/valid'
# ids = os.listdir(folder) # ids = os.listdir(folder)
# ids = os.listdir(folder + '/mesh_posed_no_orient')#[280000:] # ids = os.listdir(folder + '/mesh_posed_no_orient')#[280000:]
skirt = 0 skirt = 0
obj_folder = '/data-net/dataset/cloth3d/train'
folder = '/data-local/127/train'
count = 0 count = int(500 * 1e3)
step = 100 step = 10000
# !!!!!!! change skirt parameter in script # !!!!!!! change skirt parameter in script
while count <1200: while count < 510 * 1e3:
# os.system('nohup python convert_to_uv_cape.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &') # os.system('nohup python convert_to_uv_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.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 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('nohup python data/convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_2_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv/train_2_' + str(count) + '_' + str(count+step) + '_again.out' + ' 2>&1 &')
os.system(f'nohup python data/calculate_normals.py -f {folder} -f_obj {obj_folder} -s {count} -e {(count+step)} -n normals_train_calc_{count}_{(count+step)} > nohup/normals/normals_train_calc_{count}_{(count+step)}.out' + ' 2>&1 &')
# os.system('nohup python 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_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 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 &') # os.system('nohup python downscale.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + ' > nohup/downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
scripts/ae_norm_train.py
-218
View File
@@ -1,218 +0,0 @@
import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
import argparse
import importlib
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
parser.add_argument('--config', type=str)
parser.add_argument('--debug', action='store_true')
args = vars(parser.parse_args())
config_module = importlib.import_module(args['config'])
config = config_module.config
if args['debug']:
from configs.ae_norm_1 import config
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
if args['debug']:
os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
import time
import wandb
import random
import numpy as np
import torch
from torch import nn
# from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure as MSSSIM
from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, get_summary_str
import plotly.graph_objects as go
import plotly.express as ex
import torch.optim as optim
if config['is_skirt']:
from losses.losses import edge_loss_skirt as edge_loss
else:
from losses.losses import edge_loss
from data.dataset import Dataset
# from torch.utils.data import DataLoader
from custom_utils.custom_utils import plot_3d
from models.cond_encoder import AutoEncoder
from utils.checkpoints import CheckpointSaver, load_model
from custom_utils.custom_utils import Accumulator
from utils.utils import InfiniteDataLoader as DataLoader
from pytorch_msssim import MS_SSIM
def init_wandb():
run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127_2/temp', tags=config['tags'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('scripts/ae_norm_train.py')
artifact.add_file('losses/losses.py')
artifact.add_file('models/base.py')
artifact.add_file('models/cond_encoder.py')
artifact.add_file('models/convnext.py')
artifact.add_file('data/dataset.py')
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')]
for config_file_name in config_file_names:
file_name = '/'.join(config_file_name.split('.')) + '.py'
if os.path.exists(file_name):
artifact.add_file(file_name)
run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
def train_step(batch):
model.train()
pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.1
ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_normals'] + 1.1))
loss = 1e3 * (1 - ms_ssim)
optimizer.zero_grad()
loss.backward()
optimizer.step()
model.eval()
with torch.inference_mode():
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_normals']
l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_normals'] - pred)), axis=1))
metric = torch.mean(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1)))
return { 'train/loss': loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, }
def valid_step(batch):
model.eval()
with torch.inference_mode():
pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.1
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_normals']
l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_normals'] - pred)), axis=1))
ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_normals'] + 1.1))
loss = 1e3 * (1 - ms_ssim)
metric = torch.mean(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1)))
return { 'valid/loss': loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, }
def plot_everything(batch: dict, type: str):
model.eval()
pred = model([batch['uv_normals'], batch['uv_mask']]).detach() * 1.1
# mesh_pred = uv2mesh(pred)
# mesh_orig = batch['mesh_normals']
orig_img = torch.permute(batch['uv_normals'], [0,2,3,1])
pred_img = torch.permute(pred, [0,2,3,1])
img = torch.concat((orig_img, pred_img), axis=2)
imgs = [wandb.Image((img[i].cpu().numpy() + 1.1) / 2.2) for i in range(4)]
# F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
# garment_unposed, faces = mesh_pred.cpu().numpy(), F
# gt = [mesh_orig[i].cpu().numpy() for i in range(pred.shape[0])]
# body = uv2mesh(batch['uv_body_posed']).cpu().numpy()
# garment_posed = garment_unposed
# collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
# collision_faces = [np.concatenate((F, faces + (3638 if config['is_cape'] else 14475)), axis=0) for i in range(len(garment_posed))]
# collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
# figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces, collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
figs = {}
return { **figs, **{ type + '/images': imgs } }
def train():
valid_iterator = iter(valid_dataloader)
step = 0
for batch in train_dataloader:
train_loss = train_step(batch)
train_accumulator.update(train_loss)
if step % 25 == 0:
data_valid = next(valid_iterator)
valid_loss = valid_step(data_valid)
valid_accumulator.update(valid_loss)
if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid')
to_log = { **to_log, **train_figs, **valid_figs }
if step == 0:
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
if not args['debug']:
wandb.log(to_log, step=step)
train_accumulator.reset()
valid_accumulator.reset()
print(step)
if step % 10000 == 0:
ckpt_saver_enc(model.encoder, step, optimizer, scaler=None)
ckpt_saver_dec(model.decoder, step, optimizer, scaler=None)
step += 1
class Scale(torch.nn.Module):
def __init__(self, scale_factor):
super().__init__()
self.scale_factor = scale_factor
def forward(self, x):
return x * self.scale_factor
if __name__ == '__main__':
if not args['debug']:
init_wandb()
skirt_modifier = '_skirt' if config['is_skirt'] else ''
dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True)
valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True)
train_accumulator = Accumulator()
valid_accumulator = Accumulator()
model : AutoEncoder = AutoEncoder(config['cond_ae'])
optimizer = optim.Adam(model.parameters(), lr=config['lr'])
# _MSSSIM = MSSSIM(data_range=2).to(config['device'])
_MSSSIM = MS_SSIM(data_range=2.2, size_average=True, channel=3)
if config['pretrained_model'] != '':
load_model(model.encoder, config['pretrained_model'] + '_enc')
load_model(model.decoder, config['pretrained_model'] + '_dec')
ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_enc', top_n=5)
ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_dec', top_n=5)
run_name = config['name']
print(f'started training: {run_name}')
if not args['debug']:
# model = torch.compile(model)
wandb.watch(model, log='all', log_freq=10000, log_graph=True)
train()
scripts/cvae.py
+11 -13
View File
@@ -13,8 +13,7 @@ args = vars(parser.parse_args())
config_module = importlib.import_module(args['config']) config_module = importlib.import_module(args['config'])
config = config_module.config config = config_module.config
if args['debug']: if args['debug']:
from configs.abl_norm import config pass
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
if args['debug']: if args['debug']:
@@ -51,7 +50,7 @@ from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary
def init_wandb(): def init_wandb():
run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags']) run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags'])
artifact = wandb.Artifact('code', type='code') artifact = wandb.Artifact('code', type='code')
artifact.add_file('scripts/cvae.py') artifact.add_file('scripts/cvae.py')
artifact.add_file('losses/losses.py') artifact.add_file('losses/losses.py')
@@ -84,7 +83,7 @@ def train_step(batch):
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed'] mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
model.train() model.train()
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']]) pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
mesh_pred = uv2mesh(pred) mesh_pred = uv2mesh(pred)
if config['use_skinning']: if config['use_skinning']:
@@ -140,7 +139,7 @@ def valid_step(batch):
model.eval() model.eval()
with torch.inference_mode(): with torch.inference_mode():
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']]) pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
mesh_pred = uv2mesh(pred) mesh_pred = uv2mesh(pred)
if config['use_skinning']: if config['use_skinning']:
@@ -189,16 +188,15 @@ def plot_everything(batch: dict, type: str):
model.eval() model.eval()
with torch.inference_mode(): with torch.inference_mode():
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']]) pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
mesh_pred = uv2mesh(pred).detach().cpu().numpy() mesh_pred = uv2mesh(pred).detach().cpu().numpy()
orig_img = torch.permute(uv_orig, [0,2,3,1]) orig_img = torch.permute(uv_orig, [0,2,3,1])
uv_static_img = torch.permute(batch['uv_static'], [0,2,3,1]) uv_static_img = torch.permute(batch['uv_static'], [0,2,3,1])
uv_body_posed_img = torch.permute(batch['uv_body_posed'], [0,2,3,1]) uv_body_posed_img = torch.permute(batch['uv_body_posed'], [0,2,3,1])
uv_normals_img = (torch.permute(batch['uv_normals'], [0,2,3,1]) + 1.1) / 2.2
pred_img = torch.permute(pred, [0,2,3,1]) pred_img = torch.permute(pred, [0,2,3,1])
img = torch.concat((uv_static_img, uv_body_posed_img, uv_normals_img, orig_img, pred_img), axis=2) img = torch.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)] imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)]
F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy") F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
@@ -264,10 +262,10 @@ def sample(folder):
count = 0 count = 0
for batch in dataset_train: for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
uv_normals = tf.one_hot(batch['uv_normals'], 7) uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
nr_of_samples=2 nr_of_samples=2
pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples) pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:] pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1]) pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1])
@@ -283,7 +281,7 @@ def sample(folder):
mask = batch['mesh_mask'].numpy().astype(bool)[0] mask = batch['mesh_mask'].numpy().astype(bool)[0]
_, faces = _maskFaces(np.ones((config['smpl_res'], 3)), F, mask) _, faces = _maskFaces(np.ones((config['smpl_res'], 3)), F, mask)
pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True) pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
mesh_pred_zero = uv2mesh(pred_zero) mesh_pred_zero = uv2mesh(pred_zero)
mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy() mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
@@ -459,7 +457,7 @@ def train():
g_balancing_coeff += 1 / 10000 g_balancing_coeff += 1 / 10000
if step % 10000 == 0: if step % 10000 == 0:
ckpt_saver(model, step, optimizer, None) ckpt_saver(model, step, optimizer)
step += 1 step += 1
@@ -493,6 +491,6 @@ if __name__ == '__main__':
if not args['debug']: if not args['debug']:
# model = torch.compile(model) # model = torch.compile(model)
wandb.watch(model, log='all', log_freq=10000, log_graph=True) wandb.watch(model, log='all', log_freq=1000, log_graph=True)
train() train()
scripts/cvae_norm.py
-400
View File
@@ -1,400 +0,0 @@
import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
import importlib
import argparse
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
parser.add_argument('--config', type=str)
parser.add_argument('--debug', action='store_true')
args = vars(parser.parse_args())
config_module = importlib.import_module(args['config'])
config = config_module.config
if args['debug']:
from configs.cvae_norm_1 import config
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
if args['debug']:
os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
import time
import argparse
import wandb
import tqdm
import random
import numpy as np
import plotly.graph_objects as go
import plotly.express as ex
from PIL import Image
if config['is_skirt']:
from losses.losses import edge_loss_skirt as edge_loss
else:
from losses.losses import edge_loss
import torch
import torch.nn as nn
import torch.optim as optim
from data.dataset import Dataset
from models.cond_encoder import AutoEncoder
from utils.checkpoints import CheckpointSaver, load_model
from utils.utils import InfiniteDataLoader as DataLoader
from utils.postprocess import fix_collision
from losses.losses import collision_loss, normal_loss
from pytorch_msssim import MS_SSIM
from models.cvae_norm_model import CVAENorm
from smpl.smpl_torch import SMPL
from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary_html, maskFacesBatch, get_summary_str, plot_3d_collision
def init_wandb():
run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('scripts/cvae_norm.py')
artifact.add_file('losses/losses.py')
artifact.add_file('models/base.py')
artifact.add_file('models/cond_encoder.py')
artifact.add_file('models/convnext.py')
artifact.add_file('models/cvae_norm_model.py')
artifact.add_file('data/dataset.py')
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')]
for config_file_name in config_file_names:
file_name = '/'.join(config_file_name.split('.')) + '.py'
if os.path.exists(file_name):
artifact.add_file(file_name)
run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
def kl_loss(mean: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
kl_loss = 1 + logvar - torch.square(mean) - torch.exp(logvar)
kl_loss = torch.mean(kl_loss)
kl_loss *= -0.5
return kl_loss
def train_step(batch):
model.train()
pred, mean, logvar = model([batch['uv_normals'], batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']])
pred = pred * 1.1
mesh_pred = uv2mesh(pred)
l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_normals'] - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_normals'] - mesh_pred) ** 2, axis=-1), axis=-1))
ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_normals'] + 1.1))
kl = kl_loss(mean, logvar)
reconstr_loss = l2_loss + l2_mesh_loss + (1 - ms_ssim)
loss = reconstr_loss + g_balancing_coeff * kl
optimizer.zero_grad()
loss.backward()
optimizer.step()
model.eval()
with torch.inference_mode():
mean = torch.mean(mean)
logvar = torch.mean(logvar)
metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_normals'] - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
return { 'train/loss': loss, 'train/euclid_error': metric,
'train/reconstr_loss': reconstr_loss,
'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
'train/l2_loss': l2_loss, 'train/l2_mesh_loss': l2_mesh_loss,
'train/ms_ssim': ms_ssim,
}
def valid_step(batch):
model.eval()
with torch.inference_mode():
pred, mean, logvar = model([batch['uv_normals'], batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']])
pred = pred * 1.1
mesh_pred = uv2mesh(pred)
l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_normals'] - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_pred - batch['mesh_normals']) ** 2, axis=-1), axis=-1))
ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_static'] + 1.1))
kl = kl_loss(mean, logvar)
reconstr_loss = l2_loss + l2_mesh_loss + (1 - ms_ssim)
loss = reconstr_loss + g_balancing_coeff * kl
mean = torch.mean(mean)
logvar = torch.mean(logvar)
metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_normals'] - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
return { 'valid/loss': loss, 'valid/euclid_error': metric,
'valid/reconstr_loss': reconstr_loss, 'valid/l2_mesh_loss': l2_mesh_loss,
'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar,
'valid/ms_ssim': ms_ssim,
}
def plot_everything(batch: dict, type: str):
model.eval()
with torch.inference_mode():
pred, mean, logvar = model([batch['uv_normals'], batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']])
pred = pred * 1.1
# mesh_pred = uv2mesh(pred).detach().cpu().numpy()
uv_static_img = torch.permute(batch['uv_static'], [0,2,3,1])
uv_body_posed_img = torch.permute(batch['uv_body_posed'], [0,2,3,1])
orig_img = (torch.permute(batch['uv_normals'], [0,2,3,1]) + 1.1) / 2.2
pred_img = torch.permute(pred, [0,2,3,1])
img = torch.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)]
# F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
# mask = batch['mesh_mask'].cpu().numpy().astype(bool)
mean_hist = wandb.Histogram(sequence=mean.cpu().numpy().mean(axis=0))
logvar_hist = wandb.Histogram(sequence=logvar.cpu().numpy().mean(axis=0))
return { **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
def sample(folder):
# run = wandb.init(project='vae_plots', name=folder.split('/')[-1])
load_model(model,folder)
count = 0
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
uv_normals = tf.one_hot(batch['uv_normals'], 7)
nr_of_samples=2
pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1])
mesh_pred = uv2mesh(pred)
mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy()
mesh_orig = batch['mesh_posed'][0].numpy()
mesh_orig_reposed = smpl.run_garment(pose[:1], shape[:1], batch['mesh_unposed'][:1]).numpy()[0]
body = batch['mesh_body_posed'][0].numpy()
body_color = np.ones_like(body, dtype=np.int32)*100
F = np.load("custom_utils/smpl_faces.npy")
mask = batch['mesh_mask'].numpy().astype(bool)[0]
_, faces = _maskFaces(np.ones((config['smpl_res'], 3)), F, mask)
pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
mesh_pred_zero = uv2mesh(pred_zero)
mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
meshes = []
meshes.append(go.Mesh3d(
# 8 vertices of a cube
x=mesh_orig_reposed[mask,0],
y=mesh_orig_reposed[mask,2],
z=mesh_orig_reposed[mask,1],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
))
for i in range(nr_of_samples):
x = mesh_posed[i][mask, :].copy()
# x[:, 2] += 1 * (i + 1)
x[:, 2] += 1
meshes.append(go.Mesh3d(
# 8 vertices of a cube
x=x[:,0],
y=x[:,2],
z=x[:,1],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
))
for i in range(nr_of_samples):
x = mesh_posed_zero[i][mask, :].copy()
# x[:, 2] += 1 * (i + 1)
x[:, 2] += 2
meshes.append(go.Mesh3d(
# 8 vertices of a cube
x=x[:,0],
y=x[:,2],
z=x[:,1],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
))
# for i in range(5):
# x = body.copy()
# x[:, 2] += 1 * i
# meshes.append(go.Mesh3d(
# # 8 vertices of a cube
# x=x[:,0],
# y=x[:,2],
# z=x[:,1],
# # i, j and k give the vertices of triangles
# i = F[:,0],
# j = F[:,1],
# k = F[:,2],
# vertexcolor=body_color
# ))
fig = go.Figure(data=meshes)
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1))
# fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
# yaxis = dict(nticks=4, range=[-1,nr_of_samples+1],),
yaxis = dict(nticks=4, range=[-1,3],),
zaxis = dict(nticks=4, range=[-1,1],),)
)
# wandb.log({'plot_' + str(count): fig})
fig.show()
pred = pred.numpy()
frames = [Image.fromarray(((pred[i] + 1) * 128).astype(np.int8), mode='RGB') for i in range(pred.shape[0])]
frames[0].save("temp/garment_1.gif", format="GIF", append_images=frames[1:],
save_all=True, duration=500, loop=0)
if count > 10:
break
print(count)
count += 1
print()
def test():
str_split = config['name'].split('_')
test_idx = str_split.index('test')
checkpoint_name = '_'.join(str_split[:test_idx])
load_model(model, 'checkpoints/' + checkpoint_name)
dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False, do_repeat=False)
static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False)
step = 0
losses = []
for batch in tqdm.tqdm(dataset_test):
batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
test_loss = test_step(batch)
losses.append(test_loss)
to_log = { **{ 'current_step': step }, **test_loss }
if step % 1000 == 0:
figs = plot_everything(batch, 'test')
to_log = { **to_log, **figs }
wandb.log(to_log, step = step)
# print(step)
step += 1
# averages = dict()
path_to_plotly_html = "temp/plotly_figure.html"
histograms = dict()
for key in test_loss.keys():
accumulated = [losses[i][key].numpy() for i in range(len(losses))]
fig = ex.histogram(accumulated, title='Average: ' + str(np.asarray(accumulated).mean()))
fig.write_html(path_to_plotly_html, auto_play = False)
histograms['all/' + key.split('/')[-1] + '_hist'] = wandb.Html(path_to_plotly_html)
# averages['all/' + key.split('/')[-1] + '_avg'] = np.asarray(accumulated).mean()
wandb.log({ **histograms }, step = step)
wandb.finish()
pass
def train():
global g_balancing_coeff
global step
valid_iterator = iter(valid_dataloader)
for batch in train_dataloader:
train_loss = train_step(batch)
train_accumulator.update(train_loss)
if step % 25 == 0:
data_valid = next(valid_iterator)
valid_loss = valid_step(data_valid)
valid_accumulator.update(valid_loss)
if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid')
to_log = { **to_log, **train_figs, **valid_figs }
if step == 0:
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
if not args['debug']:
wandb.log(to_log, step=step)
train_accumulator.reset()
valid_accumulator.reset()
print(step)
if train_loss['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
g_balancing_coeff += 1 / 10000
if step % 10000 == 0:
ckpt_saver(model, step, optimizer, None)
step += 1
if __name__ == '__main__':
if not args['debug']:
init_wandb()
smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
skirt_modifier = '_skirt' if config['is_skirt'] else ''
dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'])
dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'])
train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True, )
valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True)
train_accumulator = Accumulator()
valid_accumulator = Accumulator()
model = CVAENorm(config)
optimizer = optim.Adam(model.parameters(), lr=config['lr'])
_MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
g_balancing_coeff = config['starting_balancing_coeff']
ckpt_saver = CheckpointSaver('checkpoints/iccv/' + config['name'], top_n=5)
step = 0
if config['pretrained_model'] != '':
step = load_model(model, config['pretrained_model'])
if not args['debug']:
# model = torch.compile(model)
wandb.watch(model, log='all', log_freq=10000, log_graph=True)
train()
scripts/direct_train.py
+281
View File
@@ -0,0 +1,281 @@
import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
import importlib
import argparse
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
parser.add_argument('--config', type=str)
parser.add_argument('--debug', action='store_true')
args = vars(parser.parse_args())
config_module = importlib.import_module(args['config'])
config = config_module.config
if args['debug']:
from configs.direct_3 import config
# from configs.ae_static_orig_2 import config
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
if args['debug']:
os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
import time
import argparse
import wandb
import random
import numpy as np
import torch
# from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure as MSSSIM
from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, plot_3d_collision
import plotly.graph_objects as go
import plotly.express as ex
import torch.optim as optim
if config['is_skirt']:
from losses.losses import edge_loss_skirt as edge_loss
else:
from losses.losses import edge_loss
from data.dataset import Dataset
# from torch.utils.data import DataLoader
from losses.losses import collision_loss, normal_loss
from utils.postprocess import fix_collision
from custom_utils.custom_utils import plot_3d
from utils.checkpoints import CheckpointSaver, load_model
from custom_utils.custom_utils import Accumulator
from utils.utils import InfiniteDataLoader as DataLoader
from models.cond_encoder import AutoEncoder as Model
from smpl.smpl_torch import SMPL
from pytorch_msssim import MS_SSIM
def init_wandb():
run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('scripts/direct_train.py')
artifact.add_file('losses/losses.py')
artifact.add_file('models/base.py')
artifact.add_file('models/cond_encoder.py')
artifact.add_file('models/convnext.py')
artifact.add_file('data/dataset.py')
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')]
for config_file_name in config_file_names:
file_name = '/'.join(config_file_name.split('.')) + '.py'
if os.path.exists(file_name):
artifact.add_file(file_name)
run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
def train_step(batch):
model.train()
# with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp):
if config['use_curvature']:
input = torch.cat([batch['uv_curvature'], batch['uv_static']], axis=1)
else:
input = torch.cat([batch['uv_body_posed'], batch['uv_static']], axis=1)
pred = model([input, batch['uv_mask']])
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed']
shape = batch['pose_shape'][:, 72:]
gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
# l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
loss = l2_mesh_loss + (e_loss + c_loss + n_loss) * config['reg_loss_coeff']
optimizer.zero_grad()
loss.backward()
optimizer.step()
# scaler.scale(loss).backward()
# scaler.step(optimizer)
# scaler.update()
# optimizer.zero_grad()
model.eval()
with torch.inference_mode():
ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_unposed'] + 1))
metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/euclid_error': metric, 'train/collision_ratio': c_ratio,
'train/ms_ssim': ms_ssim, 'train/collision_loss': c_loss, 'train/l2_mesh_loss': l2_mesh_loss,
'train/normal_loss': n_loss }
def valid_step(batch):
model.eval()
with torch.inference_mode():
if config['use_curvature']:
input = torch.cat([batch['uv_curvature'], batch['uv_static']], axis=1)
else:
input = torch.cat([batch['uv_body_posed'], batch['uv_static']], axis=1)
pred = model([input, batch['uv_mask']])
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed']
shape = batch['pose_shape'][:, 72:]
gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
# l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_unposed'] + 1))
loss = l2_mesh_loss + (e_loss + c_loss + n_loss) * config['reg_loss_coeff']
metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/euclid_error': metric, 'valid/collision_ratio': c_ratio,
'valid/ms_ssim': ms_ssim, 'valid/collision_loss': c_loss, 'valid/l2_mesh_loss': l2_mesh_loss,
'valid/normal_loss': n_loss }
def plot_everything(batch: dict, type: str):
model.eval()
with torch.inference_mode():
if config['use_curvature']:
input = torch.cat([batch['uv_curvature'], batch['uv_static']], axis=1)
else:
input = torch.cat([batch['uv_body_posed'], batch['uv_static']], axis=1)
pred = model([input, batch['uv_mask']]).detach()
mesh_pred = uv2mesh(pred).cpu().numpy()
mesh_orig = batch['mesh_unposed'].cpu().numpy()
orig_img = torch.permute(batch['uv_unposed'], [0,2,3,1])
uv_static_img = torch.permute(batch['uv_static'], [0,2,3,1])
uv_body_posed_img = torch.permute(batch['uv_body_posed'], [0,2,3,1])
pred_img = torch.permute(pred, [0,2,3,1])
img = torch.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)]
F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
mask = batch['mesh_mask'].cpu().numpy().astype(bool)
shape = batch['pose_shape'][:, 72:]
gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])]).cpu().numpy()
garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
garment_unposed = [garment_unposed[i] + shape_offsets[i, mask[i], :] for i in range(mask.shape[0])]
gt = [mesh_orig[i][mask[i], :] + shape_offsets[i, mask[i], :] for i in range(pred.shape[0])]
for i in range(mesh_orig.shape[0]):
mesh_pred[i, mask[i], :] += shape_offsets[i, mask[i], :]
mesh_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], torch.tensor(mesh_pred, device=config['device']))[0].cpu().numpy()
gt_posed = [batch['mesh_posed'].cpu().numpy()[i, mask[i], :] for i in range(mask.shape[0])]
mesh_posed_fixed = fix_collision(torch.tensor(mesh_posed, device=config['device']), batch['mesh_body_posed'], torch.tensor(mask, device=config['device'])).cpu().numpy()
body = batch['mesh_body_posed'].cpu().numpy()
garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
collision_faces = [np.concatenate((F, faces[i] + config['smpl_res']), axis=0) for i in range(len(garment_posed))]
collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
garment_posed_fixed = [mesh_posed_fixed[i, mask[i], :] for i in range(body.shape[0])]
collision_fixed = [np.concatenate((body[i], garment_posed_fixed[i]), axis=0) for i in range(len(garment_posed))]
collision_gt = [np.concatenate((body[i], gt_posed[i]), axis=0) for i in range(len(garment_posed))]
figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i]) for i in range(4) }
return { **figs, **figs_coll, **{ type + '/images': imgs, } }
def train():
valid_iterator = iter(valid_dataloader)
step = 0
for batch in train_dataloader:
train_loss = train_step(batch)
train_accumulator.update(train_loss)
if step % 25 == 0:
data_valid = next(valid_iterator)
valid_loss = valid_step(data_valid)
valid_accumulator.update(valid_loss)
if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid')
to_log = { **to_log, **train_figs, **valid_figs }
if step == 0:
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
if not args['debug']:
wandb.log(to_log, step=step)
train_accumulator.reset()
valid_accumulator.reset()
print(step)
if step % 10000 == 0:
ckpt_saver_enc(model.encoder, step, optimizer, scaler)
ckpt_saver_dec(model.decoder, step, optimizer, scaler)
step += 1
if __name__ == '__main__':
if not args['debug']:
init_wandb()
use_amp = True
skirt_modifier = '_skirt' if config['is_skirt'] else ''
# dataset = Dataset('/data-local/127/train', config['signature'], resolutions=[512], device=device, frames_file='frames_jumpsuit.txt')
dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True, )
valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True)
train_accumulator = Accumulator()
valid_accumulator = Accumulator()
model : Model = Model(config)
optimizer = optim.Adam(model.parameters(), lr=config['lr'])
scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
# _MSSSIM = MSSSIM(data_range=2).to(config['device'])
_MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
if config['pretrained_model'] != '':
load_model(model.encoder, config['pretrained_model'] + '_enc')
load_model(model.decoder, config['pretrained_model'] + '_dec')
ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_enc', top_n=5)
ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_dec', top_n=5)
run_name = config['name']
print(f'started training: {run_name}')
if not args['debug']:
# model = torch.compile(model)
wandb.watch(model, log='all', log_freq=1000, log_graph=True)
train()
smpl/smpl_torch.py
+4 -8
View File
@@ -163,12 +163,9 @@ class SMPL(Module):
v_shaped = torch.stack(v_shaped, axis=-1) v_shaped = torch.stack(v_shaped, axis=-1)
# J = torch.matmul(self.J_regressor, v_shaped) # J = torch.matmul(self.J_regressor, v_shaped)
if verts is None: J = torch.stack([torch.matmul(self.J_regressor[torch.greater(g, 0.5).to(torch.long)], v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)]) for i, g in enumerate(betas[:,-1])])
J = torch.stack([torch.matmul(self.J_regressor[g.to(torch.long)], v_shaped[i, :, :, g.to(torch.long)]) for i, g in enumerate(betas[:,-1])]) J0 = torch.reshape(J[:,0,:], (-1, 1, 3))
else: J = J - J0
J = torch.stack([torch.matmul(self.J_regressor[torch.greater(g, 0.5).to(torch.long)], v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)]) for i, g in enumerate(betas[:,-1])])
J0 = torch.reshape(J[:,0,:], (-1, 1, 3))
J = J - J0
R_cube_big = self.rodrigues(pose.reshape(-1, 1, 3)).reshape(batch_num, -1, 3, 3) R_cube_big = self.rodrigues(pose.reshape(-1, 1, 3)).reshape(batch_num, -1, 3, 3)
if verts is not None: if verts is not None:
@@ -211,8 +208,7 @@ class SMPL(Module):
joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0] joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
weights = self.weights if verts is not None else self.weights_all T = torch.stack([torch.tensordot(results[i], self.weights[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])])
T = torch.stack([torch.tensordot(results[i], weights[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])])
rest_shape_h = torch.cat( rest_shape_h = torch.cat(
(v_posed, torch.ones((batch_num, v_posed.shape[1], 1), dtype=torch.float32).to(self.device)), dim=2 (v_posed, torch.ones((batch_num, v_posed.shape[1], 1), dtype=torch.float32).to(self.device)), dim=2
) )
utils/postprocess.py
+1 -1
View File
@@ -22,7 +22,7 @@ from losses.losses import _nearest_neighbour
""" COLLISION FUNCTIONS START """ """ COLLISION FUNCTIONS START """
# F = np.load("smpl/smpl_faces.npy") # F = np.load("smpl/smpl_faces.npy")
to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1) to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
F_fix_coll = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[None,:,:], [4,1,1]), device='cuda') F_fix_coll = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
smpl_to_keep_orig_post = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1) smpl_to_keep_orig_post = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
smpl_to_keep_orig_post = smpl_to_keep_orig_post[smpl_to_keep_orig_post < 6890] smpl_to_keep_orig_post = smpl_to_keep_orig_post[smpl_to_keep_orig_post < 6890]
F_fix_coll_orig = torch.tensor(np.tile(np.load("smpl/smpl_faces_orig.npy")[:6890,:][None,:,:], [4,1,1]), device='cuda')[:,smpl_to_keep_orig_post,:] F_fix_coll_orig = torch.tensor(np.tile(np.load("smpl/smpl_faces_orig.npy")[:6890,:][None,:,:], [4,1,1]), device='cuda')[:,smpl_to_keep_orig_post,:]