feat: add functional center crop on mask

Author: Federico Pozzi

test/test_prototype_transforms_functional.py

@@ -1,6 +1,7 @@
 import functools
 import itertools
 import math
+from unittest.mock import patch, Mock
 
 import numpy as np
 import pytest
@@ -380,6 +381,15 @@ def pad_segmentation_mask():
         yield SampleInput(mask, padding=padding, padding_mode=padding_mode)
 
 
+@register_kernel_info_from_sample_inputs_fn
+def center_crop_segmentation_mask():
+    for mask, output_size in itertools.product(
+        make_segmentation_masks(),
+        [[4, 3], [42, 70], [4]],  # crop sizes < image sizes, crop_sizes > image sizes, single crop size
+    ):
+        yield SampleInput(mask, output_size)
+
+
 @pytest.mark.parametrize(
     "kernel",
     [
@@ -1085,3 +1095,25 @@ def parse_padding():
 
         expected_mask = _compute_expected_mask()
         torch.testing.assert_close(out_mask, expected_mask)
+
+
+@pytest.mark.parametrize("device", cpu_and_gpu())
+def test_correctness_center_crop_segmentation_mask_on_fixed_input(device):
+    mask = torch.ones((1, 6, 6), dtype=torch.long, device=device)
+    mask[:, 1:5, 2:4] = 0
+
+    out_mask = F.center_crop_segmentation_mask(mask, [2])
+    expected_mask = torch.zeros((1, 4, 2), dtype=torch.long, device=device)
+    torch.testing.assert_close(out_mask, expected_mask)
+
+
+@pytest.mark.parametrize("output_size", [[4, 3], [4], [7, 7]])
+@patch("torchvision.prototype.transforms.functional._geometry.center_crop_image_tensor")
+def test_correctness_center_crop_segmentation_mask(center_crop_mock, output_size):
+    mask, expected = Mock(spec=torch.Tensor), Mock(spec=torch.Tensor)
+    center_crop_mock.return_value = expected
+
+    out_mask = F.center_crop_segmentation_mask(mask, output_size)
+
+    center_crop_mock.assert_called_once_with(img=mask, output_size=output_size)
+    assert expected is out_mask

torchvision/prototype/transforms/functional/__init__.py

@@ -45,6 +45,7 @@
     resize_image_tensor,
     resize_image_pil,
     resize_segmentation_mask,
+    center_crop_segmentation_mask,
     center_crop_image_tensor,
     center_crop_image_pil,
     resized_crop_bounding_box,

torchvision/prototype/transforms/functional/_geometry.py

@@ -530,6 +530,10 @@ def center_crop_image_pil(img: PIL.Image.Image, output_size: List[int]) -> PIL.I
     return crop_image_pil(img, crop_top, crop_left, crop_height, crop_width)
 
 
+def center_crop_segmentation_mask(segmentation_mask: torch.Tensor, output_size: List[int]) -> torch.Tensor:
+    return center_crop_image_tensor(img=segmentation_mask, output_size=output_size)
+
+
 def resized_crop_image_tensor(
     img: torch.Tensor,
     top: int,