exapnd generate function, not tested

2023-07-04 18:22:57 +02:00
parent 9a81b82621
commit 5ab1cb5fed
1 changed files with 20 additions and 1 deletions
@@ -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():