Added the KeyPoints TVTensor #8817

Alexandre-SCHOEPP · 2024-12-17T14:21:22Z

Description

Adds and integrates the KeyPoints TVTensor (requested in #8728), which is a representation of picture-attached points (or vertices) attached point

Details

Inner workings

The KeyPoints represent a tensor of shape [..., 2], which allow for arbitrarily complex structures to be represented (polygons, skeletons, or even SAM-like points prompts). Whenever the __new__ is called, the shape of the source tensor is checked.

Tensors of shape [2] are reshaped to [1, 2], similarly to BoundingBoxes.

KeyPoints, like BoundingBoxes, carry arround a canvas_size attribute which represents the scale of a batch-typical picture.

Kernels

Kernels for all operations should be supported (if I missed one, I will fix this). It merely consists of an adaptation of the code of BoundingBoxes.

Particularities

Maintainers may notice that a TYPE_CHECKING section was added that differs significantly from the implementation:

class KeyPoints(TVTensors)


    if TYPE_CHECKING:
        # EVIL: Just so that MYPY+PYLANCE+others stop shouting that everything is wrong when initializeing the TVTensor
        # Not read or defined at Runtime (only at linting time).
        # TODO: BOUNDING BOXES needs something similar
        def __init__(
            self,
            data: Any,
            *,
            dtype: Optional[torch.dtype] = None,
            device: Optional[Union[torch.device, str, int]] = None,
            requires_grad: Optional[bool] = None,
            canvas_size: Tuple[int, int],
        ):
            ...

I marked this section as EVIL since it is a trick, but it cannot generate vulnerabilities: TYPE_CHECKING is always False at runtime, and only ever True for the linter.

For the last few months, I had issues in my weird PyLance + Mypy mix with BoundingBoxes initialization. No overload is ever detected to match it. By "re-defining" it, I got to it solved on my machine.

Convertors

Added a convertor convert_box_to_points in torchvision.transorfms.v2.functional._meta exported in torchvision.transforms.v2 which (as its name states) converts a [N, 4] BoundingBoxes TVTensor into a [N, 4, 2] KeyPoints TVTensor.

Other changes

For the purposes of my custom type checking, I also changed tv_tensors.wrap to be 3.8-compatible generics.

Since wrap only ever outputs a subclass of its like argument, I used a TypeVar bound to TVTensor to ensure that type-checking passes no matter the checker used.

Methodology

Formated using ufmt
Flake8 compliance with line-length 120 enforced by editor
Documented the classes

Discussion

Since many converters of BoundingBoxes are based on chaning the bboxes to polygons, then operating on the points, I believe that there is a possibility to lower line count and increase reliability for negligeable computational latency cost by using KeyPoints kernels and converting using the method described in the details above

pytorch-bot · 2024-12-17T14:21:26Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/8817

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

Experiencing "429: Too Many Requests" on downloading actions

❌ 1 New Failure

As of commit a30737c with merge base 428a54c ():

NEW FAILURE - The following job has failed:

Docs / build / linux-job (gh)
RuntimeError: Command docker exec -t abd4b52b09d3369cf81d4f14e8cd13b6ed2cf5daa8dd8ed35c4c036742e3ba78 /exec failed with exit code 2

This comment was automatically generated by Dr. CI and updates every 15 minutes.

facebook-github-bot · 2024-12-17T14:21:29Z

Hi @Alexandre-SCHOEPP!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at [email protected]. Thanks!

facebook-github-bot · 2024-12-17T15:08:24Z

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

broomhead · 2025-04-10T03:18:31Z

I'd love to see this get included in a future release.

Alexandre-SCHOEPP · 2025-04-10T07:03:21Z

I'd love to see this get included in a future release.

Given that the PR is about 4 months old now... me too ?

NicolasHug · 2025-04-10T07:40:32Z

Thank you for your patience @Alexandre-SCHOEPP , I don't have much bandwidth at the moment, but I will try to get this reviewed and merged for 0.23

Alexandre-SCHOEPP · 2025-04-10T10:44:07Z

No worries @NicolasHug, that project is very large given the size of the team behind it, i understant that prioritization needs to happen.

AntoineSimoulin

Hey! Thanks a lot for the cool PR. I do think it will be a very relevant addition to torchvision. I feel the PR is well aligned with the structure of the codebase. However, this is introduces a large number of changes, so I propose we make a couple iterations to fully align with the main branch before merging. How does it sound?

I left some comments in the code, but here are my high level questions and remarks.

Regarding parameter naming, let's make sure we are sticking to "keypoint" for user-facing functions and not use "kp" and "kpoint" there;
I am not sure about the _xyxy_to_points output. Can we include some tests for both _xyxy_to_points and convert_box_to_points to make sure the input/output is as expected?
We a keypoint rcnn model that support keypoints in torchvision. I'm not sure we should change the code because of the risk of breaking it, but can use this as some advanced integration test?
I think we are missing a certain number of tests for the transforms: TestResize.test_kernel_key_points, TestResize.test_functional for key_points, TestResize.test_key_points_correctness, TestResize.test_noop, TestResize.test_no_regression_5405, TestHorizontalFlip.test_kernel_key_points, TestHorizontalFlip.test_bounding_boxes_correctness ... Would it be eventually possible to add those?
More generally, the keypoints are defined here a set of independent points. Is it the usual definition, can those be used to define arbitrary polygons? In those case, will some other functionalities be needed?

Thanks again for your time submitting this PR. Let me know what you think about my comments. On my side, I am excited to work on this with you and make sure we can merge the feature into torchvision.

Best,
Antoine

AntoineSimoulin · 2025-04-29T12:33:44Z

torchvision/transforms/v2/_utils.py

+def get_all_keypoints(flat_inputs: List[Any]) -> Iterable[tv_tensors.KeyPoints]:
+    """Yields all KeyPoints in the input.
+
+    Raises:
+        ValueError: No KeyPoints can be found
+    """
+    generator = (inpt for inpt in flat_inputs if isinstance(inpt, tv_tensors.KeyPoints))
+    try:
+        yield next(generator)
+    except StopIteration:
+        raise ValueError("No Keypoints were found in the sample.")
+    return generator
+
+


It seems this get_all_keypoints function is not used elsewhere. Should we delete it if this is not needed in the current state of the PR?

This function is meant to parallel the one for bounding boxes, the schematic of which I followed here. I don't necessarily have an attachment to it, but I felt like it was better to have it if someone who worked on the BBoxes had feedback on it.

torchvision/tv_tensors/_keypoints.py

AntoineSimoulin · 2025-04-29T12:39:37Z

torchvision/tv_tensors/_keypoints.py

+        dtype: Optional[torch.dtype] = None,
+        device: Optional[Union[torch.device, str, int]] = None,
+        requires_grad: Optional[bool] = None,
+        canvas_size: Tuple[int, int],


Maybe align the order of the arguments with the BoundingBoxes class? In this case, canvas_size should be the 4th argument and not the last.

I agree for the sake of consistent code styling, however, given that this is a keyword-only argument, this difference isn't visible at runtime outside of inspecting the function's attributes

I'm putting it in third position since in BoundingBoxes the 3rd argument is the format, which does not exist in KeyPoints.

AntoineSimoulin · 2025-04-29T12:40:22Z

torchvision/tv_tensors/_keypoints.py

+            # Since a Tensor is a sequence of Tensor, had it not been the case, we may have had silent
+            # or complex errors
+            output = tuple(KeyPoints(part, canvas_size=canvas_size) for part in output)
+        elif isinstance(output, MutableSequence):


What is the purpose of this case here?

This case handles lists, and list-like objects.

Unlike in the case of tuples, (which are non-mutable sequences), I can't reuse the container object. Here, I can.

I agree that it needs to be changed, however. Custom-made non-mutable sequences are not supported, which is an oversight on my part.

The valid logic would be :

if isinstance(output, torch.Tensor) and not isinstance(output, KeyPoints): output = KeyPoints(output, canvas_size=canvas_size) elif isinstance(output, MutableSequence): # For lists and list-like object we don't try to create a new object, we just set the values in the list for i, part in enumerate(output): output[i] = KeyPoints(part, canvas_size=canvas_size) elif isinstance(output, Sequence): # Non-mutable sequences handled here (like tuples) # NB: output is checked against sequence because it has already been checked against Tensor # Since a Tensor is a sequence of Tensor, had it not been the case, we may have had silent # or complex errors output = tuple(KeyPoints(part, canvas_size=canvas_size) for part in output) return output

torchvision/transforms/v2/functional/_geometry.py

torchvision/transforms/v2/functional/_meta.py

test/common_utils.py

AntoineSimoulin · 2025-04-29T12:59:38Z

torchvision/transforms/v2/functional/_meta.py

+        new_format=BoundingBoxFormat.XYXY,
+        inplace=False,
+    )
+    return tv_tensors.KeyPoints(_xyxy_to_points(bbox), canvas_size=bounding_boxes.canvas_size)


We should also make sure this function is covered in the test case.

Alexandre-SCHOEPP

Thanks for the review. I marked every comment where I agreed with the need to change things with the word agree in my response just so I can ensure to not miss one whenever I get the time to fix it

test/common_utils.py

Alexandre-SCHOEPP · 2025-04-29T13:19:16Z

torchvision/transforms/v2/_utils.py

+def get_all_keypoints(flat_inputs: List[Any]) -> Iterable[tv_tensors.KeyPoints]:
+    """Yields all KeyPoints in the input.
+
+    Raises:
+        ValueError: No KeyPoints can be found
+    """
+    generator = (inpt for inpt in flat_inputs if isinstance(inpt, tv_tensors.KeyPoints))
+    try:
+        yield next(generator)
+    except StopIteration:
+        raise ValueError("No Keypoints were found in the sample.")
+    return generator
+
+


This function is meant to parallel the one for bounding boxes, the schematic of which I followed here. I don't necessarily have an attachment to it, but I felt like it was better to have it if someone who worked on the BBoxes had feedback on it.

torchvision/transforms/v2/functional/_meta.py

Alexandre-SCHOEPP · 2025-04-29T13:21:07Z

torchvision/transforms/v2/functional/_meta.py

+        new_format=BoundingBoxFormat.XYXY,
+        inplace=False,
+    )
+    return tv_tensors.KeyPoints(_xyxy_to_points(bbox), canvas_size=bounding_boxes.canvas_size)


torchvision/tv_tensors/_keypoints.py

Alexandre-SCHOEPP · 2025-04-29T13:23:59Z

torchvision/tv_tensors/_keypoints.py

+        dtype: Optional[torch.dtype] = None,
+        device: Optional[Union[torch.device, str, int]] = None,
+        requires_grad: Optional[bool] = None,
+        canvas_size: Tuple[int, int],


I agree for the sake of consistent code styling, however, given that this is a keyword-only argument, this difference isn't visible at runtime outside of inspecting the function's attributes

Alexandre-SCHOEPP · 2025-04-29T13:25:06Z

torchvision/tv_tensors/_keypoints.py

+        points.canvas_size = canvas_size
+        return points
+
+    if TYPE_CHECKING:


I had no need for type:ignore. This solution works, and quite frankly, I think should be the standard in the library (despite marking it jokingly as EVIL)

Alexandre-SCHOEPP · 2025-04-29T13:28:39Z

torchvision/tv_tensors/_keypoints.py

+            # Since a Tensor is a sequence of Tensor, had it not been the case, we may have had silent
+            # or complex errors
+            output = tuple(KeyPoints(part, canvas_size=canvas_size) for part in output)
+        elif isinstance(output, MutableSequence):


This case handles lists, and list-like objects.

Unlike in the case of tuples, (which are non-mutable sequences), I can't reuse the container object. Here, I can.

I agree that it needs to be changed, however. Custom-made non-mutable sequences are not supported, which is an oversight on my part.

The valid logic would be :

if isinstance(output, torch.Tensor) and not isinstance(output, KeyPoints): output = KeyPoints(output, canvas_size=canvas_size) elif isinstance(output, MutableSequence): # For lists and list-like object we don't try to create a new object, we just set the values in the list for i, part in enumerate(output): output[i] = KeyPoints(part, canvas_size=canvas_size) elif isinstance(output, Sequence): # Non-mutable sequences handled here (like tuples) # NB: output is checked against sequence because it has already been checked against Tensor # Since a Tensor is a sequence of Tensor, had it not been the case, we may have had silent # or complex errors output = tuple(KeyPoints(part, canvas_size=canvas_size) for part in output) return output

Alexandre-SCHOEPP

_geometry.py was hidden by the change reviewer, and so I failed to answer a few comments. Here's the change

torchvision/transforms/v2/functional/_geometry.py

torchvision/tv_tensors/_keypoints.py

broomhead · 2025-04-30T15:30:04Z

@Alexandre-SCHOEPP I pulled this fork and compiled it for my own project, as I really need keypoints that match the TVTensor design, particularly for transforms.

It's working well for me. One thing I did find is KeyPoints needs to be added to transforms like SanitizeBoundingBoxes. The issue arises that when a transform removes data on some criterion, like bounding box size, the associated data should be removed with it. Well, KeyPoints are not getting removed. I looked at the code and it was as simple as adding KeyPoints into a check in one line. For now I just overrode the class and made the one-line change locally:

# This is a direct copy-and-paste of SanitzeBoundingBoxes, with KeyPoints added.
@register(name="SanitizeBoundingBoxesWithKeyPoints")
class SanitizeBoundingBoxesWithKeyPoints(T.SanitizeBoundingBoxes):
    def transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
        is_label = params["labels"] is not None and any(inpt is label for label in params["labels"])
        is_bounding_boxes_or_mask_or_keypoints = isinstance(
            inpt, (BoundingBoxes, Mask, KeyPoints)
        )

        if not (is_label or is_bounding_boxes_or_mask_or_keypoints):
            return inpt

        output = inpt[params["valid"]]
        if is_label:
            return output
        else:
            wrapped = wrap(output, like=inpt)
            return wrapped

that isinstance check is all we need; I just added KeyPoints to the list. Now, I'm new to PyTorch, so I'm not sure if there are other transforms that require this sort of check as well. None of the others used in my project appear to.

Alexandre-SCHOEPP · 2025-04-30T15:37:42Z

@Alexandre-SCHOEPP I pulled this fork and compiled it for my own project, as I really need keypoints that match the TVTensor design, particularly for transforms.

Thanks ! Always good to have this feedback.

One thing I did find is KeyPoints needs to be added to transforms like SanitizeBoundingBoxes

I am waiting for feedback on it by the maintainers (@NicolasHug, @AntoineSimoulin). I am willing to do it as I correct all other mistakes, but they may want it to be another PR.

… sequences in wrap_output

…oxes_dispatch

…PP/main

This reverts commit 5825706.

Alexandre-SCHOEPP · 2025-06-03T14:12:10Z

@NicolasHug Why remove the type annotation for the init method of KeyPoints (commit 08e8843) ? Is it breaking some things ? (I use that specific trick a lot)

NicolasHug · 2025-06-03T15:59:20Z

Hi @Alexandre-SCHOEPP ,

I'm in the process of finalizing this PR, there are a few tests I need to add / fix and also for some functionalities, I need to make sure we're supporting what's strictly needed. The type annotations tend to get in the way of doing that, and they're not strictly needed for this feature. I don't mind re-considering in the future but for the sake of this PR, my job as a reviewer is a lot simpler if I can ignore the annotation logic

…PP/main

NicolasHug · 2025-06-04T13:10:49Z

torchvision/transforms/v2/functional/_geometry.py

+
+    index_xy = keypoints.to(dtype=torch.long)
+    index_x, index_y = index_xy[:, 0], index_xy[:, 1]
+    # Unlike bounding boxes, this may not work well.


@Alexandre-SCHOEPP can you share more about this comment? What may not work well, and how? Thanks!

It's a rather old comment, so i don't fully remember. The issue I see right now is that the code relies on keypoints being 2D (of shape N, 2), which is not always correct.

This comment also precedes a clamping that does not occur during the bounding boxes version. I think this is because in the bounding boxes, points outside the edge are erased, and in Keypoints they don't necessarily are.

Alexandre-SCHOEPP added 9 commits December 12, 2024 17:39

Added Keypoints to the library

8253305

Improved KeyPoints to be exported

484561d

Added kernels to support the keypoints

3255890

Added tests for keypoints

7436636

Applied ufmt formatting

b35cba6

Fixed the bugs found while testing

a19ec0b

Improved documentation to take KeyPoints into account

5f4b188

Applied ufmt check

cabce1c

Fixed the hflip not being along the right coordinate

d1b27ad

facebook-github-bot added the cla signed label Dec 17, 2024

Merge branch 'main' into main

6fa38f4

twsl mentioned this pull request Dec 25, 2024

Improve type hints in tv_tensor.wrap #8830

Open

Alexandre-SCHOEPP added 3 commits February 10, 2025 16:56

Merge branch 'main' into main

05e4ad6

Merge branch 'main' into main

03dc6c8

Merge branch 'main' into main

d4d087c

AntoineSimoulin reviewed Apr 29, 2025

View reviewed changes

Alexandre-SCHOEPP commented Apr 29, 2025

View reviewed changes

Alex-S-H-P added 3 commits April 30, 2025 19:36

Fixed order of arguments

5a8c5b4

Reworked logic of the conditions to better handle mutable/non mutable…

dea31e2

… sequences in wrap_output

Renamed out variable to be more similar with _resized_crop_bounding_b…

71e20a5

…oxes_dispatch

NicolasHug added 11 commits June 3, 2025 13:31

Merge branch 'main' of github.com:pytorch/vision into Alexandre-SCHOE…

ef36e93

…PP/main

Fix

f71bb8a

Fix

5825706

Fix

6a14fb9

Revert "Fix"

35650a4

This reverts commit 5825706.

Remove some annotations

08e8843

Align signatures

f03f958

Proper fix for test_to_tv_tensor_reference

d04a3e3

lint

0b1da89

Add more keypoints tests in test_tv_tensors

f4c8da8

Revert changes to SanitizeBoundingBoxes

48289d2

NicolasHug added 4 commits June 3, 2025 15:40

Bunch of stuff

fc528ff

Remove some stuff

ec459e6

Fix variable name

c499ef5

Cleanup untested stuff

e26a07c

NicolasHug added 10 commits June 4, 2025 10:04

Merge branch 'main' of github.com:pytorch/vision into Alexandre-SCHOE…

8e5ab22

…PP/main

Add todos and complete error messages

b91583e

Add test for get_size

6dc2e33

Simplify test

0db21e0

Add more tests

d4b130d

More tests

2e28525

Add more test_functional_signature

51bd300

Add more rotate tests

34ba90c

Add pad and crop tests

a43e534

Add tests for perspective

3441ab9

NicolasHug reviewed Jun 4, 2025

View reviewed changes

NicolasHug added 2 commits June 4, 2025 14:16

More tests

74688d6

lint + Temporarily fix the weird TestColorJitter test failure

a30737c

Added the KeyPoints TVTensor #8817

Are you sure you want to change the base?

Added the KeyPoints TVTensor #8817

Uh oh!

Conversation

Alexandre-SCHOEPP commented Dec 17, 2024

Description

Details

Inner workings

Kernels

Particularities

Convertors

Other changes

Methodology

Discussion

Uh oh!

pytorch-bot bot commented Dec 17, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/8817

❗ 1 Active SEVs

❌ 1 New Failure

Uh oh!

facebook-github-bot commented Dec 17, 2024

Action Required

Process

Uh oh!

facebook-github-bot commented Dec 17, 2024

Uh oh!

broomhead commented Apr 10, 2025

Uh oh!

Alexandre-SCHOEPP commented Apr 10, 2025

Uh oh!

NicolasHug commented Apr 10, 2025

Uh oh!

Alexandre-SCHOEPP commented Apr 10, 2025

Uh oh!

AntoineSimoulin left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Alexandre-SCHOEPP Apr 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Alexandre-SCHOEPP left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

pytorch-bot bot commented Dec 17, 2024 •

edited

Loading

Alexandre-SCHOEPP Apr 29, 2025 •

edited

Loading

Alexandre-SCHOEPP Apr 29, 2025 •

edited

Loading

Alexandre-SCHOEPP commented Apr 30, 2025 •

edited

Loading