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

direct training with normals

Browse Source
This commit is contained in:
luni
2023-05-31 21:40:18 +00:00
parent 4e759fcccc
commit f3c1e3308e
4 changed files with 174 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
configs/direct_3.py
+82
View File
@@ -0,0 +1,82 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '0'
config['name'] = 'direct_3'
config['notes'] = 'with normals input'
config['tags'] = ['torch', 'direct']
config['script'] = 'scripts/direct_train.py'
config['pretrained_model'] = ''
config['use_curvature'] = False
config['use_normals'] = True
# assert that only use curvature or normals
assert not (config['use_curvature'] and config['use_normals'])
config['batch_size'] = 4
config['train_dir'] = '/data-local/127/train'
config['valid_dir'] = '/data-local/127/valid'
config['test_dir'] = '/data-net/dataset/test_ready'
config['train_frames'] = None
config['valid_frames'] = None
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = False
config['is_cape'] = False
config['remove_global_orientation'] = True
config['device'] = 'cuda'
config['input_features'] = 9
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['latent_dim'] = 128
config['bottleneck_scale'] = 2
config['is_label'] = False
config['kernel_size'] = 7
config['bottleneck_size'] = 64 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 128, 3),
CNBlockConfig(128, 256, 9),
CNBlockConfig(256, config['bottleneck_size'], 9),
CNBlockConfig(config['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['bottleneck_size'], 256, 3),
CNBlockConfig(256, 128, 9),
CNBlockConfig(128, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
# CNBlockConfig(3, 3, 3),
]
config['encoder_block_settings'] = block_setting
config['decoder_block_settings'] = block_setting_reverse
config['activation'] = 'tanh'
config['lr'] = 0.0001
config['reg_loss_coeff'] = 0.1
config['signature'] = { 'id',
'frame_nr',
'pose_shape',
'uv_normals',
'mesh_mask',
'mesh_body_posed',
'mesh_posed',
'uv_unposed',
'mesh_unposed',
'uv_static',
'uv_mask',
'uv_body_posed',
}
configs/direct_4.py
+82
View File
@@ -0,0 +1,82 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '5'
config['name'] = 'direct_5'
config['notes'] = 'with normals input but only single garment sequence'
config['tags'] = ['torch', 'direct']
config['script'] = 'scripts/direct_train.py'
config['pretrained_model'] = ''
config['use_curvature'] = False
config['use_normals'] = True
# assert that only use curvature or normals
assert not (config['use_curvature'] and config['use_normals'])
config['batch_size'] = 4
config['train_dir'] = '/data-local/127/train'
config['valid_dir'] = '/data-local/127/train'
config['test_dir'] = '/data-net/dataset/test_ready'
config['train_frames'] = 'frames_00761.txt'
config['valid_frames'] = 'frames_00761.txt'
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = True
config['is_cape'] = False
config['remove_global_orientation'] = True
config['device'] = 'cuda'
config['input_features'] = 9
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['latent_dim'] = 128
config['bottleneck_scale'] = 2
config['is_label'] = False
config['kernel_size'] = 7
config['bottleneck_size'] = 64 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 128, 3),
CNBlockConfig(128, 256, 9),
CNBlockConfig(256, config['bottleneck_size'], 9),
CNBlockConfig(config['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['bottleneck_size'], 256, 3),
CNBlockConfig(256, 128, 9),
CNBlockConfig(128, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
# CNBlockConfig(3, 3, 3),
]
config['encoder_block_settings'] = block_setting
config['decoder_block_settings'] = block_setting_reverse
config['activation'] = 'tanh'
config['lr'] = 0.0001
config['reg_loss_coeff'] = 0.1
config['signature'] = { 'id',
'frame_nr',
'pose_shape',
'uv_normals',
'mesh_mask',
'mesh_body_posed',
'mesh_posed',
'uv_unposed',
'mesh_unposed',
'uv_static',
'uv_mask',
'uv_body_posed',
}
run.py
+1 -1
View File
@@ -1,6 +1,6 @@
import os
import sys
from configs.direct_1 import config
from configs.direct_3 import config
import argparse
# parser = argparse.ArgumentParser(description='')
scripts/direct_train.py
+9 -2
View File
@@ -76,6 +76,8 @@ def train_step(batch):
# with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp):
if config['use_curvature']:
input = torch.cat([batch['uv_curvature'], batch['uv_static']], axis=1)
elif config['use_normals']:
input = torch.cat([batch['uv_normals'], batch['uv_body_posed'], batch['uv_static']], axis=1)
else:
input = torch.cat([batch['uv_body_posed'], batch['uv_static']], axis=1)
pred = model([input, batch['uv_mask']])
@@ -121,6 +123,8 @@ def valid_step(batch):
with torch.inference_mode():
if config['use_curvature']:
input = torch.cat([batch['uv_curvature'], batch['uv_static']], axis=1)
elif config['use_normals']:
input = torch.cat([batch['uv_normals'], batch['uv_body_posed'], batch['uv_static']], axis=1)
else:
input = torch.cat([batch['uv_body_posed'], batch['uv_static']], axis=1)
pred = model([input, batch['uv_mask']])
@@ -154,6 +158,8 @@ def plot_everything(batch: dict, type: str):
with torch.inference_mode():
if config['use_curvature']:
input = torch.cat([batch['uv_curvature'], batch['uv_static']], axis=1)
elif config['use_normals']:
input = torch.cat([batch['uv_normals'], batch['uv_body_posed'], batch['uv_static']], axis=1)
else:
input = torch.cat([batch['uv_body_posed'], batch['uv_static']], axis=1)
pred = model([input, batch['uv_mask']]).detach()
@@ -164,8 +170,9 @@ def plot_everything(batch: dict, type: str):
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])
uv_normals_img = (torch.permute(batch['uv_normals'], [0,2,3,1]) + 1.1) / 2.2
pred_img = torch.permute(pred, [0,2,3,1])
img = torch.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
img = torch.concat((uv_static_img, uv_body_posed_img, uv_normals_img, orig_img, pred_img), axis=2)
imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)]
F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
@@ -273,7 +280,7 @@ if __name__ == '__main__':
if not args['debug']:
# model = torch.compile(model)
wandb.watch(model, log='all', log_freq=1000, log_graph=True)
wandb.watch(model, log='all', log_freq=10000, log_graph=True)
train()