fix(loss-func): keypoint_2d_loss score tensor not sliced (+1 more)

[andrewwhitecdw]

Small fixes in global_recon/models/loss_func.py:

fix: keypoint_2d_loss score tensor not sliced with first_frame_only

Fix: Replace:

        if first_frame_only:
            loss = loss[[0]]
            num_pose += vis_frames.sum()
        else:
            num_pose += vis_frames.sum()
        loss[:10] *= first_frame_weight
        loss = (loss.sum(-1) * (score ** 2)).sum()

with:

        if first_frame_only:
            loss = loss[[0]]
            score = score[[0]]
            num_pose += 1
        else:
            num_pose += vis_frames.sum()
        loss[:10] *= first_frame_weight
        loss = (loss.sum(-1) * (score ** 2)).sum()

fix: rel_transform_loss counts invisible frames and bad bool check

Fix: Replace:

    for (i, j), rel_transform_cam in data['rel_transform_cam'].items():
        num_pose += rel_transform_cam.shape[0]
        vis_frames = person_data[i]['vis_frames'] & person_data[j]['vis_frames']
        if sum(vis_frames) == 0:
            continue

with:

    for (i, j), rel_transform_cam in data['rel_transform_cam'].items():
        vis_frames = person_data[i]['vis_frames'] & person_data[j]['vis_frames']
        if vis_frames.sum().item() == 0:
            continue
        num_pose += vis_frames.sum()

Files changed

• global_recon/models/loss_func.py