[fbsync] port tests for F.pad and transforms.Pad (#7939)

NicolasHug · facebook-github-bot · commit 4a6100e1b38d · 2023-09-26T04:55:23.000-07:00
Reviewed By: matteobettini

Differential Revision: D49600776

fbshipit-source-id: d98db4c1415f85a210cfc5905d292ec0d76b2f65
diff --git a/test/test_transforms_v2.py b/test/test_transforms_v2.py
@@ -390,21 +390,6 @@ def was_applied(output, inpt):
         assert transform.was_applied(output, input)
 
 
-class TestPad:
-    def test_assertions(self):
-        with pytest.raises(TypeError, match="Got inappropriate padding arg"):
-            transforms.Pad("abc")
-
-        with pytest.raises(ValueError, match="Padding must be an int or a 1, 2, or 4"):
-            transforms.Pad([-0.7, 0, 0.7])
-
-        with pytest.raises(TypeError, match="Got inappropriate fill arg"):
-            transforms.Pad(12, fill="abc")
-
-        with pytest.raises(ValueError, match="Padding mode should be either"):
-            transforms.Pad(12, padding_mode="abc")
-
-
 class TestRandomZoomOut:
     def test_assertions(self):
         with pytest.raises(TypeError, match="Got inappropriate fill arg"):
diff --git a/test/test_transforms_v2_consistency.py b/test/test_transforms_v2_consistency.py
@@ -109,21 +109,6 @@ def __init__(
         ],
         make_images_kwargs=dict(DEFAULT_MAKE_IMAGES_KWARGS, sizes=[(20, 19)]),
     ),
-    ConsistencyConfig(
-        v2_transforms.Pad,
-        legacy_transforms.Pad,
-        [
-            NotScriptableArgsKwargs(3),
-            ArgsKwargs([3]),
-            ArgsKwargs([2, 3]),
-            ArgsKwargs([3, 2, 1, 4]),
-            NotScriptableArgsKwargs(5, fill=1, padding_mode="constant"),
-            ArgsKwargs([5], fill=1, padding_mode="constant"),
-            NotScriptableArgsKwargs(5, padding_mode="edge"),
-            NotScriptableArgsKwargs(5, padding_mode="reflect"),
-            NotScriptableArgsKwargs(5, padding_mode="symmetric"),
-        ],
-    ),
     *[
         ConsistencyConfig(
             v2_transforms.LinearTransformation,
diff --git a/test/test_transforms_v2_functional.py b/test/test_transforms_v2_functional.py
@@ -524,75 +524,6 @@ def test_tv_tensor_explicit_metadata(self, metadata):
 #  `transforms_v2_kernel_infos.py`
 
 
-def _parse_padding(padding):
-    if isinstance(padding, int):
-        return [padding] * 4
-    if isinstance(padding, list):
-        if len(padding) == 1:
-            return padding * 4
-        if len(padding) == 2:
-            return padding * 2  # [left, up, right, down]
-
-    return padding
-
-
-@pytest.mark.parametrize("device", cpu_and_cuda())
-@pytest.mark.parametrize("padding", [[1], [1, 1], [1, 1, 2, 2]])
-def test_correctness_pad_bounding_boxes(device, padding):
-    def _compute_expected_bbox(bbox, format, padding_):
-        pad_left, pad_up, _, _ = _parse_padding(padding_)
-
-        dtype = bbox.dtype
-        bbox = (
-            bbox.clone()
-            if format == tv_tensors.BoundingBoxFormat.XYXY
-            else convert_bounding_box_format(bbox, old_format=format, new_format=tv_tensors.BoundingBoxFormat.XYXY)
-        )
-
-        bbox[0::2] += pad_left
-        bbox[1::2] += pad_up
-
-        bbox = convert_bounding_box_format(bbox, old_format=tv_tensors.BoundingBoxFormat.XYXY, new_format=format)
-        if bbox.dtype != dtype:
-            # Temporary cast to original dtype
-            # e.g. float32 -> int
-            bbox = bbox.to(dtype)
-        return bbox
-
-    def _compute_expected_canvas_size(bbox, padding_):
-        pad_left, pad_up, pad_right, pad_down = _parse_padding(padding_)
-        height, width = bbox.canvas_size
-        return height + pad_up + pad_down, width + pad_left + pad_right
-
-    for bboxes in make_multiple_bounding_boxes(extra_dims=((4,),)):
-        bboxes = bboxes.to(device)
-        bboxes_format = bboxes.format
-        bboxes_canvas_size = bboxes.canvas_size
-
-        output_boxes, output_canvas_size = F.pad_bounding_boxes(
-            bboxes, format=bboxes_format, canvas_size=bboxes_canvas_size, padding=padding
-        )
-
-        torch.testing.assert_close(output_canvas_size, _compute_expected_canvas_size(bboxes, padding))
-
-        expected_bboxes = torch.stack(
-            [_compute_expected_bbox(b, bboxes_format, padding) for b in bboxes.reshape(-1, 4).unbind()]
-        ).reshape(bboxes.shape)
-
-        torch.testing.assert_close(output_boxes, expected_bboxes, atol=1, rtol=0)
-
-
-@pytest.mark.parametrize("device", cpu_and_cuda())
-def test_correctness_pad_segmentation_mask_on_fixed_input(device):
-    mask = torch.ones((1, 3, 3), dtype=torch.long, device=device)
-
-    out_mask = F.pad_mask(mask, padding=[1, 1, 1, 1])
-
-    expected_mask = torch.zeros((1, 5, 5), dtype=torch.long, device=device)
-    expected_mask[:, 1:-1, 1:-1] = 1
-    torch.testing.assert_close(out_mask, expected_mask)
-
-
 @pytest.mark.parametrize("device", cpu_and_cuda())
 @pytest.mark.parametrize(
     "startpoints, endpoints",
diff --git a/test/test_transforms_v2_refactored.py b/test/test_transforms_v2_refactored.py
@@ -3346,3 +3346,171 @@ def test_transform_errors_warnings(self):
         for param in ["scale", "ratio"]:
             with pytest.warns(match="Scale and ratio should be of kind"):
                 transforms.RandomResizedCrop(size=self.INPUT_SIZE, **{param: [1, 0]})
+
+
+class TestPad:
+    EXHAUSTIVE_TYPE_PADDINGS = [1, (1,), (1, 2), (1, 2, 3, 4), [1], [1, 2], [1, 2, 3, 4]]
+    CORRECTNESS_PADDINGS = [
+        padding
+        for padding in EXHAUSTIVE_TYPE_PADDINGS
+        if isinstance(padding, int) or isinstance(padding, list) and len(padding) > 1
+    ]
+    PADDING_MODES = ["constant", "symmetric", "edge", "reflect"]
+
+    @param_value_parametrization(
+        padding=EXHAUSTIVE_TYPE_PADDINGS,
+        fill=EXHAUSTIVE_TYPE_FILLS,
+        padding_mode=PADDING_MODES,
+    )
+    @pytest.mark.parametrize("dtype", [torch.uint8, torch.float32])
+    @pytest.mark.parametrize("device", cpu_and_cuda())
+    def test_kernel_image(self, param, value, dtype, device):
+        if param == "fill":
+            value = adapt_fill(value, dtype=dtype)
+        kwargs = {param: value}
+        if param != "padding":
+            kwargs["padding"] = [1]
+
+        image = make_image(dtype=dtype, device=device)
+
+        check_kernel(
+            F.pad_image,
+            image,
+            **kwargs,
+            check_scripted_vs_eager=not (
+                (param == "padding" and isinstance(value, int))
+                # See https://github.com/pytorch/vision/pull/7252#issue-1585585521 for details
+                or (
+                    param == "fill"
+                    and (
+                        isinstance(value, tuple) or (isinstance(value, list) and any(isinstance(v, int) for v in value))
+                    )
+                )
+            ),
+        )
+
+    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
+    def test_kernel_bounding_boxes(self, format):
+        bounding_boxes = make_bounding_boxes(format=format)
+        check_kernel(
+            F.pad_bounding_boxes,
+            bounding_boxes,
+            format=bounding_boxes.format,
+            canvas_size=bounding_boxes.canvas_size,
+            padding=[1],
+        )
+
+    @pytest.mark.parametrize("padding_mode", ["symmetric", "edge", "reflect"])
+    def test_kernel_bounding_boxes_errors(self, padding_mode):
+        bounding_boxes = make_bounding_boxes()
+        with pytest.raises(ValueError, match=f"'{padding_mode}' is not supported"):
+            F.pad_bounding_boxes(
+                bounding_boxes,
+                format=bounding_boxes.format,
+                canvas_size=bounding_boxes.canvas_size,
+                padding=[1],
+                padding_mode=padding_mode,
+            )
+
+    @pytest.mark.parametrize("make_mask", [make_segmentation_mask, make_detection_mask])
+    def test_kernel_mask(self, make_mask):
+        check_kernel(F.pad_mask, make_mask(), padding=[1])
+
+    @pytest.mark.parametrize("fill", [[1], (0,), [1, 0, 1], (0, 1, 0)])
+    def test_kernel_mask_errors(self, fill):
+        with pytest.raises(ValueError, match="Non-scalar fill value is not supported"):
+            check_kernel(F.pad_mask, make_segmentation_mask(), padding=[1], fill=fill)
+
+    @pytest.mark.parametrize(
+        "make_input",
+        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
+    )
+    def test_functional(self, make_input):
+        check_functional(F.pad, make_input(), padding=[1])
+
+    @pytest.mark.parametrize(
+        ("kernel", "input_type"),
+        [
+            (F.pad_image, torch.Tensor),
+            # The PIL kernel uses fill=0 as default rather than fill=None as all others.
+            # Since the whole fill story is already really inconsistent, we won't introduce yet another case to allow
+            # for this test to pass.
+            # See https://github.com/pytorch/vision/issues/6623 for a discussion.
+            # (F._pad_image_pil, PIL.Image.Image),
+            (F.pad_image, tv_tensors.Image),
+            (F.pad_bounding_boxes, tv_tensors.BoundingBoxes),
+            (F.pad_mask, tv_tensors.Mask),
+            (F.pad_video, tv_tensors.Video),
+        ],
+    )
+    def test_functional_signature(self, kernel, input_type):
+        check_functional_kernel_signature_match(F.pad, kernel=kernel, input_type=input_type)
+
+    @pytest.mark.parametrize(
+        "make_input",
+        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
+    )
+    def test_transform(self, make_input):
+        check_transform(transforms.Pad(padding=[1]), make_input())
+
+    def test_transform_errors(self):
+        with pytest.raises(TypeError, match="Got inappropriate padding arg"):
+            transforms.Pad("abc")
+
+        with pytest.raises(ValueError, match="Padding must be an int or a 1, 2, or 4"):
+            transforms.Pad([-0.7, 0, 0.7])
+
+        with pytest.raises(TypeError, match="Got inappropriate fill arg"):
+            transforms.Pad(12, fill="abc")
+
+        with pytest.raises(ValueError, match="Padding mode should be either"):
+            transforms.Pad(12, padding_mode="abc")
+
+    @pytest.mark.parametrize("padding", CORRECTNESS_PADDINGS)
+    @pytest.mark.parametrize(
+        ("padding_mode", "fill"),
+        [
+            *[("constant", fill) for fill in CORRECTNESS_FILLS],
+            *[(padding_mode, None) for padding_mode in ["symmetric", "edge", "reflect"]],
+        ],
+    )
+    @pytest.mark.parametrize("fn", [F.pad, transform_cls_to_functional(transforms.Pad)])
+    def test_image_correctness(self, padding, padding_mode, fill, fn):
+        image = make_image(dtype=torch.uint8, device="cpu")
+
+        actual = fn(image, padding=padding, padding_mode=padding_mode, fill=fill)
+        expected = F.to_image(F.pad(F.to_pil_image(image), padding=padding, padding_mode=padding_mode, fill=fill))
+
+        assert_equal(actual, expected)
+
+    def _reference_pad_bounding_boxes(self, bounding_boxes, *, padding):
+        if isinstance(padding, int):
+            padding = [padding]
+        left, top, right, bottom = padding * (4 // len(padding))
+
+        affine_matrix = np.array(
+            [
+                [1, 0, left],
+                [0, 1, top],
+            ],
+        )
+
+        height = bounding_boxes.canvas_size[0] + top + bottom
+        width = bounding_boxes.canvas_size[1] + left + right
+
+        return reference_affine_bounding_boxes_helper(
+            bounding_boxes, affine_matrix=affine_matrix, new_canvas_size=(height, width)
+        )
+
+    @pytest.mark.parametrize("padding", CORRECTNESS_PADDINGS)
+    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
+    @pytest.mark.parametrize("dtype", [torch.int64, torch.float32])
+    @pytest.mark.parametrize("device", cpu_and_cuda())
+    @pytest.mark.parametrize("fn", [F.pad, transform_cls_to_functional(transforms.Pad)])
+    def test_bounding_boxes_correctness(self, padding, format, dtype, device, fn):
+        bounding_boxes = make_bounding_boxes(format=format, dtype=dtype, device=device)
+
+        actual = fn(bounding_boxes, padding=padding)
+        expected = self._reference_pad_bounding_boxes(bounding_boxes, padding=padding)
+
+        assert_equal(actual, expected)
diff --git a/test/transforms_v2_dispatcher_infos.py b/test/transforms_v2_dispatcher_infos.py
@@ -1,7 +1,7 @@
 import pytest
 import torchvision.transforms.v2.functional as F
 from torchvision import tv_tensors
-from transforms_v2_kernel_infos import KERNEL_INFOS, pad_xfail_jit_fill_condition
+from transforms_v2_kernel_infos import KERNEL_INFOS
 from transforms_v2_legacy_utils import InfoBase, TestMark
 
 __all__ = ["DispatcherInfo", "DISPATCHER_INFOS"]
@@ -111,20 +111,6 @@ def xfail_jit_python_scalar_arg(name, *, reason=None):
 
 
 DISPATCHER_INFOS = [
-    DispatcherInfo(
-        F.pad,
-        kernels={
-            tv_tensors.Image: F.pad_image,
-            tv_tensors.Video: F.pad_video,
-            tv_tensors.BoundingBoxes: F.pad_bounding_boxes,
-            tv_tensors.Mask: F.pad_mask,
-        },
-        pil_kernel_info=PILKernelInfo(F._pad_image_pil, kernel_name="pad_image_pil"),
-        test_marks=[
-            xfail_jit("F.pad only supports vector fills for list of floats", condition=pad_xfail_jit_fill_condition),
-            xfail_jit_python_scalar_arg("padding"),
-        ],
-    ),
     DispatcherInfo(
         F.perspective,
         kernels={
diff --git a/test/transforms_v2_kernel_infos.py b/test/transforms_v2_kernel_infos.py
