Added the KeyPoints TVTensor

docs/source/tv_tensors.rst

@@ -21,6 +21,7 @@ info.
 
     Image
     Video
+    KeyPoints
     BoundingBoxFormat
     BoundingBoxes
     Mask

gallery/transforms/plot_tv_tensors.py

@@ -46,11 +46,12 @@
 # Under the hood, they are needed in :mod:`torchvision.transforms.v2` to correctly dispatch to the appropriate function
 # for the input data.
 #
-# :mod:`torchvision.tv_tensors` supports four types of TVTensors:
+# :mod:`torchvision.tv_tensors` supports five types of TVTensors:
 #
 # * :class:`~torchvision.tv_tensors.Image`
 # * :class:`~torchvision.tv_tensors.Video`
 # * :class:`~torchvision.tv_tensors.BoundingBoxes`
+# * :class:`~torchvision.tv_tensors.KeyPoints`
 # * :class:`~torchvision.tv_tensors.Mask`
 #
 # What can I do with a TVTensor?
@@ -96,6 +97,7 @@
 # :class:`~torchvision.tv_tensors.BoundingBoxes` requires the coordinate format as well as the size of the
 # corresponding image (``canvas_size``) alongside the actual values. These
 # metadata are required to properly transform the bounding boxes.
+# In a similar fashion, :class:`~torchvision.tv_tensors.KeyPoints` also require the ``canvas_size`` metadata to be added.
 
 bboxes = tv_tensors.BoundingBoxes(
     [[17, 16, 344, 495], [0, 10, 0, 10]],
@@ -104,6 +106,13 @@
 )
 print(bboxes)
 
+
+keypoints = tv_tensors.KeyPoints(
+    [[17, 16], [344, 495], [0, 10], [0, 10]],
+    canvas_size=image.shape[-2:]
+)
+print(keypoints)
+
 # %%
 # Using ``tv_tensors.wrap()``
 # ^^^^^^^^^^^^^^^^^^^^^^^^^^^

test/common_utils.py

@@ -9,6 +9,7 @@
 import sys
 import tempfile
 import warnings
+from collections.abc import Sequence
 from subprocess import CalledProcessError, check_output, STDOUT
 
 import numpy as np
@@ -400,6 +401,20 @@ def make_image_pil(*args, **kwargs):
     return to_pil_image(make_image(*args, **kwargs))
 
 
+def make_keypoints(
+    canvas_size: tuple[int, int] = DEFAULT_SIZE, *, num_points: int | Sequence[int] = 4, dtype=None, device="cpu"
+) -> tv_tensors.KeyPoints:
+    """Make the KeyPoints for testing purposes"""
+    if isinstance(num_points, int):
+        num_points = [num_points]
+    single_coord_shape: tuple[int, ...] = tuple(num_points) + (1,)
+    y = torch.randint(0, canvas_size[0] - 1, single_coord_shape, dtype=dtype, device=device)
+    x = torch.randint(0, canvas_size[1] - 1, single_coord_shape, dtype=dtype, device=device)
+    points = torch.cat((x, y), dim=-1)
+    keypoints = tv_tensors.KeyPoints(points, canvas_size=canvas_size)
+    return keypoints
+
+
 def make_bounding_boxes(
     canvas_size=DEFAULT_SIZE,
     *,