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
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()