Support removing keypoints outside the image instead of clamping them to the edges

[Callidior]

🚀 The feature

When keypoints fall outside of the visible area of an image after a transformation such as cropping or rotation, they are currently clamped to one of the edges of the image, i.e., a keypoint at [-10, 25] becomes [0, 25]. It would be useful to have an option to remove such invisible keypoints instead.

Motivation, pitch

Often it makes more sense to omit keypoints that are not visible in the image instead of changing their location, which the current clamping mode does.

It would be useful to be able to specify an alternative clamping_mode for KeyPoints tensors that does not move them to the image edges but allows them to have coordinates outside of the image. This clamping mode could be combined with a new transformation SanitizeKeyPoints, similar to the existing SanitizeBoundingBoxes, which removes such invisible keypoints. Keypoints can sometimes have labels associated with them in a separate tensor, which this transformation should be able to remove as well.

Alternatives

No response

Additional context

No response

[AntoineSimoulin]

@Callidior this seems like a valid concern to me. Thanks for reporting the issue. Could you provide a minimal example to reproduce the issue?

Also, if I understand you correctly, you think the best way to address the issue would be to modify the behavior of the _clamp_keypoints and remove invisible keypoints instead of changing their location?

[Callidior]

@AntoineSimoulin Here is an example illustrating the issue:

import torch
import torchvision.transforms.v2 as T
from torchvision.tv_tensors import Image, KeyPoints

image = Image(torch.randn(3, 256, 320))
keypoints = KeyPoints(torch.tensor([[10, 10], [260, 240]]), canvas_size=(256, 320))

transforms = T.CenterCrop(224)
image_aug, keypoints_aug = transforms(image, keypoints)

print(image_aug.shape)  # [3, 224, 224]
print(keypoints_aug)  # [[0, 0], [212, 223]]

Imagine the image contains a face, for example, and the keypoints are the locations of the two eyes. After cropping, the eyes would not be visible in the image any more and their actual locations in this example would be [[-38, -6], [212, 234]].
The clamping, which is currently applied to these locations, moves them to a completely different position in the cropped image, which is not that I would need in this case.

The problem is not so much the behavior of the function _clamp_keypoints itself - it does what it is supposed to do - but that this function is called at the end of all geometric keypoint transformations, e.g., crop_keypoints, elastic_keypoints etc.
Thus, we cannot clearly determine any more which keypoints are still valid and which are not.

I would propose to handle this in a similar way as bounding boxes: The BoundingBoxes tensor class has an argument clamping_mode, which controls what clamping should be applied automatically after transforms. By default, bounding box coordinates are clamped to image boundaries, but clamping_mode can be set to None to disable automatic clamping. The user can still apply it explicitly at the end of the transformation chain by adding a ClampBoundingBoxes transform or they can remove bounding boxes outside of the image by adding a SanitizeBoundingBoxes transform.

In the case of bounding boxes, we can still identify and remove invalid bounding boxes after clamping because they have 0 area after clamping all coordinates to the image boundaries. This is not possible for keypoints, which do not have an area.

I would propose an API for keypoints similar to that of bounding boxes:

• Add a clamping_mode to KeyPoints, which can be None to disable automatic clamping of KeyPoints after transforms.
  • I currently cannot think of other options for keypoints than no clamping or clamping to image boundaries, so this argument could also be a boolean clamping flag.
• Add a SanitizeKeyPoints transform, which can be used to remove keypoints that are outside of the image area.
  • Like bounding boxes, keypoints can also be associated with labels (e.g., left eye, right eye), which are provided in a separate tensor. The transform would need to be able to remove the label belonging to an invalid keypoint as well. Functionality for this is already implemented in SanitizeBoundingBoxes.

[AntoineSimoulin]

@Callidior thanks a lot for the very detail answer, this is very clear. based on your comment I proposed a first PR (#9234) to add clamping_mode to KeyPoints. Let me know if you have time to go through it and review. This would be extremely helpful. Also, let me know how you would feel to prepare a PR to implement the SanitizeKeyPoints once we have merged this first stage.
Thanks a lot for your time!

[Callidior]

@AntoineSimoulin Thanks a lot for resolving this issue so quickly! I added my comments to the PR.

I'll try to find time to work on a SanitizeKeypoints transform.

[AntoineSimoulin]

@Callidior we probably would like to merge those changes before the next release next week. Do you think you will be able to propose a PR by early next week for the SanitizeKeypoints transform? Otherwise, I can have a look at it! Thanks for your review of #9234 also!

[Callidior]

@AntoineSimoulin I think I can open a PR on Monday or Tuesday next week.