From 5ab1cb5feda296dfbd5a8be64e7c693467d240cf Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Tue, 4 Jul 2023 18:22:57 +0200 Subject: [PATCH] exapnd generate function, not tested --- scripts/generate_results.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/scripts/generate_results.py b/scripts/generate_results.py index 0107780..a57d483 100644 --- a/scripts/generate_results.py +++ b/scripts/generate_results.py @@ -78,7 +78,26 @@ def step_pyramid(batch): def step(batch): - return None + if config['use_skinning']: + uv_orig = batch['uv_unposed'] + else: + uv_orig = batch['uv_posed'] + + model.eval() + pred, _, _ = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']]) + mesh_pred = uv2mesh(pred) + + if config['use_skinning']: + shape = batch['pose_shape'][:, 72:] + gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long) + shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])]) + + mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None] + mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0] + else: + mesh_pred_posed = mesh_pred + + return mesh_pred_posed def generate():