How to calculate xy-keypoints during training?

I’d like to implement a new loss for the keypoint head. For this loss, I need the xy-coordinates of the predicted keypoints. However, in _forward_keypoint and keypoint_rcnn_loss, I only have access to keypoint_logits. To convert keypoint_logits into xy-keypoints, I think I need to apply the steps in keypoint_rcnn_inference.

Now I have two questions:

1. Is this approach correct so far?
2. For the conversion from keypoint_logits into xy-keypoints, I need to specify the bounding boxes of the instances. In keypoint_rcnn_inference the predicted boxes from the box head are used. However, during training I don’t have access to this information yet. Can I just use the proposal boxes?

Thanks in advance.

[maxfrei750]

@cognitiveRobot I used the function heatmaps_to_keypoints with the proposed bounding boxes and implemented a custom_loss_function which operates on the xy-keypoints:

def custom_keypoint_rcnn_loss(pred_keypoint_logits, instances):
    """
    Arguments:
        pred_keypoint_logits (Tensor): A tensor of shape (N, K, S, S) where N is the total number
            of instances in the batch, K is the number of keypoints, and S is the side length
            of the keypoint heatmap. The values are spatial logits.
        instances (list[Instances]): A list of M Instances, where M is the batch size.
            These instances are predictions from the model
            that are in 1:1 correspondence with pred_keypoint_logits.
            Each Instances should contain a `gt_keypoints` field containing a `structures.Keypoint`
            instance.
    Returns a scalar tensor containing the loss.
    """

    device = pred_keypoint_logits.device

    # flatten all bboxes from all images together (list[Boxes] -> Rx4 tensor)
    bboxes_flat = cat([b.proposal_boxes.tensor for b in instances], dim=0)

    keypoint_results = heatmaps_to_keypoints(pred_keypoint_logits, bboxes_flat)
    num_instances_per_image = [len(i) for i in instances]
    keypoint_results = keypoint_results.split(num_instances_per_image, dim=0)

    instance_losses = torch.Tensor().to(device)
    instance_weights = torch.Tensor().to(device)

    for keypoint_results_per_image, instances_per_image in zip(keypoint_results, instances):
        # keypoint_results_per_image is (num instances)x(num keypoints)x(x, y, score, prob)
        keypoints_gt_xy = instances_per_image.gt_keypoints.tensor[:, :, :2]
        keypoints_xy = keypoint_results_per_image[:, :, [0, 1]]

        keypoints_p = keypoint_results_per_image[:, :, 3]
        instance_weights = torch.cat([instance_weights, keypoints_p.sum(1)])
        instance_losses_per_image = torch.Tensor().to(device)

        for instance_keypoints_xy, instance_keypoints_gt_xy in zip(keypoints_xy, keypoints_gt_xy):

            # You need to implement the custom_loss_function yourself.
            instance_loss = custom_loss_function(instance_keypoints_xy, instance_keypoints_gt_xy)

            instance_losses_per_image = torch.cat(
                [instance_losses_per_image, instance_loss.unsqueeze(0)]
            )

        instance_losses = torch.cat([instance_losses, instance_losses_per_image])

    keypoint_loss = torch.sum(instance_losses * instance_weights)

    # Maybe you need normalization.
    # normalizer = sum(num_instances_per_image)
    # keypoint_loss /= normalizer

    return keypoint_loss

Ultimately, I did not end up using this loss function but stuck to the default keypoint loss, so this function is not very well tested. Also, it is obviously shamelessly copied from various places in the original code.

[cognitiveRobot]

@maxfrei750, thanks for sharing. If I test, I will let you know.