diff --git a/test/test_transforms_v2.py b/test/test_transforms_v2.py index a47cf3bd408..093c378aa72 100644 --- a/test/test_transforms_v2.py +++ b/test/test_transforms_v2.py @@ -541,80 +541,6 @@ def test__transform_image_mask(self, fill, mocker): fn.assert_has_calls(calls) -class TestRandomRotation: - def test_assertions(self): - with pytest.raises(ValueError, match="is a single number, it must be positive"): - transforms.RandomRotation(-0.7) - - for d in [[-0.7], [-0.7, 0, 0.7]]: - with pytest.raises(ValueError, match="degrees should be a sequence of length 2"): - transforms.RandomRotation(d) - - with pytest.raises(TypeError, match="Got inappropriate fill arg"): - transforms.RandomRotation(12, fill="abc") - - with pytest.raises(TypeError, match="center should be a sequence of length"): - transforms.RandomRotation(12, center=12) - - with pytest.raises(ValueError, match="center should be a sequence of length"): - transforms.RandomRotation(12, center=[1, 2, 3]) - - def test__get_params(self): - angle_bound = 34 - transform = transforms.RandomRotation(angle_bound) - - params = transform._get_params(None) - assert -angle_bound <= params["angle"] <= angle_bound - - angle_bounds = [12, 34] - transform = transforms.RandomRotation(angle_bounds) - - params = transform._get_params(None) - assert angle_bounds[0] <= params["angle"] <= angle_bounds[1] - - @pytest.mark.parametrize("degrees", [23, [0, 45], (0, 45)]) - @pytest.mark.parametrize("expand", [False, True]) - @pytest.mark.parametrize("fill", [0, [1, 2, 3], (2, 3, 4)]) - @pytest.mark.parametrize("center", [None, [2.0, 3.0]]) - def test__transform(self, degrees, expand, fill, center, mocker): - interpolation = InterpolationMode.BILINEAR - transform = transforms.RandomRotation( - degrees, interpolation=interpolation, expand=expand, fill=fill, center=center - ) - - if isinstance(degrees, (tuple, list)): - assert transform.degrees == [float(degrees[0]), float(degrees[1])] - else: - assert transform.degrees == [float(-degrees), float(degrees)] - - fn = mocker.patch("torchvision.transforms.v2.functional.rotate") - inpt = mocker.MagicMock(spec=datapoints.Image) - # vfdev-5, Feature Request: let's store params as Transform attribute - # This could be also helpful for users - # Otherwise, we can mock transform._get_params - torch.manual_seed(12) - _ = transform(inpt) - torch.manual_seed(12) - params = transform._get_params(inpt) - - fill = transforms._utils._convert_fill_arg(fill) - fn.assert_called_once_with(inpt, **params, interpolation=interpolation, expand=expand, fill=fill, center=center) - - @pytest.mark.parametrize("angle", [34, -87]) - @pytest.mark.parametrize("expand", [False, True]) - def test_boundingbox_spatial_size(self, angle, expand): - # Specific test for BoundingBox.rotate - bbox = datapoints.BoundingBox( - torch.tensor([1, 2, 3, 4]), format=datapoints.BoundingBoxFormat.XYXY, spatial_size=(32, 32) - ) - img = datapoints.Image(torch.rand(1, 3, 32, 32)) - - out_img = img.rotate(angle, expand=expand) - out_bbox = bbox.rotate(angle, expand=expand) - - assert out_img.spatial_size == out_bbox.spatial_size - - class TestRandomCrop: def test_assertions(self): with pytest.raises(ValueError, match="Please provide only two dimensions"): diff --git a/test/test_transforms_v2_refactored.py b/test/test_transforms_v2_refactored.py index b737a7f0102..2130a8cf50a 100644 --- a/test/test_transforms_v2_refactored.py +++ b/test/test_transforms_v2_refactored.py @@ -308,6 +308,97 @@ def wrapper(input, *args, **kwargs): return wrapper +def make_input(input_type, *, dtype=None, device="cpu", spatial_size=(17, 11), mask_type="segmentation", **kwargs): + if input_type in {torch.Tensor, PIL.Image.Image, datapoints.Image}: + input = make_image(size=spatial_size, dtype=dtype or torch.uint8, device=device, **kwargs) + if input_type is torch.Tensor: + input = input.as_subclass(torch.Tensor) + elif input_type is PIL.Image.Image: + input = F.to_image_pil(input) + elif input_type is datapoints.BoundingBox: + kwargs.setdefault("format", datapoints.BoundingBoxFormat.XYXY) + input = make_bounding_box( + dtype=dtype or torch.float32, + device=device, + spatial_size=spatial_size, + **kwargs, + ) + elif input_type is datapoints.Mask: + if mask_type == "segmentation": + make_mask = make_segmentation_mask + default_dtype = torch.uint8 + elif mask_type == "detection": + make_mask = make_detection_mask + default_dtype = torch.bool + else: + raise ValueError(f"`mask_type` can be `'segmentation'` or `'detection'`, but got {mask_type}.") + input = make_mask(size=spatial_size, dtype=dtype or default_dtype, device=device, **kwargs) + elif input_type is datapoints.Video: + input = make_video(size=spatial_size, dtype=dtype or torch.uint8, device=device, **kwargs) + else: + raise TypeError( + f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, " + f"but got {input_type} instead." + ) + + return input + + +def param_value_parametrization(**kwargs): + """Helper function to turn + + @pytest.mark.parametrize( + ("param", "value"), + ("a", 1), + ("a", 2), + ("a", 3), + ("b", -1.0) + ("b", 1.0) + ) + + into + + @param_value_parametrization(a=[1, 2, 3], b=[-1.0, 1.0]) + """ + return pytest.mark.parametrize( + ("param", "value"), + [(param, value) for param, values in kwargs.items() for value in values], + ) + + +def adapt_fill(value, *, dtype): + """Adapt fill values in the range [0.0, 1.0] to the value range of the dtype""" + if value is None: + return value + + max_value = get_max_value(dtype) + + if isinstance(value, (int, float)): + return type(value)(value * max_value) + elif isinstance(value, (list, tuple)): + return type(value)(type(v)(v * max_value) for v in value) + else: + raise ValueError(f"fill should be an int or float, or a list or tuple of the former, but got '{value}'.") + + +EXHAUSTIVE_TYPE_FILLS = [ + None, + 1, + 0.5, + [1], + [0.2], + (0,), + (0.7,), + [1, 0, 1], + [0.1, 0.2, 0.3], + (0, 1, 0), + (0.9, 0.234, 0.314), +] +CORRECTNESS_FILLS = [ + v for v in EXHAUSTIVE_TYPE_FILLS if v is None or isinstance(v, float) or (isinstance(v, list) and len(v) > 1) +] + + # We cannot use `list(transforms.InterpolationMode)` here, since it includes some PIL-only ones as well INTERPOLATION_MODES = [ transforms.InterpolationMode.NEAREST, @@ -380,28 +471,6 @@ def _make_max_size_kwarg(self, *, use_max_size, size): return dict(max_size=max_size) - def _make_input(self, input_type, *, dtype=None, device="cpu", **kwargs): - if input_type in {torch.Tensor, PIL.Image.Image, datapoints.Image}: - input = make_image(size=self.INPUT_SIZE, dtype=dtype or torch.uint8, device=device, **kwargs) - if input_type is torch.Tensor: - input = input.as_subclass(torch.Tensor) - elif input_type is PIL.Image.Image: - input = F.to_image_pil(input) - elif input_type is datapoints.BoundingBox: - kwargs.setdefault("format", datapoints.BoundingBoxFormat.XYXY) - input = make_bounding_box( - spatial_size=self.INPUT_SIZE, - dtype=dtype or torch.float32, - device=device, - **kwargs, - ) - elif input_type is datapoints.Mask: - input = make_segmentation_mask(size=self.INPUT_SIZE, dtype=dtype or torch.uint8, device=device, **kwargs) - elif input_type is datapoints.Video: - input = make_video(size=self.INPUT_SIZE, dtype=dtype or torch.uint8, device=device, **kwargs) - - return input - def _compute_output_size(self, *, input_size, size, max_size): if not (isinstance(size, int) or len(size) == 1): return tuple(size) @@ -447,7 +516,7 @@ def test_kernel_image_tensor(self, size, interpolation, use_max_size, antialias, check_kernel( F.resize_image_tensor, - self._make_input(datapoints.Image, dtype=dtype, device=device), + make_input(datapoints.Image, dtype=dtype, device=device, spatial_size=self.INPUT_SIZE), size=size, interpolation=interpolation, **max_size_kwarg, @@ -465,7 +534,9 @@ def test_kernel_bounding_box(self, format, size, use_max_size, dtype, device): if not (max_size_kwarg := self._make_max_size_kwarg(use_max_size=use_max_size, size=size)): return - bounding_box = self._make_input(datapoints.BoundingBox, dtype=dtype, device=device, format=format) + bounding_box = make_input( + datapoints.BoundingBox, dtype=dtype, device=device, format=format, spatial_size=self.INPUT_SIZE + ) check_kernel( F.resize_bounding_box, bounding_box, @@ -475,14 +546,21 @@ def test_kernel_bounding_box(self, format, size, use_max_size, dtype, device): check_scripted_vs_eager=not isinstance(size, int), ) - @pytest.mark.parametrize( - ("dtype", "make_mask"), [(torch.uint8, make_segmentation_mask), (torch.bool, make_detection_mask)] - ) - def test_kernel_mask(self, dtype, make_mask): - check_kernel(F.resize_mask, make_mask(dtype=dtype), size=self.OUTPUT_SIZES[-1]) + @pytest.mark.parametrize("mask_type", ["segmentation", "detection"]) + def test_kernel_mask(self, mask_type): + check_kernel( + F.resize_mask, + make_input(datapoints.Mask, spatial_size=self.INPUT_SIZE, mask_type=mask_type), + size=self.OUTPUT_SIZES[-1], + ) def test_kernel_video(self): - check_kernel(F.resize_video, self._make_input(datapoints.Video), size=self.OUTPUT_SIZES[-1], antialias=True) + check_kernel( + F.resize_video, + make_input(datapoints.Video, spatial_size=self.INPUT_SIZE), + size=self.OUTPUT_SIZES[-1], + antialias=True, + ) @pytest.mark.parametrize("size", OUTPUT_SIZES) @pytest.mark.parametrize( @@ -500,7 +578,7 @@ def test_dispatcher(self, size, input_type, kernel): check_dispatcher( F.resize, kernel, - self._make_input(input_type), + make_input(input_type, spatial_size=self.INPUT_SIZE), size=size, antialias=True, check_scripted_smoke=not isinstance(size, int), @@ -527,7 +605,7 @@ def test_dispatcher_signature(self, kernel, input_type): [torch.Tensor, PIL.Image.Image, datapoints.Image, datapoints.BoundingBox, datapoints.Mask, datapoints.Video], ) def test_transform(self, size, device, input_type): - input = self._make_input(input_type, device=device) + input = make_input(input_type, device=device, spatial_size=self.INPUT_SIZE) check_transform( transforms.Resize, @@ -551,7 +629,7 @@ def test_image_correctness(self, size, interpolation, use_max_size, fn): if not (max_size_kwarg := self._make_max_size_kwarg(use_max_size=use_max_size, size=size)): return - image = self._make_input(torch.Tensor, dtype=torch.uint8, device="cpu") + image = make_input(torch.Tensor, dtype=torch.uint8, device="cpu", spatial_size=self.INPUT_SIZE) actual = fn(image, size=size, interpolation=interpolation, **max_size_kwarg, antialias=True) expected = F.to_image_tensor( @@ -594,7 +672,7 @@ def test_bounding_box_correctness(self, format, size, use_max_size, fn): if not (max_size_kwarg := self._make_max_size_kwarg(use_max_size=use_max_size, size=size)): return - bounding_box = self._make_input(datapoints.BoundingBox) + bounding_box = make_input(datapoints.BoundingBox, spatial_size=self.INPUT_SIZE) actual = fn(bounding_box, size=size, **max_size_kwarg) expected = self._reference_resize_bounding_box(bounding_box, size=size, **max_size_kwarg) @@ -608,7 +686,7 @@ def test_bounding_box_correctness(self, format, size, use_max_size, fn): [torch.Tensor, PIL.Image.Image, datapoints.Image, datapoints.Video], ) def test_pil_interpolation_compat_smoke(self, interpolation, input_type): - input = self._make_input(input_type) + input = make_input(input_type, spatial_size=self.INPUT_SIZE) with ( contextlib.nullcontext() @@ -624,7 +702,9 @@ def test_pil_interpolation_compat_smoke(self, interpolation, input_type): def test_dispatcher_pil_antialias_warning(self): with pytest.warns(UserWarning, match="Anti-alias option is always applied for PIL Image input"): - F.resize(self._make_input(PIL.Image.Image), size=self.OUTPUT_SIZES[0], antialias=False) + F.resize( + make_input(PIL.Image.Image, spatial_size=self.INPUT_SIZE), size=self.OUTPUT_SIZES[0], antialias=False + ) @pytest.mark.parametrize("size", OUTPUT_SIZES) @pytest.mark.parametrize( @@ -641,7 +721,7 @@ def test_max_size_error(self, size, input_type): match = "size should be an int or a sequence of length 1" with pytest.raises(ValueError, match=match): - F.resize(self._make_input(input_type), size=size, max_size=max_size, antialias=True) + F.resize(make_input(input_type, spatial_size=self.INPUT_SIZE), size=size, max_size=max_size, antialias=True) @pytest.mark.parametrize("interpolation", INTERPOLATION_MODES) @pytest.mark.parametrize( @@ -654,7 +734,11 @@ def test_antialias_warning(self, interpolation, input_type): if interpolation in {transforms.InterpolationMode.BILINEAR, transforms.InterpolationMode.BICUBIC} else assert_no_warnings() ): - F.resize(self._make_input(input_type), size=self.OUTPUT_SIZES[0], interpolation=interpolation) + F.resize( + make_input(input_type, spatial_size=self.INPUT_SIZE), + size=self.OUTPUT_SIZES[0], + interpolation=interpolation, + ) @pytest.mark.parametrize("interpolation", INTERPOLATION_MODES) @pytest.mark.parametrize( @@ -668,7 +752,7 @@ def test_interpolation_int(self, interpolation, input_type): if issubclass(input_type, torch.Tensor) and interpolation is transforms.InterpolationMode.NEAREST_EXACT: return - input = self._make_input(input_type) + input = make_input(input_type, spatial_size=self.INPUT_SIZE) expected = F.resize(input, size=self.OUTPUT_SIZES[0], interpolation=interpolation, antialias=True) actual = F.resize( @@ -689,7 +773,7 @@ def test_transform_unknown_size_error(self): [torch.Tensor, PIL.Image.Image, datapoints.Image, datapoints.BoundingBox, datapoints.Mask, datapoints.Video], ) def test_noop(self, size, input_type): - input = self._make_input(input_type) + input = make_input(input_type, spatial_size=self.INPUT_SIZE) output = F.resize(input, size=size, antialias=True) @@ -711,7 +795,7 @@ def test_no_regression_5405(self, input_type): # Checks that `max_size` is not ignored if `size == small_edge_size` # See https://github.com/pytorch/vision/issues/5405 - input = self._make_input(input_type) + input = make_input(input_type, spatial_size=self.INPUT_SIZE) size = min(F.get_spatial_size(input)) max_size = size + 1 @@ -721,38 +805,16 @@ def test_no_regression_5405(self, input_type): class TestHorizontalFlip: - def _make_input(self, input_type, *, dtype=None, device="cpu", spatial_size=(17, 11), **kwargs): - if input_type in {torch.Tensor, PIL.Image.Image, datapoints.Image}: - input = make_image(size=spatial_size, dtype=dtype or torch.uint8, device=device, **kwargs) - if input_type is torch.Tensor: - input = input.as_subclass(torch.Tensor) - elif input_type is PIL.Image.Image: - input = F.to_image_pil(input) - elif input_type is datapoints.BoundingBox: - kwargs.setdefault("format", datapoints.BoundingBoxFormat.XYXY) - input = make_bounding_box( - dtype=dtype or torch.float32, - device=device, - spatial_size=spatial_size, - **kwargs, - ) - elif input_type is datapoints.Mask: - input = make_segmentation_mask(size=spatial_size, dtype=dtype or torch.uint8, device=device, **kwargs) - elif input_type is datapoints.Video: - input = make_video(size=spatial_size, dtype=dtype or torch.uint8, device=device, **kwargs) - - return input - @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8]) @pytest.mark.parametrize("device", cpu_and_cuda()) def test_kernel_image_tensor(self, dtype, device): - check_kernel(F.horizontal_flip_image_tensor, self._make_input(torch.Tensor, dtype=dtype, device=device)) + check_kernel(F.horizontal_flip_image_tensor, make_input(torch.Tensor, dtype=dtype, device=device)) @pytest.mark.parametrize("format", list(datapoints.BoundingBoxFormat)) @pytest.mark.parametrize("dtype", [torch.float32, torch.int64]) @pytest.mark.parametrize("device", cpu_and_cuda()) def test_kernel_bounding_box(self, format, dtype, device): - bounding_box = self._make_input(datapoints.BoundingBox, dtype=dtype, device=device, format=format) + bounding_box = make_input(datapoints.BoundingBox, dtype=dtype, device=device, format=format) check_kernel( F.horizontal_flip_bounding_box, bounding_box, @@ -760,15 +822,12 @@ def test_kernel_bounding_box(self, format, dtype, device): spatial_size=bounding_box.spatial_size, ) - @pytest.mark.parametrize( - "dtype_and_make_mask", [(torch.uint8, make_segmentation_mask), (torch.bool, make_detection_mask)] - ) - def test_kernel_mask(self, dtype_and_make_mask): - dtype, make_mask = dtype_and_make_mask - check_kernel(F.horizontal_flip_mask, make_mask(dtype=dtype)) + @pytest.mark.parametrize("mask_type", ["segmentation", "detection"]) + def test_kernel_mask(self, mask_type): + check_kernel(F.horizontal_flip_mask, make_input(datapoints.Mask, mask_type=mask_type)) def test_kernel_video(self): - check_kernel(F.horizontal_flip_video, self._make_input(datapoints.Video)) + check_kernel(F.horizontal_flip_video, make_input(datapoints.Video)) @pytest.mark.parametrize( ("input_type", "kernel"), @@ -782,7 +841,7 @@ def test_kernel_video(self): ], ) def test_dispatcher(self, kernel, input_type): - check_dispatcher(F.horizontal_flip, kernel, self._make_input(input_type)) + check_dispatcher(F.horizontal_flip, kernel, make_input(input_type)) @pytest.mark.parametrize( ("input_type", "kernel"), @@ -804,7 +863,7 @@ def test_dispatcher_signature(self, kernel, input_type): ) @pytest.mark.parametrize("device", cpu_and_cuda()) def test_transform(self, input_type, device): - input = self._make_input(input_type, device=device) + input = make_input(input_type, device=device) check_transform(transforms.RandomHorizontalFlip, input, p=1) @@ -812,7 +871,7 @@ def test_transform(self, input_type, device): "fn", [F.horizontal_flip, transform_cls_to_functional(transforms.RandomHorizontalFlip, p=1)] ) def test_image_correctness(self, fn): - image = self._make_input(torch.Tensor, dtype=torch.uint8, device="cpu") + image = make_input(torch.Tensor, dtype=torch.uint8, device="cpu") actual = fn(image) expected = F.to_image_tensor(F.horizontal_flip(F.to_image_pil(image))) @@ -842,7 +901,7 @@ def _reference_horizontal_flip_bounding_box(self, bounding_box): "fn", [F.horizontal_flip, transform_cls_to_functional(transforms.RandomHorizontalFlip, p=1)] ) def test_bounding_box_correctness(self, format, fn): - bounding_box = self._make_input(datapoints.BoundingBox, format=format) + bounding_box = make_input(datapoints.BoundingBox, format=format) actual = fn(bounding_box) expected = self._reference_horizontal_flip_bounding_box(bounding_box) @@ -855,7 +914,7 @@ def test_bounding_box_correctness(self, format, fn): ) @pytest.mark.parametrize("device", cpu_and_cuda()) def test_transform_noop(self, input_type, device): - input = self._make_input(input_type, device=device) + input = make_input(input_type, device=device) transform = transforms.RandomHorizontalFlip(p=0) @@ -865,50 +924,6 @@ def test_transform_noop(self, input_type, device): class TestAffine: - def _make_input( - self, input_type, *, dtype=None, device="cpu", spatial_size=(17, 11), mask_type="segmentation", **kwargs - ): - if input_type in {torch.Tensor, PIL.Image.Image, datapoints.Image}: - input = make_image(size=spatial_size, dtype=dtype or torch.uint8, device=device, **kwargs) - if input_type is torch.Tensor: - input = input.as_subclass(torch.Tensor) - elif input_type is PIL.Image.Image: - input = F.to_image_pil(input) - elif input_type is datapoints.BoundingBox: - kwargs.setdefault("format", datapoints.BoundingBoxFormat.XYXY) - input = make_bounding_box( - dtype=dtype or torch.float32, - device=device, - spatial_size=spatial_size, - **kwargs, - ) - elif input_type is datapoints.Mask: - if mask_type == "segmentation": - make_mask = make_segmentation_mask - default_dtype = torch.uint8 - elif mask_type == "detection": - make_mask = make_detection_mask - default_dtype = torch.bool - input = make_mask(size=spatial_size, dtype=dtype or default_dtype, device=device, **kwargs) - elif input_type is datapoints.Video: - input = make_video(size=spatial_size, dtype=dtype or torch.uint8, device=device, **kwargs) - - return input - - def _adapt_fill(self, value, *, dtype): - """Adapt fill values in the range [0.0, 1.0] to the value range of the dtype""" - if value is None: - return value - - max_value = get_max_value(dtype) - - if isinstance(value, (int, float)): - return type(value)(value * max_value) - elif isinstance(value, (list, tuple)): - return type(value)(type(v)(v * max_value) for v in value) - else: - raise ValueError(f"fill should be an int or float, or a list or tuple of the former, but got '{value}'") - _EXHAUSTIVE_TYPE_AFFINE_KWARGS = dict( # float, int angle=[-10.9, 18], @@ -934,23 +949,6 @@ def _adapt_fill(self, value, *, dtype): for k, vs in _EXHAUSTIVE_TYPE_AFFINE_KWARGS.items() } - _EXHAUSTIVE_TYPE_FILLS = [ - None, - 1, - 0.5, - [1], - [0.2], - (0,), - (0.7,), - [1, 0, 1], - [0.1, 0.2, 0.3], - (0, 1, 0), - (0.9, 0.234, 0.314), - ] - _CORRECTNESS_FILL = [ - v for v in _EXHAUSTIVE_TYPE_FILLS if v is None or isinstance(v, float) or (isinstance(v, list) and len(v) > 1) - ] - _EXHAUSTIVE_TYPE_TRANSFORM_AFFINE_RANGES = dict( degrees=[30, (-15, 20)], translate=[None, (0.5, 0.5)], @@ -966,29 +964,22 @@ def _check_kernel(self, kernel, input, *args, **kwargs): kwargs_.update(kwargs) check_kernel(kernel, input, *args, **kwargs_) - @pytest.mark.parametrize( - ("param", "value"), - [ - (param, value) - for param, values in [ - ("angle", _EXHAUSTIVE_TYPE_AFFINE_KWARGS["angle"]), - ("translate", _EXHAUSTIVE_TYPE_AFFINE_KWARGS["translate"]), - ("shear", _EXHAUSTIVE_TYPE_AFFINE_KWARGS["shear"]), - ("center", _EXHAUSTIVE_TYPE_AFFINE_KWARGS["center"]), - ("interpolation", [transforms.InterpolationMode.NEAREST, transforms.InterpolationMode.BILINEAR]), - ("fill", _EXHAUSTIVE_TYPE_FILLS), - ] - for value in values - ], + @param_value_parametrization( + angle=_EXHAUSTIVE_TYPE_AFFINE_KWARGS["angle"], + translate=_EXHAUSTIVE_TYPE_AFFINE_KWARGS["translate"], + shear=_EXHAUSTIVE_TYPE_AFFINE_KWARGS["shear"], + center=_EXHAUSTIVE_TYPE_AFFINE_KWARGS["center"], + interpolation=[transforms.InterpolationMode.NEAREST, transforms.InterpolationMode.BILINEAR], + fill=EXHAUSTIVE_TYPE_FILLS, ) @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8]) @pytest.mark.parametrize("device", cpu_and_cuda()) def test_kernel_image_tensor(self, param, value, dtype, device): if param == "fill": - value = self._adapt_fill(value, dtype=dtype) + value = adapt_fill(value, dtype=dtype) self._check_kernel( F.affine_image_tensor, - self._make_input(torch.Tensor, dtype=dtype, device=device), + make_input(torch.Tensor, dtype=dtype, device=device), **{param: value}, check_scripted_vs_eager=not (param in {"shear", "fill"} and isinstance(value, (int, float))), check_cuda_vs_cpu=dict(atol=1, rtol=0) @@ -996,27 +987,20 @@ def test_kernel_image_tensor(self, param, value, dtype, device): else True, ) - @pytest.mark.parametrize( - ("param", "value"), - [ - (param, value) - for param, values in [ - ("angle", _EXHAUSTIVE_TYPE_AFFINE_KWARGS["angle"]), - ("translate", _EXHAUSTIVE_TYPE_AFFINE_KWARGS["translate"]), - ("shear", _EXHAUSTIVE_TYPE_AFFINE_KWARGS["shear"]), - ("center", _EXHAUSTIVE_TYPE_AFFINE_KWARGS["center"]), - ] - for value in values - ], + @param_value_parametrization( + angle=_EXHAUSTIVE_TYPE_AFFINE_KWARGS["angle"], + translate=_EXHAUSTIVE_TYPE_AFFINE_KWARGS["translate"], + shear=_EXHAUSTIVE_TYPE_AFFINE_KWARGS["shear"], + center=_EXHAUSTIVE_TYPE_AFFINE_KWARGS["center"], ) @pytest.mark.parametrize("format", list(datapoints.BoundingBoxFormat)) @pytest.mark.parametrize("dtype", [torch.float32, torch.int64]) @pytest.mark.parametrize("device", cpu_and_cuda()) def test_kernel_bounding_box(self, param, value, format, dtype, device): - bounding_box = self._make_input(datapoints.BoundingBox, format=format, dtype=dtype, device=device) + bounding_box = make_input(datapoints.BoundingBox, format=format, dtype=dtype, device=device) self._check_kernel( F.affine_bounding_box, - self._make_input(datapoints.BoundingBox, format=format, dtype=dtype, device=device), + make_input(datapoints.BoundingBox, format=format, dtype=dtype, device=device), format=format, spatial_size=bounding_box.spatial_size, **{param: value}, @@ -1025,10 +1009,10 @@ def test_kernel_bounding_box(self, param, value, format, dtype, device): @pytest.mark.parametrize("mask_type", ["segmentation", "detection"]) def test_kernel_mask(self, mask_type): - self._check_kernel(F.affine_mask, self._make_input(datapoints.Mask, mask_type=mask_type)) + self._check_kernel(F.affine_mask, make_input(datapoints.Mask, mask_type=mask_type)) def test_kernel_video(self): - self._check_kernel(F.affine_video, self._make_input(datapoints.Video)) + self._check_kernel(F.affine_video, make_input(datapoints.Video)) @pytest.mark.parametrize( ("input_type", "kernel"), @@ -1042,7 +1026,7 @@ def test_kernel_video(self): ], ) def test_dispatcher(self, kernel, input_type): - check_dispatcher(F.affine, kernel, self._make_input(input_type), **self._MINIMAL_AFFINE_KWARGS) + check_dispatcher(F.affine, kernel, make_input(input_type), **self._MINIMAL_AFFINE_KWARGS) @pytest.mark.parametrize( ("input_type", "kernel"), @@ -1064,7 +1048,7 @@ def test_dispatcher_signature(self, kernel, input_type): ) @pytest.mark.parametrize("device", cpu_and_cuda()) def test_transform(self, input_type, device): - input = self._make_input(input_type, device=device) + input = make_input(input_type, device=device) check_transform(transforms.RandomAffine, input, **self._CORRECTNESS_TRANSFORM_AFFINE_RANGES) @@ -1076,11 +1060,11 @@ def test_transform(self, input_type, device): @pytest.mark.parametrize( "interpolation", [transforms.InterpolationMode.NEAREST, transforms.InterpolationMode.BILINEAR] ) - @pytest.mark.parametrize("fill", _CORRECTNESS_FILL) + @pytest.mark.parametrize("fill", CORRECTNESS_FILLS) def test_functional_image_correctness(self, angle, translate, scale, shear, center, interpolation, fill): - image = self._make_input(torch.Tensor, dtype=torch.uint8, device="cpu") + image = make_input(torch.Tensor, dtype=torch.uint8, device="cpu") - fill = self._adapt_fill(fill, dtype=torch.uint8) + fill = adapt_fill(fill, dtype=torch.uint8) actual = F.affine( image, @@ -1112,12 +1096,12 @@ def test_functional_image_correctness(self, angle, translate, scale, shear, cent @pytest.mark.parametrize( "interpolation", [transforms.InterpolationMode.NEAREST, transforms.InterpolationMode.BILINEAR] ) - @pytest.mark.parametrize("fill", _CORRECTNESS_FILL) + @pytest.mark.parametrize("fill", CORRECTNESS_FILLS) @pytest.mark.parametrize("seed", list(range(5))) def test_transform_image_correctness(self, center, interpolation, fill, seed): - image = self._make_input(torch.Tensor, dtype=torch.uint8, device="cpu") + image = make_input(torch.Tensor, dtype=torch.uint8, device="cpu") - fill = self._adapt_fill(fill, dtype=torch.uint8) + fill = adapt_fill(fill, dtype=torch.uint8) transform = transforms.RandomAffine( **self._CORRECTNESS_TRANSFORM_AFFINE_RANGES, center=center, interpolation=interpolation, fill=fill @@ -1179,7 +1163,7 @@ def _reference_affine_bounding_box(self, bounding_box, *, angle, translate, scal @pytest.mark.parametrize("shear", _CORRECTNESS_AFFINE_KWARGS["shear"]) @pytest.mark.parametrize("center", _CORRECTNESS_AFFINE_KWARGS["center"]) def test_functional_bounding_box_correctness(self, format, angle, translate, scale, shear, center): - bounding_box = self._make_input(datapoints.BoundingBox, format=format) + bounding_box = make_input(datapoints.BoundingBox, format=format) actual = F.affine( bounding_box, @@ -1204,7 +1188,7 @@ def test_functional_bounding_box_correctness(self, format, angle, translate, sca @pytest.mark.parametrize("center", _CORRECTNESS_AFFINE_KWARGS["center"]) @pytest.mark.parametrize("seed", list(range(5))) def test_transform_bounding_box_correctness(self, format, center, seed): - bounding_box = self._make_input(datapoints.BoundingBox, format=format) + bounding_box = make_input(datapoints.BoundingBox, format=format) transform = transforms.RandomAffine(**self._CORRECTNESS_TRANSFORM_AFFINE_RANGES, center=center) @@ -1224,7 +1208,7 @@ def test_transform_bounding_box_correctness(self, format, center, seed): @pytest.mark.parametrize("shear", _EXHAUSTIVE_TYPE_TRANSFORM_AFFINE_RANGES["shear"]) @pytest.mark.parametrize("seed", list(range(10))) def test_transform_get_params_bounds(self, degrees, translate, scale, shear, seed): - image = self._make_input(torch.Tensor) + image = make_input(torch.Tensor) height, width = F.get_spatial_size(image) transform = transforms.RandomAffine(degrees=degrees, translate=translate, scale=scale, shear=shear) @@ -1302,38 +1286,16 @@ def test_transform_unknown_fill_error(self): class TestVerticalFlip: - def _make_input(self, input_type, *, dtype=None, device="cpu", spatial_size=(17, 11), **kwargs): - if input_type in {torch.Tensor, PIL.Image.Image, datapoints.Image}: - input = make_image(size=spatial_size, dtype=dtype or torch.uint8, device=device, **kwargs) - if input_type is torch.Tensor: - input = input.as_subclass(torch.Tensor) - elif input_type is PIL.Image.Image: - input = F.to_image_pil(input) - elif input_type is datapoints.BoundingBox: - kwargs.setdefault("format", datapoints.BoundingBoxFormat.XYXY) - input = make_bounding_box( - dtype=dtype or torch.float32, - device=device, - spatial_size=spatial_size, - **kwargs, - ) - elif input_type is datapoints.Mask: - input = make_segmentation_mask(size=spatial_size, dtype=dtype or torch.uint8, device=device, **kwargs) - elif input_type is datapoints.Video: - input = make_video(size=spatial_size, dtype=dtype or torch.uint8, device=device, **kwargs) - - return input - @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8]) @pytest.mark.parametrize("device", cpu_and_cuda()) def test_kernel_image_tensor(self, dtype, device): - check_kernel(F.vertical_flip_image_tensor, self._make_input(torch.Tensor, dtype=dtype, device=device)) + check_kernel(F.vertical_flip_image_tensor, make_input(torch.Tensor, dtype=dtype, device=device)) @pytest.mark.parametrize("format", list(datapoints.BoundingBoxFormat)) @pytest.mark.parametrize("dtype", [torch.float32, torch.int64]) @pytest.mark.parametrize("device", cpu_and_cuda()) def test_kernel_bounding_box(self, format, dtype, device): - bounding_box = self._make_input(datapoints.BoundingBox, dtype=dtype, device=device, format=format) + bounding_box = make_input(datapoints.BoundingBox, dtype=dtype, device=device, format=format) check_kernel( F.vertical_flip_bounding_box, bounding_box, @@ -1341,15 +1303,12 @@ def test_kernel_bounding_box(self, format, dtype, device): spatial_size=bounding_box.spatial_size, ) - @pytest.mark.parametrize( - "dtype_and_make_mask", [(torch.uint8, make_segmentation_mask), (torch.bool, make_detection_mask)] - ) - def test_kernel_mask(self, dtype_and_make_mask): - dtype, make_mask = dtype_and_make_mask - check_kernel(F.vertical_flip_mask, make_mask(dtype=dtype)) + @pytest.mark.parametrize("mask_type", ["segmentation", "detection"]) + def test_kernel_mask(self, mask_type): + check_kernel(F.vertical_flip_mask, make_input(datapoints.Mask, mask_type=mask_type)) def test_kernel_video(self): - check_kernel(F.vertical_flip_video, self._make_input(datapoints.Video)) + check_kernel(F.vertical_flip_video, make_input(datapoints.Video)) @pytest.mark.parametrize( ("input_type", "kernel"), @@ -1363,7 +1322,7 @@ def test_kernel_video(self): ], ) def test_dispatcher(self, kernel, input_type): - check_dispatcher(F.vertical_flip, kernel, self._make_input(input_type)) + check_dispatcher(F.vertical_flip, kernel, make_input(input_type)) @pytest.mark.parametrize( ("input_type", "kernel"), @@ -1385,13 +1344,13 @@ def test_dispatcher_signature(self, kernel, input_type): ) @pytest.mark.parametrize("device", cpu_and_cuda()) def test_transform(self, input_type, device): - input = self._make_input(input_type, device=device) + input = make_input(input_type, device=device) check_transform(transforms.RandomVerticalFlip, input, p=1) @pytest.mark.parametrize("fn", [F.vertical_flip, transform_cls_to_functional(transforms.RandomVerticalFlip, p=1)]) def test_image_correctness(self, fn): - image = self._make_input(torch.Tensor, dtype=torch.uint8, device="cpu") + image = make_input(torch.Tensor, dtype=torch.uint8, device="cpu") actual = fn(image) expected = F.to_image_tensor(F.vertical_flip(F.to_image_pil(image))) @@ -1419,7 +1378,7 @@ def _reference_vertical_flip_bounding_box(self, bounding_box): @pytest.mark.parametrize("format", list(datapoints.BoundingBoxFormat)) @pytest.mark.parametrize("fn", [F.vertical_flip, transform_cls_to_functional(transforms.RandomVerticalFlip, p=1)]) def test_bounding_box_correctness(self, format, fn): - bounding_box = self._make_input(datapoints.BoundingBox, format=format) + bounding_box = make_input(datapoints.BoundingBox, format=format) actual = fn(bounding_box) expected = self._reference_vertical_flip_bounding_box(bounding_box) @@ -1432,10 +1391,267 @@ def test_bounding_box_correctness(self, format, fn): ) @pytest.mark.parametrize("device", cpu_and_cuda()) def test_transform_noop(self, input_type, device): - input = self._make_input(input_type, device=device) + input = make_input(input_type, device=device) transform = transforms.RandomVerticalFlip(p=0) output = transform(input) assert_equal(output, input) + + +class TestRotate: + _EXHAUSTIVE_TYPE_AFFINE_KWARGS = dict( + # float, int + angle=[-10.9, 18], + # None + # two-list of float, two-list of int, two-tuple of float, two-tuple of int + center=[None, [1.2, 4.9], [-3, 1], (2.5, -4.7), (3, 2)], + ) + _MINIMAL_AFFINE_KWARGS = {k: vs[0] for k, vs in _EXHAUSTIVE_TYPE_AFFINE_KWARGS.items()} + _CORRECTNESS_AFFINE_KWARGS = { + k: [v for v in vs if v is None or isinstance(v, float) or isinstance(v, list)] + for k, vs in _EXHAUSTIVE_TYPE_AFFINE_KWARGS.items() + } + + _EXHAUSTIVE_TYPE_TRANSFORM_AFFINE_RANGES = dict( + degrees=[30, (-15, 20)], + ) + _CORRECTNESS_TRANSFORM_AFFINE_RANGES = {k: vs[0] for k, vs in _EXHAUSTIVE_TYPE_TRANSFORM_AFFINE_RANGES.items()} + + @param_value_parametrization( + angle=_EXHAUSTIVE_TYPE_AFFINE_KWARGS["angle"], + interpolation=[transforms.InterpolationMode.NEAREST, transforms.InterpolationMode.BILINEAR], + expand=[False, True], + center=_EXHAUSTIVE_TYPE_AFFINE_KWARGS["center"], + fill=EXHAUSTIVE_TYPE_FILLS, + ) + @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8]) + @pytest.mark.parametrize("device", cpu_and_cuda()) + def test_kernel_image_tensor(self, param, value, dtype, device): + kwargs = {param: value} + if param != "angle": + kwargs["angle"] = self._MINIMAL_AFFINE_KWARGS["angle"] + check_kernel( + F.rotate_image_tensor, + make_input(torch.Tensor, dtype=dtype, device=device), + **kwargs, + check_scripted_vs_eager=not (param == "fill" and isinstance(value, (int, float))), + ) + + @param_value_parametrization( + angle=_EXHAUSTIVE_TYPE_AFFINE_KWARGS["angle"], + expand=[False, True], + center=_EXHAUSTIVE_TYPE_AFFINE_KWARGS["center"], + ) + @pytest.mark.parametrize("format", list(datapoints.BoundingBoxFormat)) + @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8]) + @pytest.mark.parametrize("device", cpu_and_cuda()) + def test_kernel_bounding_box(self, param, value, format, dtype, device): + kwargs = {param: value} + if param != "angle": + kwargs["angle"] = self._MINIMAL_AFFINE_KWARGS["angle"] + + bounding_box = make_input(datapoints.BoundingBox, dtype=dtype, device=device, format=format) + + check_kernel( + F.rotate_bounding_box, + bounding_box, + format=format, + spatial_size=bounding_box.spatial_size, + **kwargs, + ) + + @pytest.mark.parametrize("mask_type", ["segmentation", "detection"]) + def test_kernel_mask(self, mask_type): + check_kernel(F.rotate_mask, make_input(datapoints.Mask, mask_type=mask_type), **self._MINIMAL_AFFINE_KWARGS) + + def test_kernel_video(self): + check_kernel(F.rotate_video, make_input(datapoints.Video), **self._MINIMAL_AFFINE_KWARGS) + + @pytest.mark.parametrize( + ("input_type", "kernel"), + [ + (torch.Tensor, F.rotate_image_tensor), + (PIL.Image.Image, F.rotate_image_pil), + (datapoints.Image, F.rotate_image_tensor), + (datapoints.BoundingBox, F.rotate_bounding_box), + (datapoints.Mask, F.rotate_mask), + (datapoints.Video, F.rotate_video), + ], + ) + def test_dispatcher(self, kernel, input_type): + check_dispatcher(F.rotate, kernel, make_input(input_type), **self._MINIMAL_AFFINE_KWARGS) + + @pytest.mark.parametrize( + ("input_type", "kernel"), + [ + (torch.Tensor, F.rotate_image_tensor), + (PIL.Image.Image, F.rotate_image_pil), + (datapoints.Image, F.rotate_image_tensor), + (datapoints.BoundingBox, F.rotate_bounding_box), + (datapoints.Mask, F.rotate_mask), + (datapoints.Video, F.rotate_video), + ], + ) + def test_dispatcher_signature(self, kernel, input_type): + check_dispatcher_signatures_match(F.rotate, kernel=kernel, input_type=input_type) + + @pytest.mark.parametrize( + "input_type", + [torch.Tensor, PIL.Image.Image, datapoints.Image, datapoints.BoundingBox, datapoints.Mask, datapoints.Video], + ) + @pytest.mark.parametrize("device", cpu_and_cuda()) + def test_transform(self, input_type, device): + input = make_input(input_type, device=device) + + check_transform(transforms.RandomRotation, input, **self._CORRECTNESS_TRANSFORM_AFFINE_RANGES) + + @pytest.mark.parametrize("angle", _CORRECTNESS_AFFINE_KWARGS["angle"]) + @pytest.mark.parametrize("center", _CORRECTNESS_AFFINE_KWARGS["center"]) + @pytest.mark.parametrize( + "interpolation", [transforms.InterpolationMode.NEAREST, transforms.InterpolationMode.BILINEAR] + ) + @pytest.mark.parametrize("expand", [False, True]) + @pytest.mark.parametrize("fill", CORRECTNESS_FILLS) + def test_functional_image_correctness(self, angle, center, interpolation, expand, fill): + image = make_input(torch.Tensor, dtype=torch.uint8, device="cpu") + + fill = adapt_fill(fill, dtype=torch.uint8) + + actual = F.rotate(image, angle=angle, center=center, interpolation=interpolation, expand=expand, fill=fill) + expected = F.to_image_tensor( + F.rotate( + F.to_image_pil(image), angle=angle, center=center, interpolation=interpolation, expand=expand, fill=fill + ) + ) + + mae = (actual.float() - expected.float()).abs().mean() + assert mae < 1 if interpolation is transforms.InterpolationMode.NEAREST else 6 + + @pytest.mark.parametrize("center", _CORRECTNESS_AFFINE_KWARGS["center"]) + @pytest.mark.parametrize( + "interpolation", [transforms.InterpolationMode.NEAREST, transforms.InterpolationMode.BILINEAR] + ) + @pytest.mark.parametrize("expand", [False, True]) + @pytest.mark.parametrize("fill", CORRECTNESS_FILLS) + @pytest.mark.parametrize("seed", list(range(5))) + def test_transform_image_correctness(self, center, interpolation, expand, fill, seed): + image = make_input(torch.Tensor, dtype=torch.uint8, device="cpu") + + fill = adapt_fill(fill, dtype=torch.uint8) + + transform = transforms.RandomRotation( + **self._CORRECTNESS_TRANSFORM_AFFINE_RANGES, + center=center, + interpolation=interpolation, + expand=expand, + fill=fill, + ) + + torch.manual_seed(seed) + actual = transform(image) + + torch.manual_seed(seed) + expected = F.to_image_tensor(transform(F.to_image_pil(image))) + + mae = (actual.float() - expected.float()).abs().mean() + assert mae < 1 if interpolation is transforms.InterpolationMode.NEAREST else 6 + + def _reference_rotate_bounding_box(self, bounding_box, *, angle, expand, center): + # FIXME + if expand: + raise ValueError("This reference currently does not support expand=True") + + if center is None: + center = [s * 0.5 for s in bounding_box.spatial_size[::-1]] + + a = np.cos(angle * np.pi / 180.0) + b = np.sin(angle * np.pi / 180.0) + cx = center[0] + cy = center[1] + affine_matrix = np.array( + [ + [a, b, cx - cx * a - b * cy], + [-b, a, cy + cx * b - a * cy], + ], + dtype="float64" if bounding_box.dtype == torch.float64 else "float32", + ) + + expected_bboxes = reference_affine_bounding_box_helper( + bounding_box, + format=bounding_box.format, + spatial_size=bounding_box.spatial_size, + affine_matrix=affine_matrix, + ) + + return expected_bboxes + + @pytest.mark.parametrize("format", list(datapoints.BoundingBoxFormat)) + @pytest.mark.parametrize("angle", _CORRECTNESS_AFFINE_KWARGS["angle"]) + # TODO: add support for expand=True in the reference + @pytest.mark.parametrize("expand", [False]) + @pytest.mark.parametrize("center", _CORRECTNESS_AFFINE_KWARGS["center"]) + def test_functional_bounding_box_correctness(self, format, angle, expand, center): + bounding_box = make_input(datapoints.BoundingBox, format=format) + + actual = F.rotate(bounding_box, angle=angle, expand=expand, center=center) + expected = self._reference_rotate_bounding_box(bounding_box, angle=angle, expand=expand, center=center) + + torch.testing.assert_close(actual, expected) + + @pytest.mark.parametrize("format", list(datapoints.BoundingBoxFormat)) + # TODO: add support for expand=True in the reference + @pytest.mark.parametrize("expand", [False]) + @pytest.mark.parametrize("center", _CORRECTNESS_AFFINE_KWARGS["center"]) + @pytest.mark.parametrize("seed", list(range(5))) + def test_transform_bounding_box_correctness(self, format, expand, center, seed): + bounding_box = make_input(datapoints.BoundingBox, format=format) + + transform = transforms.RandomRotation(**self._CORRECTNESS_TRANSFORM_AFFINE_RANGES, expand=expand, center=center) + + torch.manual_seed(seed) + params = transform._get_params([bounding_box]) + + torch.manual_seed(seed) + actual = transform(bounding_box) + + expected = self._reference_rotate_bounding_box(bounding_box, **params, expand=expand, center=center) + + torch.testing.assert_close(actual, expected) + + @pytest.mark.parametrize("degrees", _EXHAUSTIVE_TYPE_TRANSFORM_AFFINE_RANGES["degrees"]) + @pytest.mark.parametrize("seed", list(range(10))) + def test_transform_get_params_bounds(self, degrees, seed): + transform = transforms.RandomRotation(degrees=degrees) + + torch.manual_seed(seed) + params = transform._get_params([]) + + if isinstance(degrees, (int, float)): + assert -degrees <= params["angle"] <= degrees + else: + assert degrees[0] <= params["angle"] <= degrees[1] + + @pytest.mark.parametrize("param", ["degrees", "center"]) + @pytest.mark.parametrize("value", [0, [0], [0, 0, 0]]) + def test_transform_sequence_len_errors(self, param, value): + if param == "degrees" and not isinstance(value, list): + return + + kwargs = {param: value} + if param != "degrees": + kwargs["degrees"] = 0 + + with pytest.raises( + ValueError if isinstance(value, list) else TypeError, match=f"{param} should be a sequence of length 2" + ): + transforms.RandomRotation(**kwargs) + + def test_transform_negative_degrees_error(self): + with pytest.raises(ValueError, match="If degrees is a single number, it must be positive"): + transforms.RandomAffine(degrees=-1) + + def test_transform_unknown_fill_error(self): + with pytest.raises(TypeError, match="Got inappropriate fill arg"): + transforms.RandomAffine(degrees=0, fill="fill") diff --git a/test/transforms_v2_dispatcher_infos.py b/test/transforms_v2_dispatcher_infos.py index 6b13ad33861..6f61526f382 100644 --- a/test/transforms_v2_dispatcher_infos.py +++ b/test/transforms_v2_dispatcher_infos.py @@ -138,20 +138,6 @@ def fill_sequence_needs_broadcast(args_kwargs): DISPATCHER_INFOS = [ - DispatcherInfo( - F.rotate, - kernels={ - datapoints.Image: F.rotate_image_tensor, - datapoints.Video: F.rotate_video, - datapoints.BoundingBox: F.rotate_bounding_box, - datapoints.Mask: F.rotate_mask, - }, - pil_kernel_info=PILKernelInfo(F.rotate_image_pil), - test_marks=[ - xfail_jit_python_scalar_arg("fill"), - *xfails_pil_if_fill_sequence_needs_broadcast, - ], - ), DispatcherInfo( F.crop, kernels={ diff --git a/test/transforms_v2_kernel_infos.py b/test/transforms_v2_kernel_infos.py index b28b514fa38..cae8d3157e9 100644 --- a/test/transforms_v2_kernel_infos.py +++ b/test/transforms_v2_kernel_infos.py @@ -264,129 +264,6 @@ def reference_inputs_convert_format_bounding_box(): ) -_ROTATE_ANGLES = [-87, 15, 90] - - -def sample_inputs_rotate_image_tensor(): - make_rotate_image_loaders = functools.partial( - make_image_loaders, sizes=["random"], color_spaces=["RGB"], dtypes=[torch.float32] - ) - - for image_loader in make_rotate_image_loaders(): - yield ArgsKwargs(image_loader, angle=15.0, expand=True) - - for image_loader, center in itertools.product( - make_rotate_image_loaders(), [None, [1.0, 0.5], [1, 2], (1.0, 0.5), (1, 2)] - ): - yield ArgsKwargs(image_loader, angle=15.0, center=center) - - for image_loader in make_rotate_image_loaders(): - for fill in get_fills(num_channels=image_loader.num_channels, dtype=image_loader.dtype): - yield ArgsKwargs(image_loader, angle=15.0, fill=fill) - - for image_loader, interpolation in itertools.product( - make_rotate_image_loaders(), - [F.InterpolationMode.NEAREST, F.InterpolationMode.BILINEAR], - ): - yield ArgsKwargs(image_loader, angle=15.0, fill=0) - - -def reference_inputs_rotate_image_tensor(): - for image_loader, angle in itertools.product(make_image_loaders_for_interpolation(), _ROTATE_ANGLES): - yield ArgsKwargs(image_loader, angle=angle) - - -def sample_inputs_rotate_bounding_box(): - for bounding_box_loader in make_bounding_box_loaders(): - yield ArgsKwargs( - bounding_box_loader, - format=bounding_box_loader.format, - spatial_size=bounding_box_loader.spatial_size, - angle=_ROTATE_ANGLES[0], - ) - - -def reference_inputs_rotate_bounding_box(): - for bounding_box_loader, angle in itertools.product( - make_bounding_box_loaders(extra_dims=((), (4,))), _ROTATE_ANGLES - ): - yield ArgsKwargs( - bounding_box_loader, - format=bounding_box_loader.format, - spatial_size=bounding_box_loader.spatial_size, - angle=angle, - ) - - # TODO: add samples with expand=True and center - - -def reference_rotate_bounding_box(bounding_box, *, format, spatial_size, angle, expand=False, center=None): - - if center is None: - center = [spatial_size[1] * 0.5, spatial_size[0] * 0.5] - - a = np.cos(angle * np.pi / 180.0) - b = np.sin(angle * np.pi / 180.0) - cx = center[0] - cy = center[1] - affine_matrix = np.array( - [ - [a, b, cx - cx * a - b * cy], - [-b, a, cy + cx * b - a * cy], - ], - dtype="float64" if bounding_box.dtype == torch.float64 else "float32", - ) - - expected_bboxes = reference_affine_bounding_box_helper( - bounding_box, format=format, spatial_size=spatial_size, affine_matrix=affine_matrix - ) - return expected_bboxes, spatial_size - - -def sample_inputs_rotate_mask(): - for mask_loader in make_mask_loaders(sizes=["random"], num_categories=["random"], num_objects=["random"]): - yield ArgsKwargs(mask_loader, angle=15.0) - - -def sample_inputs_rotate_video(): - for video_loader in make_video_loaders(sizes=["random"], num_frames=["random"]): - yield ArgsKwargs(video_loader, angle=15.0) - - -KERNEL_INFOS.extend( - [ - KernelInfo( - F.rotate_image_tensor, - sample_inputs_fn=sample_inputs_rotate_image_tensor, - reference_fn=pil_reference_wrapper(F.rotate_image_pil), - reference_inputs_fn=reference_inputs_rotate_image_tensor, - float32_vs_uint8=True, - closeness_kwargs=pil_reference_pixel_difference(1, mae=True), - test_marks=[ - xfail_jit_python_scalar_arg("fill"), - ], - ), - KernelInfo( - F.rotate_bounding_box, - sample_inputs_fn=sample_inputs_rotate_bounding_box, - reference_fn=reference_rotate_bounding_box, - reference_inputs_fn=reference_inputs_rotate_bounding_box, - closeness_kwargs={ - **scripted_vs_eager_float64_tolerances("cpu", atol=1e-4, rtol=1e-4), - **scripted_vs_eager_float64_tolerances("cuda", atol=1e-4, rtol=1e-4), - }, - ), - KernelInfo( - F.rotate_mask, - sample_inputs_fn=sample_inputs_rotate_mask, - ), - KernelInfo( - F.rotate_video, - sample_inputs_fn=sample_inputs_rotate_video, - ), - ] -) - _CROP_PARAMS = combinations_grid(top=[-8, 0, 9], left=[-8, 0, 9], height=[12, 20], width=[12, 20])