Skip to content

GitLab

Unverified Commit 2d4484fb authored by Philip Meier's avatar Philip Meier Committed by GitHub
Browse files

port rotate (#7713)

parent c3e92565
Hide whitespace changes
Inline Side-by-side
Showing with 433 additions and 428 deletions
test/test_transforms_v2.py
View file @ 2d4484fb
......@@ -541,80 +541,6 @@ class TestRandomZoomOut:
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"):
......
test/test_transforms_v2_refactored.py
View file @ 2d4484fb
......@@ -308,6 +308,97 @@ def transform_cls_to_functional(transform_cls, **transform_specific_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 @@ class TestResize:
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 @@ class TestResize:
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 @@ class TestResize:
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 @@ class TestResize:
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 @@ class TestResize:
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 @@ class TestResize:
[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 @@ class TestResize:
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 @@ class TestResize:
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 @@ class TestResize:
[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 @@ class TestResize:
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 @@ class TestResize:
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 @@ class TestResize:
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 @@ class TestResize:
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 @@ class TestResize:
[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 @@ class TestResize:
# 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 @@ class TestResize:
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 @@ class TestHorizontalFlip:
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 @@ class TestHorizontalFlip:
],
)
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 @@ class TestHorizontalFlip:
)
@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 @@ class TestHorizontalFlip:
"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 @@ class TestHorizontalFlip:
"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 @@ class TestHorizontalFlip:
)
@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 @@ class TestHorizontalFlip:
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 @@ class TestAffine:
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 @@ class TestAffine:
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 @@ class TestAffine:
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 @@ class TestAffine:
@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 @@ class TestAffine:
],
)
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 @@ class TestAffine:
)
@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 @@ class TestAffine:
@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 @@ class TestAffine:
@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 @@ class TestAffine:
@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 @@ class TestAffine:
@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 @@ class TestAffine:
@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 @@ class TestAffine:
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 @@ class TestVerticalFlip:
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 @@ class TestVerticalFlip:
],
)
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 @@ class TestVerticalFlip:
)
@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 @@ class TestVerticalFlip:
@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 @@ class TestVerticalFlip:
)
@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")
test/transforms_v2_dispatcher_infos.py
View file @ 2d4484fb
......@@ -138,20 +138,6 @@ xfails_pil_if_fill_sequence_needs_broadcast = xfails_pil(
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={
......
test/transforms_v2_kernel_infos.py
View file @ 2d4484fb
......@@ -264,129 +264,6 @@ KERNEL_INFOS.append(
)
_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])
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment