def forward_step(th_scan_meshes, smpl, th_pose_3d=None):
"""
Performs a forward step, given smpl and scan meshes.
Then computes the losses.
"""
# Get pose prior
prior = get_prior(smpl.gender)
# forward
verts, _, _, _ = smpl()
th_smpl_meshes = [
tm.from_tensors(vertices=v, faces=smpl.faces) for v in verts
]
# losses
loss = dict()
loss['s2m'] = batch_point_to_surface(
[sm.vertices for sm in th_scan_meshes], th_smpl_meshes)
loss['m2s'] = batch_point_to_surface(
[sm.vertices for sm in th_smpl_meshes], th_scan_meshes)
loss['betas'] = torch.mean(smpl.betas**2, axis=1)
loss['pose_pr'] = prior(smpl.pose)
if th_pose_3d is not None:
loss['pose_obj'] = batch_get_pose_obj(th_pose_3d, smpl)
return loss
def forward_step_pose_only(smpl, th_pose_3d, prior_weight):
"""
Performs a forward step, given smpl and scan meshes.
Then computes the losses.
"""
# Get pose prior
prior = get_prior(smpl.gender)
# losses
loss = dict()
loss['pose_pr'] = prior(smpl.pose, prior_weight)
loss['pose_obj'] = batch_get_pose_obj(th_pose_3d, smpl, init_pose=False)
return loss