cherry-pick missing functions from pytorch#6401

Author: pmeier

torchvision/prototype/transforms/_geometry.py

@@ -13,7 +13,7 @@
 from typing_extensions import Literal
 
 from ._transform import _RandomApplyTransform
-from ._utils import get_image_dimensions, has_any, is_simple_tensor, query_bboxes, query_image
+from ._utils import get_image_dimensions, has_any, is_simple_tensor, query_bounding_box, query_image
 
 
 class RandomHorizontalFlip(_RandomApplyTransform):
@@ -699,7 +699,7 @@ def _get_params(self, sample: Any) -> Dict[str, Any]:
         left = int(offset_width * r)
 
         if needs_crop:
-            bounding_boxes = query_bboxes(sample)
+            bounding_boxes = query_bounding_box(sample)
             bounding_boxes = F.crop(bounding_boxes, top=top, left=left, height=height, width=width)
             bounding_boxes = features.BoundingBox.new_like(
                 bounding_boxes,

torchvision/prototype/transforms/_utils.py

@@ -17,6 +17,15 @@ def query_image(sample: Any) -> Union[PIL.Image.Image, torch.Tensor, features.Im
     raise TypeError("No image was found in the sample")
 
 
+def query_bounding_box(sample: Any) -> features.BoundingBox:
+    flat_sample, _ = tree_flatten(sample)
+    for i in flat_sample:
+        if isinstance(i, features.BoundingBox):
+            return i
+
+    raise TypeError("No bounding box was found in the sample")
+
+
 def get_image_dimensions(image: Union[PIL.Image.Image, torch.Tensor, features.Image]) -> Tuple[int, int, int]:
     if isinstance(image, features.Image):
         channels = image.num_channels

torchvision/prototype/transforms/functional/__init__.py

@@ -1,5 +1,6 @@
 from torchvision.transforms import InterpolationMode  # usort: skip
 from ._meta import (
+    clamp_bounding_box,
     convert_bounding_box_format,
     convert_color_space_image_tensor,
     convert_color_space_image_pil,

torchvision/prototype/transforms/functional/_meta.py

@@ -61,6 +61,15 @@ def convert_bounding_box_format(
     return bounding_box
 
 
+def clamp_bounding_box(
+    bounding_box: torch.Tensor, format: BoundingBoxFormat, image_size: Tuple[int, int]
+) -> torch.Tensor:
+    xyxy_boxes = convert_bounding_box_format(bounding_box, format, BoundingBoxFormat.XYXY)
+    xyxy_boxes[..., 0::2].clamp_(min=0, max=image_size[1])
+    xyxy_boxes[..., 1::2].clamp_(min=0, max=image_size[0])
+    return convert_bounding_box_format(xyxy_boxes, BoundingBoxFormat.XYXY, format, copy=False)
+
+
 def _split_alpha(image: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
     return image[..., :-1, :, :], image[..., -1:, :, :]