get_all_keypoints is now get_keypoints and returns the only keypoints object in the sample (as is assumed)

Author: Alex-S-H-P

torchvision/transforms/v2/_utils.py

@@ -2,7 +2,7 @@
 
 import collections.abc
 import numbers
-from collections.abc import Iterable, Sequence
+from collections.abc import Sequence
 from contextlib import suppress
 
 from typing import Any, Callable, Literal
@@ -165,18 +165,16 @@ def get_bounding_boxes(flat_inputs: list[Any]) -> tv_tensors.BoundingBoxes:
         raise ValueError("No bounding boxes were found in the sample")
 
 
-def get_all_keypoints(flat_inputs: list[Any]) -> Iterable[tv_tensors.KeyPoints]:
-    """Yields all KeyPoints in the input.
+def get_keypoints(flat_inputs: list[Any]) -> tv_tensors.KeyPoints:
+    """Returns the KeyPoints in the input.
 
-    Raises:
-        ValueError: No KeyPoints can be found
+    Assumes only one ``KeyPoints`` object is present
     """
     generator = (inpt for inpt in flat_inputs if isinstance(inpt, tv_tensors.KeyPoints))
     try:
-        yield next(generator)
+        return next(generator)
     except StopIteration:
         raise ValueError("No Keypoints were found in the sample.")
-    return generator
 
 
 def query_chw(flat_inputs: list[Any]) -> tuple[int, int, int]: