def forward_step(th_scan_meshes, smplx, init_smplx_meshes, search_tree,
pen_distance, tri_filtering_module):
"""
Performs a forward step, given smplx and scan meshes.
Then computes the losses.
"""
# forward
# verts, _, _, _ = smplx()
verts = smplx()
th_SMPLX_meshes = [
tm.from_tensors(vertices=v, faces=smplx.faces) for v in verts
]
p3d_meshes = Meshes(verts=verts, faces=smplx.faces.expand(1, -1, -1))
# losses
loss = dict()
loss['s2m'] = batch_point_to_surface(
[sm.vertices for sm in th_scan_meshes], th_SMPLX_meshes)
loss['m2s'] = batch_point_to_surface(
[sm.vertices for sm in th_SMPLX_meshes], th_scan_meshes)
loss['lap'] = torch.stack([
laplacian_loss(sc, sm)
for sc, sm in zip(init_smplx_meshes, th_SMPLX_meshes)
])
loss['offsets'] = torch.mean(torch.mean(smplx.offsets**2, axis=1), axis=1)
# loss['normal'] = mesh_normal_consistency(p3d_meshes).unsqueeze(0)
# loss['interpenetration'] = interpenetration_loss(verts, smplx.faces, search_tree, pen_distance, tri_filtering_module, 1.0)
return loss
def forward_step(th_scan_meshes, smpl, init_smpl_meshes):
"""
Performs a forward step, given smpl and scan meshes.
Then computes the losses.
"""
# forward
verts, _, _, _ = smpl()
th_smpl_meshes = [
tm.from_tensors(vertices=v, faces=smpl.faces) for v in verts
]
# p3d_meshes = Meshes(verts=verts, faces=smpl.faces.expand(1,-1,-1))
# 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['lap'] = torch.stack([
laplacian_loss(sc, sm)
for sc, sm in zip(init_smpl_meshes, th_smpl_meshes)
])
# lap1 = init_smpl_meshes[0].compute_laplacian()
# lap2 = th_smpl_meshes[0].compute_laplacian()
# loss['lap'] = (torch.sum((lap1 - lap2) ** 2, 1) * lap_weight).sum().unsqueeze(0)
loss['offsets'] = torch.mean(torch.mean(smpl.offsets**2, axis=1), axis=1)
# loss['edge'] = (edge_length(th_smpl_meshes[0]) - edge_length(init_smpl_meshes[0])).unsqueeze(0)
# loss['normal'] = mesh_normal_consistency(p3d_meshes).unsqueeze(0)
# loss['pen'] = interpenetration_loss(verts, smpl.faces.expand(1,-1,-1), search_tree, pen_distance, tri_filtering_module)
return loss
def forward_step(th_scan_meshes, smpl, init_smpl_meshes):
"""
Performs a forward step, given smpl and scan meshes.
Then computes the losses.
"""
# 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['lap'] = torch.stack([laplacian_loss(sc, sm) for sc, sm in zip(init_smpl_meshes, th_smpl_meshes)])
loss['offsets'] = torch.mean(torch.mean(smpl.offsets**2, axis=1), axis=1)
return loss
def lap_loss(pred_mesh, gt_mesh):
return torch.stack(
[laplacian_loss(sc, sm) for sc, sm in zip(pred_mesh, gt_mesh)]).mean()