test_transforms_v2_refactored.py 115 KB
Newer Older
1
import contextlib
2
import decimal
3
import inspect
Philip Meier's avatar
Philip Meier committed
4
import math
5
import pickle
6
import re
7
from pathlib import Path
8
9
10
11
12
13
14
15
16
17
18
19
20
from unittest import mock

import numpy as np
import PIL.Image
import pytest

import torch
import torchvision.transforms.v2 as transforms
from common_utils import (
    assert_equal,
    assert_no_warnings,
    cache,
    cpu_and_cuda,
21
    freeze_rng_state,
22
    ignore_jit_no_profile_information_warning,
23
    make_bounding_boxes,
24
25
    make_detection_mask,
    make_image,
26
27
    make_image_pil,
    make_image_tensor,
28
29
    make_segmentation_mask,
    make_video,
30
    make_video_tensor,
31
    needs_cuda,
Nicolas Hug's avatar
Nicolas Hug committed
32
    set_rng_seed,
33
)
34
35

from torch import nn
36
from torch.testing import assert_close
37
from torch.utils._pytree import tree_map
38
from torch.utils.data import DataLoader, default_collate
39
from torchvision import tv_tensors
Philip Meier's avatar
Philip Meier committed
40
41

from torchvision.transforms._functional_tensor import _max_value as get_max_value
42
43
from torchvision.transforms.functional import pil_modes_mapping
from torchvision.transforms.v2 import functional as F
44
from torchvision.transforms.v2.functional._utils import _get_kernel, _register_kernel_internal
45
46


Nicolas Hug's avatar
Nicolas Hug committed
47
48
49
50
51
52
@pytest.fixture(autouse=True)
def fix_rng_seed():
    set_rng_seed(0)
    yield


53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def _to_tolerances(maybe_tolerance_dict):
    if not isinstance(maybe_tolerance_dict, dict):
        return dict(rtol=None, atol=None)

    tolerances = dict(rtol=0, atol=0)
    tolerances.update(maybe_tolerance_dict)
    return tolerances


def _check_kernel_cuda_vs_cpu(kernel, input, *args, rtol, atol, **kwargs):
    """Checks if the kernel produces closes results for inputs on GPU and CPU."""
    if input.device.type != "cuda":
        return

    input_cuda = input.as_subclass(torch.Tensor)
    input_cpu = input_cuda.to("cpu")

70
71
72
73
    with freeze_rng_state():
        actual = kernel(input_cuda, *args, **kwargs)
    with freeze_rng_state():
        expected = kernel(input_cpu, *args, **kwargs)
74
75
76
77
78

    assert_close(actual, expected, check_device=False, rtol=rtol, atol=atol)


@cache
79
def _script(obj):
80
    try:
81
        return torch.jit.script(obj)
82
    except Exception as error:
83
84
        name = getattr(obj, "__name__", obj.__class__.__name__)
        raise AssertionError(f"Trying to `torch.jit.script` '{name}' raised the error above.") from error
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140


def _check_kernel_scripted_vs_eager(kernel, input, *args, rtol, atol, **kwargs):
    """Checks if the kernel is scriptable and if the scripted output is close to the eager one."""
    if input.device.type != "cpu":
        return

    kernel_scripted = _script(kernel)

    input = input.as_subclass(torch.Tensor)
    with ignore_jit_no_profile_information_warning():
        actual = kernel_scripted(input, *args, **kwargs)
    expected = kernel(input, *args, **kwargs)

    assert_close(actual, expected, rtol=rtol, atol=atol)


def _check_kernel_batched_vs_unbatched(kernel, input, *args, rtol, atol, **kwargs):
    """Checks if the kernel produces close results for batched and unbatched inputs."""
    unbatched_input = input.as_subclass(torch.Tensor)

    for batch_dims in [(2,), (2, 1)]:
        repeats = [*batch_dims, *[1] * input.ndim]

        actual = kernel(unbatched_input.repeat(repeats), *args, **kwargs)

        expected = kernel(unbatched_input, *args, **kwargs)
        # We can't directly call `.repeat()` on the output, since some kernel also return some additional metadata
        if isinstance(expected, torch.Tensor):
            expected = expected.repeat(repeats)
        else:
            tensor, *metadata = expected
            expected = (tensor.repeat(repeats), *metadata)

        assert_close(actual, expected, rtol=rtol, atol=atol)

    for degenerate_batch_dims in [(0,), (5, 0), (0, 5)]:
        degenerate_batched_input = torch.empty(
            degenerate_batch_dims + input.shape, dtype=input.dtype, device=input.device
        )

        output = kernel(degenerate_batched_input, *args, **kwargs)
        # Most kernels just return a tensor, but some also return some additional metadata
        if not isinstance(output, torch.Tensor):
            output, *_ = output

        assert output.shape[: -input.ndim] == degenerate_batch_dims


def check_kernel(
    kernel,
    input,
    *args,
    check_cuda_vs_cpu=True,
    check_scripted_vs_eager=True,
    check_batched_vs_unbatched=True,
141
    expect_same_dtype=True,
142
143
144
145
146
147
148
149
150
151
152
153
    **kwargs,
):
    initial_input_version = input._version

    output = kernel(input.as_subclass(torch.Tensor), *args, **kwargs)
    # Most kernels just return a tensor, but some also return some additional metadata
    if not isinstance(output, torch.Tensor):
        output, *_ = output

    # check that no inplace operation happened
    assert input._version == initial_input_version

154
155
    if expect_same_dtype:
        assert output.dtype == input.dtype
156
157
158
159
160
161
162
163
164
165
166
167
    assert output.device == input.device

    if check_cuda_vs_cpu:
        _check_kernel_cuda_vs_cpu(kernel, input, *args, **kwargs, **_to_tolerances(check_cuda_vs_cpu))

    if check_scripted_vs_eager:
        _check_kernel_scripted_vs_eager(kernel, input, *args, **kwargs, **_to_tolerances(check_scripted_vs_eager))

    if check_batched_vs_unbatched:
        _check_kernel_batched_vs_unbatched(kernel, input, *args, **kwargs, **_to_tolerances(check_batched_vs_unbatched))


Nicolas Hug's avatar
Nicolas Hug committed
168
169
def _check_functional_scripted_smoke(functional, input, *args, **kwargs):
    """Checks if the functional can be scripted and the scripted version can be called without error."""
170
    if not isinstance(input, tv_tensors.Image):
171
172
        return

Nicolas Hug's avatar
Nicolas Hug committed
173
    functional_scripted = _script(functional)
174
    with ignore_jit_no_profile_information_warning():
Nicolas Hug's avatar
Nicolas Hug committed
175
        functional_scripted(input.as_subclass(torch.Tensor), *args, **kwargs)
176
177


Nicolas Hug's avatar
Nicolas Hug committed
178
def check_functional(functional, input, *args, check_scripted_smoke=True, **kwargs):
179
    unknown_input = object()
180
    with pytest.raises(TypeError, match=re.escape(str(type(unknown_input)))):
Nicolas Hug's avatar
Nicolas Hug committed
181
        functional(unknown_input, *args, **kwargs)
182

183
    with mock.patch("torch._C._log_api_usage_once", wraps=torch._C._log_api_usage_once) as spy:
Nicolas Hug's avatar
Nicolas Hug committed
184
        output = functional(input, *args, **kwargs)
185

Nicolas Hug's avatar
Nicolas Hug committed
186
        spy.assert_any_call(f"{functional.__module__}.{functional.__name__}")
187

188
189
    assert isinstance(output, type(input))

190
    if isinstance(input, tv_tensors.BoundingBoxes):
191
192
        assert output.format == input.format

193
    if check_scripted_smoke:
Nicolas Hug's avatar
Nicolas Hug committed
194
        _check_functional_scripted_smoke(functional, input, *args, **kwargs)
195
196


Nicolas Hug's avatar
Nicolas Hug committed
197
198
199
def check_functional_kernel_signature_match(functional, *, kernel, input_type):
    """Checks if the signature of the functional matches the kernel signature."""
    functional_params = list(inspect.signature(functional).parameters.values())[1:]
200
    kernel_params = list(inspect.signature(kernel).parameters.values())[1:]
201

202
203
    if issubclass(input_type, tv_tensors.TVTensor):
        # We filter out metadata that is implicitly passed to the functional through the input tv_tensor, but has to be
204
        # explicitly passed to the kernel.
205
        explicit_metadata = {
206
            tv_tensors.BoundingBoxes: {"format", "canvas_size"},
207
208
        }
        kernel_params = [param for param in kernel_params if param.name not in explicit_metadata.get(input_type, set())]
209

Nicolas Hug's avatar
Nicolas Hug committed
210
211
    functional_params = iter(functional_params)
    for functional_param, kernel_param in zip(functional_params, kernel_params):
212
        try:
Nicolas Hug's avatar
Nicolas Hug committed
213
214
215
216
            # In general, the functional parameters are a superset of the kernel parameters. Thus, we filter out
            # functional parameters that have no kernel equivalent while keeping the order intact.
            while functional_param.name != kernel_param.name:
                functional_param = next(functional_params)
217
218
219
        except StopIteration:
            raise AssertionError(
                f"Parameter `{kernel_param.name}` of kernel `{kernel.__name__}` "
Nicolas Hug's avatar
Nicolas Hug committed
220
                f"has no corresponding parameter on the functional `{functional.__name__}`."
221
222
223
224
225
            ) from None

        if issubclass(input_type, PIL.Image.Image):
            # PIL kernels often have more correct annotations, since they are not limited by JIT. Thus, we don't check
            # them in the first place.
Nicolas Hug's avatar
Nicolas Hug committed
226
            functional_param._annotation = kernel_param._annotation = inspect.Parameter.empty
227

Nicolas Hug's avatar
Nicolas Hug committed
228
        assert functional_param == kernel_param
229
230


231
def _check_transform_v1_compatibility(transform, input, *, rtol, atol):
232
    """If the transform defines the ``_v1_transform_cls`` attribute, checks if the transform has a public, static
233
234
235
    ``get_params`` method that is the v1 equivalent, the output is close to v1, is scriptable, and the scripted version
    can be called without error."""
    if type(input) is not torch.Tensor or isinstance(input, PIL.Image.Image):
236
237
        return

238
239
    v1_transform_cls = transform._v1_transform_cls
    if v1_transform_cls is None:
240
241
        return

242
243
    if hasattr(v1_transform_cls, "get_params"):
        assert type(transform).get_params is v1_transform_cls.get_params
244

245
246
247
248
249
250
251
252
253
    v1_transform = v1_transform_cls(**transform._extract_params_for_v1_transform())

    with freeze_rng_state():
        output_v2 = transform(input)

    with freeze_rng_state():
        output_v1 = v1_transform(input)

    assert_close(output_v2, output_v1, rtol=rtol, atol=atol)
254

255
256
257
258
    if isinstance(input, PIL.Image.Image):
        return

    _script(v1_transform)(input)
259
260


261
def check_transform(transform, input, check_v1_compatibility=True):
262
263
    pickle.loads(pickle.dumps(transform))

264
265
266
    output = transform(input)
    assert isinstance(output, type(input))

267
    if isinstance(input, tv_tensors.BoundingBoxes):
268
269
        assert output.format == input.format

270
271
    if check_v1_compatibility:
        _check_transform_v1_compatibility(transform, input, **_to_tolerances(check_v1_compatibility))
272
273


274
def transform_cls_to_functional(transform_cls, **transform_specific_kwargs):
275
    def wrapper(input, *args, **kwargs):
276
        transform = transform_cls(*args, **transform_specific_kwargs, **kwargs)
277
278
279
280
281
282
283
        return transform(input)

    wrapper.__name__ = transform_cls.__name__

    return wrapper


Philip Meier's avatar
Philip Meier committed
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
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)
]


339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
# We cannot use `list(transforms.InterpolationMode)` here, since it includes some PIL-only ones as well
INTERPOLATION_MODES = [
    transforms.InterpolationMode.NEAREST,
    transforms.InterpolationMode.NEAREST_EXACT,
    transforms.InterpolationMode.BILINEAR,
    transforms.InterpolationMode.BICUBIC,
]


@contextlib.contextmanager
def assert_warns_antialias_default_value():
    with pytest.warns(UserWarning, match="The default value of the antialias parameter of all the resizing transforms"):
        yield


354
355
356
357
358
359
def reference_affine_bounding_boxes_helper(bounding_boxes, *, affine_matrix, new_canvas_size=None, clamp=True):
    format = bounding_boxes.format
    canvas_size = new_canvas_size or bounding_boxes.canvas_size

    def affine_bounding_boxes(bounding_boxes):
        dtype = bounding_boxes.dtype
360
        device = bounding_boxes.device
361

362
        # Go to float before converting to prevent precision loss in case of CXCYWH -> XYXY and W or H is 1
363
        input_xyxy = F.convert_bounding_box_format(
364
            bounding_boxes.to(dtype=torch.float64, device="cpu", copy=True),
365
            old_format=format,
366
            new_format=tv_tensors.BoundingBoxFormat.XYXY,
367
368
            inplace=True,
        )
369
370
        x1, y1, x2, y2 = input_xyxy.squeeze(0).tolist()

371
372
        points = np.array(
            [
373
374
375
376
                [x1, y1, 1.0],
                [x2, y1, 1.0],
                [x1, y2, 1.0],
                [x2, y2, 1.0],
377
378
            ]
        )
379
380
381
        transformed_points = np.matmul(points, affine_matrix.astype(points.dtype).T)

        output_xyxy = torch.Tensor(
382
            [
383
384
385
386
387
                float(np.min(transformed_points[:, 0])),
                float(np.min(transformed_points[:, 1])),
                float(np.max(transformed_points[:, 0])),
                float(np.max(transformed_points[:, 1])),
            ]
388
        )
389
390

        output = F.convert_bounding_box_format(
391
            output_xyxy, old_format=tv_tensors.BoundingBoxFormat.XYXY, new_format=format
392
393
        )

394
395
396
397
398
399
        if clamp:
            # It is important to clamp before casting, especially for CXCYWH format, dtype=int64
            output = F.clamp_bounding_boxes(
                output,
                format=format,
                canvas_size=canvas_size,
400
401
402
403
404
            )
        else:
            # We leave the bounding box as float64 so the caller gets the full precision to perform any additional
            # operation
            dtype = output.dtype
405

406
        return output.to(dtype=dtype, device=device)
407

408
    return tv_tensors.BoundingBoxes(
409
410
411
412
413
414
        torch.cat([affine_bounding_boxes(b) for b in bounding_boxes.reshape(-1, 4).unbind()], dim=0).reshape(
            bounding_boxes.shape
        ),
        format=format,
        canvas_size=canvas_size,
    )
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466


class TestResize:
    INPUT_SIZE = (17, 11)
    OUTPUT_SIZES = [17, [17], (17,), [12, 13], (12, 13)]

    def _make_max_size_kwarg(self, *, use_max_size, size):
        if use_max_size:
            if not (isinstance(size, int) or len(size) == 1):
                # This would result in an `ValueError`
                return None

            max_size = (size if isinstance(size, int) else size[0]) + 1
        else:
            max_size = None

        return dict(max_size=max_size)

    def _compute_output_size(self, *, input_size, size, max_size):
        if not (isinstance(size, int) or len(size) == 1):
            return tuple(size)

        if not isinstance(size, int):
            size = size[0]

        old_height, old_width = input_size
        ratio = old_width / old_height
        if ratio > 1:
            new_height = size
            new_width = int(ratio * new_height)
        else:
            new_width = size
            new_height = int(new_width / ratio)

        if max_size is not None and max(new_height, new_width) > max_size:
            # Need to recompute the aspect ratio, since it might have changed due to rounding
            ratio = new_width / new_height
            if ratio > 1:
                new_width = max_size
                new_height = int(new_width / ratio)
            else:
                new_height = max_size
                new_width = int(new_height * ratio)

        return new_height, new_width

    @pytest.mark.parametrize("size", OUTPUT_SIZES)
    @pytest.mark.parametrize("interpolation", INTERPOLATION_MODES)
    @pytest.mark.parametrize("use_max_size", [True, False])
    @pytest.mark.parametrize("antialias", [True, False])
    @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8])
    @pytest.mark.parametrize("device", cpu_and_cuda())
467
    def test_kernel_image(self, size, interpolation, use_max_size, antialias, dtype, device):
468
469
470
471
472
473
474
475
476
        if not (max_size_kwarg := self._make_max_size_kwarg(use_max_size=use_max_size, size=size)):
            return

        # In contrast to CPU, there is no native `InterpolationMode.BICUBIC` implementation for uint8 images on CUDA.
        # Internally, it uses the float path. Thus, we need to test with an enormous tolerance here to account for that.
        atol = 30 if transforms.InterpolationMode.BICUBIC and dtype is torch.uint8 else 1
        check_cuda_vs_cpu_tolerances = dict(rtol=0, atol=atol / 255 if dtype.is_floating_point else atol)

        check_kernel(
477
            F.resize_image,
478
            make_image(self.INPUT_SIZE, dtype=dtype, device=device),
479
480
481
482
483
484
485
486
            size=size,
            interpolation=interpolation,
            **max_size_kwarg,
            antialias=antialias,
            check_cuda_vs_cpu=check_cuda_vs_cpu_tolerances,
            check_scripted_vs_eager=not isinstance(size, int),
        )

487
    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
488
489
490
491
    @pytest.mark.parametrize("size", OUTPUT_SIZES)
    @pytest.mark.parametrize("use_max_size", [True, False])
    @pytest.mark.parametrize("dtype", [torch.float32, torch.int64])
    @pytest.mark.parametrize("device", cpu_and_cuda())
492
    def test_kernel_bounding_boxes(self, format, size, use_max_size, dtype, device):
493
494
495
        if not (max_size_kwarg := self._make_max_size_kwarg(use_max_size=use_max_size, size=size)):
            return

496
        bounding_boxes = make_bounding_boxes(
497
            format=format,
Philip Meier's avatar
Philip Meier committed
498
            canvas_size=self.INPUT_SIZE,
499
500
            dtype=dtype,
            device=device,
Philip Meier's avatar
Philip Meier committed
501
        )
502
        check_kernel(
503
504
            F.resize_bounding_boxes,
            bounding_boxes,
Philip Meier's avatar
Philip Meier committed
505
            canvas_size=bounding_boxes.canvas_size,
506
507
508
509
510
            size=size,
            **max_size_kwarg,
            check_scripted_vs_eager=not isinstance(size, int),
        )

511
512
513
    @pytest.mark.parametrize("make_mask", [make_segmentation_mask, make_detection_mask])
    def test_kernel_mask(self, make_mask):
        check_kernel(F.resize_mask, make_mask(self.INPUT_SIZE), size=self.OUTPUT_SIZES[-1])
514
515

    def test_kernel_video(self):
516
        check_kernel(F.resize_video, make_video(self.INPUT_SIZE), size=self.OUTPUT_SIZES[-1], antialias=True)
517
518
519

    @pytest.mark.parametrize("size", OUTPUT_SIZES)
    @pytest.mark.parametrize(
Philip Meier's avatar
Philip Meier committed
520
        "make_input",
521
        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
522
    )
Nicolas Hug's avatar
Nicolas Hug committed
523
524
    def test_functional(self, size, make_input):
        check_functional(
525
            F.resize,
526
            make_input(self.INPUT_SIZE),
527
528
529
530
531
532
            size=size,
            antialias=True,
            check_scripted_smoke=not isinstance(size, int),
        )

    @pytest.mark.parametrize(
533
        ("kernel", "input_type"),
534
        [
535
536
            (F.resize_image, torch.Tensor),
            (F._resize_image_pil, PIL.Image.Image),
537
538
539
540
            (F.resize_image, tv_tensors.Image),
            (F.resize_bounding_boxes, tv_tensors.BoundingBoxes),
            (F.resize_mask, tv_tensors.Mask),
            (F.resize_video, tv_tensors.Video),
541
542
        ],
    )
Nicolas Hug's avatar
Nicolas Hug committed
543
544
    def test_functional_signature(self, kernel, input_type):
        check_functional_kernel_signature_match(F.resize, kernel=kernel, input_type=input_type)
545
546
547
548

    @pytest.mark.parametrize("size", OUTPUT_SIZES)
    @pytest.mark.parametrize("device", cpu_and_cuda())
    @pytest.mark.parametrize(
549
550
551
552
553
        "make_input",
        [
            make_image_tensor,
            make_image_pil,
            make_image,
554
            make_bounding_boxes,
555
556
557
558
            make_segmentation_mask,
            make_detection_mask,
            make_video,
        ],
559
    )
560
    def test_transform(self, size, device, make_input):
561
562
563
564
565
566
        check_transform(
            transforms.Resize(size=size, antialias=True),
            make_input(self.INPUT_SIZE, device=device),
            # atol=1 due to Resize v2 is using native uint8 interpolate path for bilinear and nearest modes
            check_v1_compatibility=dict(rtol=0, atol=1),
        )
567
568

    def _check_output_size(self, input, output, *, size, max_size):
Philip Meier's avatar
Philip Meier committed
569
570
        assert tuple(F.get_size(output)) == self._compute_output_size(
            input_size=F.get_size(input), size=size, max_size=max_size
571
572
573
574
575
576
577
578
579
580
581
582
        )

    @pytest.mark.parametrize("size", OUTPUT_SIZES)
    # `InterpolationMode.NEAREST` is modeled after the buggy `INTER_NEAREST` interpolation of CV2.
    # The PIL equivalent of `InterpolationMode.NEAREST` is `InterpolationMode.NEAREST_EXACT`
    @pytest.mark.parametrize("interpolation", set(INTERPOLATION_MODES) - {transforms.InterpolationMode.NEAREST})
    @pytest.mark.parametrize("use_max_size", [True, False])
    @pytest.mark.parametrize("fn", [F.resize, transform_cls_to_functional(transforms.Resize)])
    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

583
        image = make_image(self.INPUT_SIZE, dtype=torch.uint8)
584
585

        actual = fn(image, size=size, interpolation=interpolation, **max_size_kwarg, antialias=True)
586
        expected = F.to_image(F.resize(F.to_pil_image(image), size=size, interpolation=interpolation, **max_size_kwarg))
587
588
589
590

        self._check_output_size(image, actual, size=size, **max_size_kwarg)
        torch.testing.assert_close(actual, expected, atol=1, rtol=0)

591
    def _reference_resize_bounding_boxes(self, bounding_boxes, *, size, max_size=None):
Philip Meier's avatar
Philip Meier committed
592
        old_height, old_width = bounding_boxes.canvas_size
593
        new_height, new_width = self._compute_output_size(
Philip Meier's avatar
Philip Meier committed
594
            input_size=bounding_boxes.canvas_size, size=size, max_size=max_size
595
596
597
        )

        if (old_height, old_width) == (new_height, new_width):
598
            return bounding_boxes
599
600
601
602
603
604
605
606

        affine_matrix = np.array(
            [
                [new_width / old_width, 0, 0],
                [0, new_height / old_height, 0],
            ],
        )

607
        return reference_affine_bounding_boxes_helper(
608
            bounding_boxes,
609
            affine_matrix=affine_matrix,
610
            new_canvas_size=(new_height, new_width),
611
612
        )

613
    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
614
615
616
    @pytest.mark.parametrize("size", OUTPUT_SIZES)
    @pytest.mark.parametrize("use_max_size", [True, False])
    @pytest.mark.parametrize("fn", [F.resize, transform_cls_to_functional(transforms.Resize)])
617
    def test_bounding_boxes_correctness(self, format, size, use_max_size, fn):
618
619
620
        if not (max_size_kwarg := self._make_max_size_kwarg(use_max_size=use_max_size, size=size)):
            return

621
        bounding_boxes = make_bounding_boxes(format=format, canvas_size=self.INPUT_SIZE)
622

623
624
        actual = fn(bounding_boxes, size=size, **max_size_kwarg)
        expected = self._reference_resize_bounding_boxes(bounding_boxes, size=size, **max_size_kwarg)
625

626
        self._check_output_size(bounding_boxes, actual, size=size, **max_size_kwarg)
627
628
629
630
        torch.testing.assert_close(actual, expected)

    @pytest.mark.parametrize("interpolation", set(transforms.InterpolationMode) - set(INTERPOLATION_MODES))
    @pytest.mark.parametrize(
631
632
        "make_input",
        [make_image_tensor, make_image_pil, make_image, make_video],
633
    )
634
635
    def test_pil_interpolation_compat_smoke(self, interpolation, make_input):
        input = make_input(self.INPUT_SIZE)
636
637
638
639
640
641
642
643
644
645
646
647
648

        with (
            contextlib.nullcontext()
            if isinstance(input, PIL.Image.Image)
            # This error is triggered in PyTorch core
            else pytest.raises(NotImplementedError, match=f"got {interpolation.value.lower()}")
        ):
            F.resize(
                input,
                size=self.OUTPUT_SIZES[0],
                interpolation=interpolation,
            )

Nicolas Hug's avatar
Nicolas Hug committed
649
    def test_functional_pil_antialias_warning(self):
650
        with pytest.warns(UserWarning, match="Anti-alias option is always applied for PIL Image input"):
651
            F.resize(make_image_pil(self.INPUT_SIZE), size=self.OUTPUT_SIZES[0], antialias=False)
652
653
654

    @pytest.mark.parametrize("size", OUTPUT_SIZES)
    @pytest.mark.parametrize(
655
656
657
658
659
        "make_input",
        [
            make_image_tensor,
            make_image_pil,
            make_image,
660
            make_bounding_boxes,
661
662
663
664
            make_segmentation_mask,
            make_detection_mask,
            make_video,
        ],
665
    )
666
    def test_max_size_error(self, size, make_input):
667
668
669
670
671
672
673
674
675
        if isinstance(size, int) or len(size) == 1:
            max_size = (size if isinstance(size, int) else size[0]) - 1
            match = "must be strictly greater than the requested size"
        else:
            # value can be anything other than None
            max_size = -1
            match = "size should be an int or a sequence of length 1"

        with pytest.raises(ValueError, match=match):
676
            F.resize(make_input(self.INPUT_SIZE), size=size, max_size=max_size, antialias=True)
677
678
679

    @pytest.mark.parametrize("interpolation", INTERPOLATION_MODES)
    @pytest.mark.parametrize(
680
681
        "make_input",
        [make_image_tensor, make_image, make_video],
682
    )
683
    def test_antialias_warning(self, interpolation, make_input):
684
685
686
687
688
        with (
            assert_warns_antialias_default_value()
            if interpolation in {transforms.InterpolationMode.BILINEAR, transforms.InterpolationMode.BICUBIC}
            else assert_no_warnings()
        ):
Philip Meier's avatar
Philip Meier committed
689
            F.resize(
690
                make_input(self.INPUT_SIZE),
Philip Meier's avatar
Philip Meier committed
691
692
693
                size=self.OUTPUT_SIZES[0],
                interpolation=interpolation,
            )
694
695
696

    @pytest.mark.parametrize("interpolation", INTERPOLATION_MODES)
    @pytest.mark.parametrize(
697
698
        "make_input",
        [make_image_tensor, make_image_pil, make_image, make_video],
699
    )
700
701
702
    def test_interpolation_int(self, interpolation, make_input):
        input = make_input(self.INPUT_SIZE)

703
704
705
        # `InterpolationMode.NEAREST_EXACT` has no proper corresponding integer equivalent. Internally, we map it to
        # `0` to be the same as `InterpolationMode.NEAREST` for PIL. However, for the tensor backend there is a
        # difference and thus we don't test it here.
706
        if isinstance(input, torch.Tensor) and interpolation is transforms.InterpolationMode.NEAREST_EXACT:
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
            return

        expected = F.resize(input, size=self.OUTPUT_SIZES[0], interpolation=interpolation, antialias=True)
        actual = F.resize(
            input, size=self.OUTPUT_SIZES[0], interpolation=pil_modes_mapping[interpolation], antialias=True
        )

        assert_equal(actual, expected)

    def test_transform_unknown_size_error(self):
        with pytest.raises(ValueError, match="size can either be an integer or a list or tuple of one or two integers"):
            transforms.Resize(size=object())

    @pytest.mark.parametrize(
        "size", [min(INPUT_SIZE), [min(INPUT_SIZE)], (min(INPUT_SIZE),), list(INPUT_SIZE), tuple(INPUT_SIZE)]
    )
    @pytest.mark.parametrize(
724
725
726
727
728
        "make_input",
        [
            make_image_tensor,
            make_image_pil,
            make_image,
729
            make_bounding_boxes,
730
731
732
733
            make_segmentation_mask,
            make_detection_mask,
            make_video,
        ],
734
    )
735
736
    def test_noop(self, size, make_input):
        input = make_input(self.INPUT_SIZE)
737

Philip Meier's avatar
Philip Meier committed
738
        output = F.resize(input, size=F.get_size(input), antialias=True)
739
740
741

        # This identity check is not a requirement. It is here to avoid breaking the behavior by accident. If there
        # is a good reason to break this, feel free to downgrade to an equality check.
742
        if isinstance(input, tv_tensors.TVTensor):
743
            # We can't test identity directly, since that checks for the identity of the Python object. Since all
744
            # tv_tensors unwrap before a kernel and wrap again afterwards, the Python object changes. Thus, we check
745
746
747
748
749
750
            # that the underlying storage is the same
            assert output.data_ptr() == input.data_ptr()
        else:
            assert output is input

    @pytest.mark.parametrize(
751
752
753
754
755
        "make_input",
        [
            make_image_tensor,
            make_image_pil,
            make_image,
756
            make_bounding_boxes,
757
758
759
760
            make_segmentation_mask,
            make_detection_mask,
            make_video,
        ],
761
    )
762
    def test_no_regression_5405(self, make_input):
763
764
765
        # Checks that `max_size` is not ignored if `size == small_edge_size`
        # See https://github.com/pytorch/vision/issues/5405

766
        input = make_input(self.INPUT_SIZE)
767

Philip Meier's avatar
Philip Meier committed
768
        size = min(F.get_size(input))
769
770
771
        max_size = size + 1
        output = F.resize(input, size=size, max_size=max_size, antialias=True)

Philip Meier's avatar
Philip Meier committed
772
        assert max(F.get_size(output)) == max_size
773

774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
    def _make_image(self, *args, batch_dims=(), memory_format=torch.contiguous_format, **kwargs):
        # torch.channels_last memory_format is only available for 4D tensors, i.e. (B, C, H, W). However, images coming
        # from PIL or our own I/O functions do not have a batch dimensions and are thus 3D, i.e. (C, H, W). Still, the
        # layout of the data in memory is channels last. To emulate this when a 3D input is requested here, we create
        # the image as 4D and create a view with the right shape afterwards. With this the layout in memory is channels
        # last although PyTorch doesn't recognizes it as such.
        emulate_channels_last = memory_format is torch.channels_last and len(batch_dims) != 1

        image = make_image(
            *args,
            batch_dims=(math.prod(batch_dims),) if emulate_channels_last else batch_dims,
            memory_format=memory_format,
            **kwargs,
        )

        if emulate_channels_last:
790
            image = tv_tensors.wrap(image.view(*batch_dims, *image.shape[-3:]), like=image)
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828

        return image

    def _check_stride(self, image, *, memory_format):
        C, H, W = F.get_dimensions(image)
        if memory_format is torch.contiguous_format:
            expected_stride = (H * W, W, 1)
        elif memory_format is torch.channels_last:
            expected_stride = (1, W * C, C)
        else:
            raise ValueError(f"Unknown memory_format: {memory_format}")

        assert image.stride() == expected_stride

    # TODO: We can remove this test and related torchvision workaround
    #  once we fixed related pytorch issue: https://github.com/pytorch/pytorch/issues/68430
    @pytest.mark.parametrize("interpolation", INTERPOLATION_MODES)
    @pytest.mark.parametrize("antialias", [True, False])
    @pytest.mark.parametrize("memory_format", [torch.contiguous_format, torch.channels_last])
    @pytest.mark.parametrize("dtype", [torch.uint8, torch.float32])
    @pytest.mark.parametrize("device", cpu_and_cuda())
    def test_kernel_image_memory_format_consistency(self, interpolation, antialias, memory_format, dtype, device):
        size = self.OUTPUT_SIZES[0]

        input = self._make_image(self.INPUT_SIZE, dtype=dtype, device=device, memory_format=memory_format)

        # Smoke test to make sure we aren't starting with wrong assumptions
        self._check_stride(input, memory_format=memory_format)

        output = F.resize_image(input, size=size, interpolation=interpolation, antialias=antialias)

        self._check_stride(output, memory_format=memory_format)

    def test_float16_no_rounding(self):
        # Make sure Resize() doesn't round float16 images
        # Non-regression test for https://github.com/pytorch/vision/issues/7667

        input = make_image_tensor(self.INPUT_SIZE, dtype=torch.float16)
829
        output = F.resize_image(input, size=self.OUTPUT_SIZES[0], antialias=True)
830
831
832
833

        assert output.dtype is torch.float16
        assert (output.round() - output).abs().sum() > 0

834
835
836
837

class TestHorizontalFlip:
    @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8])
    @pytest.mark.parametrize("device", cpu_and_cuda())
838
    def test_kernel_image(self, dtype, device):
839
        check_kernel(F.horizontal_flip_image, make_image(dtype=dtype, device=device))
840

841
    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
842
843
    @pytest.mark.parametrize("dtype", [torch.float32, torch.int64])
    @pytest.mark.parametrize("device", cpu_and_cuda())
844
    def test_kernel_bounding_boxes(self, format, dtype, device):
845
        bounding_boxes = make_bounding_boxes(format=format, dtype=dtype, device=device)
846
        check_kernel(
847
848
            F.horizontal_flip_bounding_boxes,
            bounding_boxes,
849
            format=format,
Philip Meier's avatar
Philip Meier committed
850
            canvas_size=bounding_boxes.canvas_size,
851
852
        )

853
854
855
    @pytest.mark.parametrize("make_mask", [make_segmentation_mask, make_detection_mask])
    def test_kernel_mask(self, make_mask):
        check_kernel(F.horizontal_flip_mask, make_mask())
856
857

    def test_kernel_video(self):
858
        check_kernel(F.horizontal_flip_video, make_video())
859
860

    @pytest.mark.parametrize(
Philip Meier's avatar
Philip Meier committed
861
        "make_input",
862
        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
863
    )
Nicolas Hug's avatar
Nicolas Hug committed
864
865
    def test_functional(self, make_input):
        check_functional(F.horizontal_flip, make_input())
866
867

    @pytest.mark.parametrize(
868
        ("kernel", "input_type"),
869
        [
870
871
            (F.horizontal_flip_image, torch.Tensor),
            (F._horizontal_flip_image_pil, PIL.Image.Image),
872
873
874
875
            (F.horizontal_flip_image, tv_tensors.Image),
            (F.horizontal_flip_bounding_boxes, tv_tensors.BoundingBoxes),
            (F.horizontal_flip_mask, tv_tensors.Mask),
            (F.horizontal_flip_video, tv_tensors.Video),
876
877
        ],
    )
Nicolas Hug's avatar
Nicolas Hug committed
878
879
    def test_functional_signature(self, kernel, input_type):
        check_functional_kernel_signature_match(F.horizontal_flip, kernel=kernel, input_type=input_type)
880
881

    @pytest.mark.parametrize(
882
        "make_input",
883
        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
884
885
    )
    @pytest.mark.parametrize("device", cpu_and_cuda())
886
    def test_transform(self, make_input, device):
887
        check_transform(transforms.RandomHorizontalFlip(p=1), make_input(device=device))
888
889
890
891
892

    @pytest.mark.parametrize(
        "fn", [F.horizontal_flip, transform_cls_to_functional(transforms.RandomHorizontalFlip, p=1)]
    )
    def test_image_correctness(self, fn):
893
        image = make_image(dtype=torch.uint8, device="cpu")
894
895

        actual = fn(image)
896
        expected = F.to_image(F.horizontal_flip(F.to_pil_image(image)))
897
898
899

        torch.testing.assert_close(actual, expected)

900
    def _reference_horizontal_flip_bounding_boxes(self, bounding_boxes):
901
902
        affine_matrix = np.array(
            [
Philip Meier's avatar
Philip Meier committed
903
                [-1, 0, bounding_boxes.canvas_size[1]],
904
905
906
907
                [0, 1, 0],
            ],
        )

908
        return reference_affine_bounding_boxes_helper(bounding_boxes, affine_matrix=affine_matrix)
909

910
    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
911
912
913
    @pytest.mark.parametrize(
        "fn", [F.horizontal_flip, transform_cls_to_functional(transforms.RandomHorizontalFlip, p=1)]
    )
914
    def test_bounding_boxes_correctness(self, format, fn):
915
        bounding_boxes = make_bounding_boxes(format=format)
916

917
918
        actual = fn(bounding_boxes)
        expected = self._reference_horizontal_flip_bounding_boxes(bounding_boxes)
919
920
921
922

        torch.testing.assert_close(actual, expected)

    @pytest.mark.parametrize(
923
        "make_input",
924
        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
925
926
    )
    @pytest.mark.parametrize("device", cpu_and_cuda())
927
928
    def test_transform_noop(self, make_input, device):
        input = make_input(device=device)
929
930
931
932
933
934

        transform = transforms.RandomHorizontalFlip(p=0)

        output = transform(input)

        assert_equal(output, input)
Philip Meier's avatar
Philip Meier committed
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977


class TestAffine:
    _EXHAUSTIVE_TYPE_AFFINE_KWARGS = dict(
        # float, int
        angle=[-10.9, 18],
        # two-list of float, two-list of int, two-tuple of float, two-tuple of int
        translate=[[6.3, -0.6], [1, -3], (16.6, -6.6), (-2, 4)],
        # float
        scale=[0.5],
        # float, int,
        # one-list of float, one-list of int, one-tuple of float, one-tuple of int
        # two-list of float, two-list of int, two-tuple of float, two-tuple of int
        shear=[35.6, 38, [-37.7], [-23], (5.3,), (-52,), [5.4, 21.8], [-47, 51], (-11.2, 36.7), (8, -53)],
        # 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)],
    )
    # The special case for shear makes sure we pick a value that is supported while JIT scripting
    _MINIMAL_AFFINE_KWARGS = {
        k: vs[0] if k != "shear" else next(v for v in vs if isinstance(v, list))
        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) and len(v) > 1)]
        for k, vs in _EXHAUSTIVE_TYPE_AFFINE_KWARGS.items()
    }

    _EXHAUSTIVE_TYPE_TRANSFORM_AFFINE_RANGES = dict(
        degrees=[30, (-15, 20)],
        translate=[None, (0.5, 0.5)],
        scale=[None, (0.75, 1.25)],
        shear=[None, (12, 30, -17, 5), 10, (-5, 12)],
    )
    _CORRECTNESS_TRANSFORM_AFFINE_RANGES = {
        k: next(v for v in vs if v is not None) for k, vs in _EXHAUSTIVE_TYPE_TRANSFORM_AFFINE_RANGES.items()
    }

    def _check_kernel(self, kernel, input, *args, **kwargs):
        kwargs_ = self._MINIMAL_AFFINE_KWARGS.copy()
        kwargs_.update(kwargs)
        check_kernel(kernel, input, *args, **kwargs_)

Philip Meier's avatar
Philip Meier committed
978
979
980
981
982
983
984
    @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,
Philip Meier's avatar
Philip Meier committed
985
986
987
    )
    @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8])
    @pytest.mark.parametrize("device", cpu_and_cuda())
988
    def test_kernel_image(self, param, value, dtype, device):
Philip Meier's avatar
Philip Meier committed
989
        if param == "fill":
Philip Meier's avatar
Philip Meier committed
990
            value = adapt_fill(value, dtype=dtype)
Philip Meier's avatar
Philip Meier committed
991
        self._check_kernel(
992
            F.affine_image,
993
            make_image(dtype=dtype, device=device),
Philip Meier's avatar
Philip Meier committed
994
995
996
997
998
999
1000
            **{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)
            if dtype is torch.uint8 and param == "interpolation" and value is transforms.InterpolationMode.BILINEAR
            else True,
        )

Philip Meier's avatar
Philip Meier committed
1001
1002
1003
1004
1005
    @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"],
Philip Meier's avatar
Philip Meier committed
1006
    )
1007
    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
Philip Meier's avatar
Philip Meier committed
1008
1009
    @pytest.mark.parametrize("dtype", [torch.float32, torch.int64])
    @pytest.mark.parametrize("device", cpu_and_cuda())
1010
    def test_kernel_bounding_boxes(self, param, value, format, dtype, device):
1011
        bounding_boxes = make_bounding_boxes(format=format, dtype=dtype, device=device)
Philip Meier's avatar
Philip Meier committed
1012
        self._check_kernel(
1013
1014
            F.affine_bounding_boxes,
            bounding_boxes,
Philip Meier's avatar
Philip Meier committed
1015
            format=format,
Philip Meier's avatar
Philip Meier committed
1016
            canvas_size=bounding_boxes.canvas_size,
Philip Meier's avatar
Philip Meier committed
1017
1018
1019
1020
            **{param: value},
            check_scripted_vs_eager=not (param == "shear" and isinstance(value, (int, float))),
        )

1021
1022
1023
    @pytest.mark.parametrize("make_mask", [make_segmentation_mask, make_detection_mask])
    def test_kernel_mask(self, make_mask):
        self._check_kernel(F.affine_mask, make_mask())
Philip Meier's avatar
Philip Meier committed
1024
1025

    def test_kernel_video(self):
1026
        self._check_kernel(F.affine_video, make_video())
Philip Meier's avatar
Philip Meier committed
1027
1028

    @pytest.mark.parametrize(
Philip Meier's avatar
Philip Meier committed
1029
        "make_input",
1030
        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
Philip Meier's avatar
Philip Meier committed
1031
    )
Nicolas Hug's avatar
Nicolas Hug committed
1032
1033
    def test_functional(self, make_input):
        check_functional(F.affine, make_input(), **self._MINIMAL_AFFINE_KWARGS)
Philip Meier's avatar
Philip Meier committed
1034
1035

    @pytest.mark.parametrize(
1036
        ("kernel", "input_type"),
Philip Meier's avatar
Philip Meier committed
1037
        [
1038
1039
            (F.affine_image, torch.Tensor),
            (F._affine_image_pil, PIL.Image.Image),
1040
1041
1042
1043
            (F.affine_image, tv_tensors.Image),
            (F.affine_bounding_boxes, tv_tensors.BoundingBoxes),
            (F.affine_mask, tv_tensors.Mask),
            (F.affine_video, tv_tensors.Video),
Philip Meier's avatar
Philip Meier committed
1044
1045
        ],
    )
Nicolas Hug's avatar
Nicolas Hug committed
1046
1047
    def test_functional_signature(self, kernel, input_type):
        check_functional_kernel_signature_match(F.affine, kernel=kernel, input_type=input_type)
Philip Meier's avatar
Philip Meier committed
1048
1049

    @pytest.mark.parametrize(
1050
        "make_input",
1051
        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
Philip Meier's avatar
Philip Meier committed
1052
1053
    )
    @pytest.mark.parametrize("device", cpu_and_cuda())
1054
1055
    def test_transform(self, make_input, device):
        input = make_input(device=device)
Philip Meier's avatar
Philip Meier committed
1056

1057
        check_transform(transforms.RandomAffine(**self._CORRECTNESS_TRANSFORM_AFFINE_RANGES), input)
Philip Meier's avatar
Philip Meier committed
1058
1059
1060
1061
1062
1063
1064
1065
1066

    @pytest.mark.parametrize("angle", _CORRECTNESS_AFFINE_KWARGS["angle"])
    @pytest.mark.parametrize("translate", _CORRECTNESS_AFFINE_KWARGS["translate"])
    @pytest.mark.parametrize("scale", _CORRECTNESS_AFFINE_KWARGS["scale"])
    @pytest.mark.parametrize("shear", _CORRECTNESS_AFFINE_KWARGS["shear"])
    @pytest.mark.parametrize("center", _CORRECTNESS_AFFINE_KWARGS["center"])
    @pytest.mark.parametrize(
        "interpolation", [transforms.InterpolationMode.NEAREST, transforms.InterpolationMode.BILINEAR]
    )
Philip Meier's avatar
Philip Meier committed
1067
    @pytest.mark.parametrize("fill", CORRECTNESS_FILLS)
Philip Meier's avatar
Philip Meier committed
1068
    def test_functional_image_correctness(self, angle, translate, scale, shear, center, interpolation, fill):
1069
        image = make_image(dtype=torch.uint8, device="cpu")
Philip Meier's avatar
Philip Meier committed
1070

Philip Meier's avatar
Philip Meier committed
1071
        fill = adapt_fill(fill, dtype=torch.uint8)
Philip Meier's avatar
Philip Meier committed
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082

        actual = F.affine(
            image,
            angle=angle,
            translate=translate,
            scale=scale,
            shear=shear,
            center=center,
            interpolation=interpolation,
            fill=fill,
        )
1083
        expected = F.to_image(
Philip Meier's avatar
Philip Meier committed
1084
            F.affine(
1085
                F.to_pil_image(image),
Philip Meier's avatar
Philip Meier committed
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
                angle=angle,
                translate=translate,
                scale=scale,
                shear=shear,
                center=center,
                interpolation=interpolation,
                fill=fill,
            )
        )

        mae = (actual.float() - expected.float()).abs().mean()
        assert mae < 2 if interpolation is transforms.InterpolationMode.NEAREST else 8

    @pytest.mark.parametrize("center", _CORRECTNESS_AFFINE_KWARGS["center"])
    @pytest.mark.parametrize(
        "interpolation", [transforms.InterpolationMode.NEAREST, transforms.InterpolationMode.BILINEAR]
    )
Philip Meier's avatar
Philip Meier committed
1103
    @pytest.mark.parametrize("fill", CORRECTNESS_FILLS)
Philip Meier's avatar
Philip Meier committed
1104
1105
    @pytest.mark.parametrize("seed", list(range(5)))
    def test_transform_image_correctness(self, center, interpolation, fill, seed):
1106
        image = make_image(dtype=torch.uint8, device="cpu")
Philip Meier's avatar
Philip Meier committed
1107

Philip Meier's avatar
Philip Meier committed
1108
        fill = adapt_fill(fill, dtype=torch.uint8)
Philip Meier's avatar
Philip Meier committed
1109
1110
1111
1112
1113
1114
1115
1116
1117

        transform = transforms.RandomAffine(
            **self._CORRECTNESS_TRANSFORM_AFFINE_RANGES, center=center, interpolation=interpolation, fill=fill
        )

        torch.manual_seed(seed)
        actual = transform(image)

        torch.manual_seed(seed)
1118
        expected = F.to_image(transform(F.to_pil_image(image)))
Philip Meier's avatar
Philip Meier committed
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142

        mae = (actual.float() - expected.float()).abs().mean()
        assert mae < 2 if interpolation is transforms.InterpolationMode.NEAREST else 8

    def _compute_affine_matrix(self, *, angle, translate, scale, shear, center):
        rot = math.radians(angle)
        cx, cy = center
        tx, ty = translate
        sx, sy = [math.radians(s) for s in ([shear, 0.0] if isinstance(shear, (int, float)) else shear)]

        c_matrix = np.array([[1, 0, cx], [0, 1, cy], [0, 0, 1]])
        t_matrix = np.array([[1, 0, tx], [0, 1, ty], [0, 0, 1]])
        c_matrix_inv = np.linalg.inv(c_matrix)
        rs_matrix = np.array(
            [
                [scale * math.cos(rot), -scale * math.sin(rot), 0],
                [scale * math.sin(rot), scale * math.cos(rot), 0],
                [0, 0, 1],
            ]
        )
        shear_x_matrix = np.array([[1, -math.tan(sx), 0], [0, 1, 0], [0, 0, 1]])
        shear_y_matrix = np.array([[1, 0, 0], [-math.tan(sy), 1, 0], [0, 0, 1]])
        rss_matrix = np.matmul(rs_matrix, np.matmul(shear_y_matrix, shear_x_matrix))
        true_matrix = np.matmul(t_matrix, np.matmul(c_matrix, np.matmul(rss_matrix, c_matrix_inv)))
1143
        return true_matrix[:2, :]
Philip Meier's avatar
Philip Meier committed
1144

1145
    def _reference_affine_bounding_boxes(self, bounding_boxes, *, angle, translate, scale, shear, center):
Philip Meier's avatar
Philip Meier committed
1146
        if center is None:
Philip Meier's avatar
Philip Meier committed
1147
            center = [s * 0.5 for s in bounding_boxes.canvas_size[::-1]]
Philip Meier's avatar
Philip Meier committed
1148

1149
        return reference_affine_bounding_boxes_helper(
1150
            bounding_boxes,
1151
1152
1153
            affine_matrix=self._compute_affine_matrix(
                angle=angle, translate=translate, scale=scale, shear=shear, center=center
            ),
Philip Meier's avatar
Philip Meier committed
1154
1155
        )

1156
    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
Philip Meier's avatar
Philip Meier committed
1157
1158
1159
1160
1161
    @pytest.mark.parametrize("angle", _CORRECTNESS_AFFINE_KWARGS["angle"])
    @pytest.mark.parametrize("translate", _CORRECTNESS_AFFINE_KWARGS["translate"])
    @pytest.mark.parametrize("scale", _CORRECTNESS_AFFINE_KWARGS["scale"])
    @pytest.mark.parametrize("shear", _CORRECTNESS_AFFINE_KWARGS["shear"])
    @pytest.mark.parametrize("center", _CORRECTNESS_AFFINE_KWARGS["center"])
1162
    def test_functional_bounding_boxes_correctness(self, format, angle, translate, scale, shear, center):
1163
        bounding_boxes = make_bounding_boxes(format=format)
Philip Meier's avatar
Philip Meier committed
1164
1165

        actual = F.affine(
1166
            bounding_boxes,
Philip Meier's avatar
Philip Meier committed
1167
1168
1169
1170
1171
1172
            angle=angle,
            translate=translate,
            scale=scale,
            shear=shear,
            center=center,
        )
1173
1174
        expected = self._reference_affine_bounding_boxes(
            bounding_boxes,
Philip Meier's avatar
Philip Meier committed
1175
1176
1177
1178
1179
1180
1181
1182
1183
            angle=angle,
            translate=translate,
            scale=scale,
            shear=shear,
            center=center,
        )

        torch.testing.assert_close(actual, expected)

1184
    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
Philip Meier's avatar
Philip Meier committed
1185
1186
    @pytest.mark.parametrize("center", _CORRECTNESS_AFFINE_KWARGS["center"])
    @pytest.mark.parametrize("seed", list(range(5)))
1187
    def test_transform_bounding_boxes_correctness(self, format, center, seed):
1188
        bounding_boxes = make_bounding_boxes(format=format)
Philip Meier's avatar
Philip Meier committed
1189
1190
1191
1192

        transform = transforms.RandomAffine(**self._CORRECTNESS_TRANSFORM_AFFINE_RANGES, center=center)

        torch.manual_seed(seed)
1193
        params = transform._get_params([bounding_boxes])
Philip Meier's avatar
Philip Meier committed
1194
1195

        torch.manual_seed(seed)
1196
        actual = transform(bounding_boxes)
Philip Meier's avatar
Philip Meier committed
1197

1198
        expected = self._reference_affine_bounding_boxes(bounding_boxes, **params, center=center)
Philip Meier's avatar
Philip Meier committed
1199
1200
1201
1202
1203
1204
1205
1206
1207

        torch.testing.assert_close(actual, expected)

    @pytest.mark.parametrize("degrees", _EXHAUSTIVE_TYPE_TRANSFORM_AFFINE_RANGES["degrees"])
    @pytest.mark.parametrize("translate", _EXHAUSTIVE_TYPE_TRANSFORM_AFFINE_RANGES["translate"])
    @pytest.mark.parametrize("scale", _EXHAUSTIVE_TYPE_TRANSFORM_AFFINE_RANGES["scale"])
    @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):
1208
        image = make_image()
Philip Meier's avatar
Philip Meier committed
1209
        height, width = F.get_size(image)
Philip Meier's avatar
Philip Meier committed
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282

        transform = transforms.RandomAffine(degrees=degrees, translate=translate, scale=scale, shear=shear)

        torch.manual_seed(seed)
        params = transform._get_params([image])

        if isinstance(degrees, (int, float)):
            assert -degrees <= params["angle"] <= degrees
        else:
            assert degrees[0] <= params["angle"] <= degrees[1]

        if translate is not None:
            width_max = int(round(translate[0] * width))
            height_max = int(round(translate[1] * height))
            assert -width_max <= params["translate"][0] <= width_max
            assert -height_max <= params["translate"][1] <= height_max
        else:
            assert params["translate"] == (0, 0)

        if scale is not None:
            assert scale[0] <= params["scale"] <= scale[1]
        else:
            assert params["scale"] == 1.0

        if shear is not None:
            if isinstance(shear, (int, float)):
                assert -shear <= params["shear"][0] <= shear
                assert params["shear"][1] == 0.0
            elif len(shear) == 2:
                assert shear[0] <= params["shear"][0] <= shear[1]
                assert params["shear"][1] == 0.0
            elif len(shear) == 4:
                assert shear[0] <= params["shear"][0] <= shear[1]
                assert shear[2] <= params["shear"][1] <= shear[3]
        else:
            assert params["shear"] == (0, 0)

    @pytest.mark.parametrize("param", ["degrees", "translate", "scale", "shear", "center"])
    @pytest.mark.parametrize("value", [0, [0], [0, 0, 0]])
    def test_transform_sequence_len_errors(self, param, value):
        if param in {"degrees", "shear"} 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.RandomAffine(**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)

    @pytest.mark.parametrize("translate", [[-1, 0], [2, 0], [-1, 2]])
    def test_transform_translate_range_error(self, translate):
        with pytest.raises(ValueError, match="translation values should be between 0 and 1"):
            transforms.RandomAffine(degrees=0, translate=translate)

    @pytest.mark.parametrize("scale", [[-1, 0], [0, -1], [-1, -1]])
    def test_transform_scale_range_error(self, scale):
        with pytest.raises(ValueError, match="scale values should be positive"):
            transforms.RandomAffine(degrees=0, scale=scale)

    def test_transform_negative_shear_error(self):
        with pytest.raises(ValueError, match="If shear is a single number, it must be positive"):
            transforms.RandomAffine(degrees=0, shear=-1)

    def test_transform_unknown_fill_error(self):
        with pytest.raises(TypeError, match="Got inappropriate fill arg"):
            transforms.RandomAffine(degrees=0, fill="fill")
Philip Meier's avatar
Philip Meier committed
1283
1284
1285
1286
1287


class TestVerticalFlip:
    @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8])
    @pytest.mark.parametrize("device", cpu_and_cuda())
1288
    def test_kernel_image(self, dtype, device):
1289
        check_kernel(F.vertical_flip_image, make_image(dtype=dtype, device=device))
Philip Meier's avatar
Philip Meier committed
1290

1291
    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
Philip Meier's avatar
Philip Meier committed
1292
1293
    @pytest.mark.parametrize("dtype", [torch.float32, torch.int64])
    @pytest.mark.parametrize("device", cpu_and_cuda())
1294
    def test_kernel_bounding_boxes(self, format, dtype, device):
1295
        bounding_boxes = make_bounding_boxes(format=format, dtype=dtype, device=device)
Philip Meier's avatar
Philip Meier committed
1296
        check_kernel(
1297
1298
            F.vertical_flip_bounding_boxes,
            bounding_boxes,
Philip Meier's avatar
Philip Meier committed
1299
            format=format,
Philip Meier's avatar
Philip Meier committed
1300
            canvas_size=bounding_boxes.canvas_size,
Philip Meier's avatar
Philip Meier committed
1301
1302
        )

1303
1304
1305
    @pytest.mark.parametrize("make_mask", [make_segmentation_mask, make_detection_mask])
    def test_kernel_mask(self, make_mask):
        check_kernel(F.vertical_flip_mask, make_mask())
Philip Meier's avatar
Philip Meier committed
1306
1307

    def test_kernel_video(self):
1308
        check_kernel(F.vertical_flip_video, make_video())
Philip Meier's avatar
Philip Meier committed
1309
1310

    @pytest.mark.parametrize(
Philip Meier's avatar
Philip Meier committed
1311
        "make_input",
1312
        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
Philip Meier's avatar
Philip Meier committed
1313
    )
Nicolas Hug's avatar
Nicolas Hug committed
1314
1315
    def test_functional(self, make_input):
        check_functional(F.vertical_flip, make_input())
Philip Meier's avatar
Philip Meier committed
1316
1317

    @pytest.mark.parametrize(
1318
        ("kernel", "input_type"),
Philip Meier's avatar
Philip Meier committed
1319
        [
1320
1321
            (F.vertical_flip_image, torch.Tensor),
            (F._vertical_flip_image_pil, PIL.Image.Image),
1322
1323
1324
1325
            (F.vertical_flip_image, tv_tensors.Image),
            (F.vertical_flip_bounding_boxes, tv_tensors.BoundingBoxes),
            (F.vertical_flip_mask, tv_tensors.Mask),
            (F.vertical_flip_video, tv_tensors.Video),
Philip Meier's avatar
Philip Meier committed
1326
1327
        ],
    )
Nicolas Hug's avatar
Nicolas Hug committed
1328
1329
    def test_functional_signature(self, kernel, input_type):
        check_functional_kernel_signature_match(F.vertical_flip, kernel=kernel, input_type=input_type)
Philip Meier's avatar
Philip Meier committed
1330
1331

    @pytest.mark.parametrize(
1332
        "make_input",
1333
        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
Philip Meier's avatar
Philip Meier committed
1334
1335
    )
    @pytest.mark.parametrize("device", cpu_and_cuda())
1336
    def test_transform(self, make_input, device):
1337
        check_transform(transforms.RandomVerticalFlip(p=1), make_input(device=device))
Philip Meier's avatar
Philip Meier committed
1338
1339
1340

    @pytest.mark.parametrize("fn", [F.vertical_flip, transform_cls_to_functional(transforms.RandomVerticalFlip, p=1)])
    def test_image_correctness(self, fn):
1341
        image = make_image(dtype=torch.uint8, device="cpu")
Philip Meier's avatar
Philip Meier committed
1342
1343

        actual = fn(image)
1344
        expected = F.to_image(F.vertical_flip(F.to_pil_image(image)))
Philip Meier's avatar
Philip Meier committed
1345
1346
1347

        torch.testing.assert_close(actual, expected)

1348
    def _reference_vertical_flip_bounding_boxes(self, bounding_boxes):
Philip Meier's avatar
Philip Meier committed
1349
1350
1351
        affine_matrix = np.array(
            [
                [1, 0, 0],
Philip Meier's avatar
Philip Meier committed
1352
                [0, -1, bounding_boxes.canvas_size[0]],
Philip Meier's avatar
Philip Meier committed
1353
1354
1355
            ],
        )

1356
        return reference_affine_bounding_boxes_helper(bounding_boxes, affine_matrix=affine_matrix)
Philip Meier's avatar
Philip Meier committed
1357

1358
    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
Philip Meier's avatar
Philip Meier committed
1359
    @pytest.mark.parametrize("fn", [F.vertical_flip, transform_cls_to_functional(transforms.RandomVerticalFlip, p=1)])
1360
    def test_bounding_boxes_correctness(self, format, fn):
1361
        bounding_boxes = make_bounding_boxes(format=format)
Philip Meier's avatar
Philip Meier committed
1362

1363
1364
        actual = fn(bounding_boxes)
        expected = self._reference_vertical_flip_bounding_boxes(bounding_boxes)
Philip Meier's avatar
Philip Meier committed
1365
1366
1367
1368

        torch.testing.assert_close(actual, expected)

    @pytest.mark.parametrize(
1369
        "make_input",
1370
        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
Philip Meier's avatar
Philip Meier committed
1371
1372
    )
    @pytest.mark.parametrize("device", cpu_and_cuda())
1373
1374
    def test_transform_noop(self, make_input, device):
        input = make_input(device=device)
Philip Meier's avatar
Philip Meier committed
1375
1376
1377
1378
1379
1380

        transform = transforms.RandomVerticalFlip(p=0)

        output = transform(input)

        assert_equal(output, input)
Philip Meier's avatar
Philip Meier committed
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410


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())
1411
    def test_kernel_image(self, param, value, dtype, device):
Philip Meier's avatar
Philip Meier committed
1412
1413
1414
1415
        kwargs = {param: value}
        if param != "angle":
            kwargs["angle"] = self._MINIMAL_AFFINE_KWARGS["angle"]
        check_kernel(
1416
            F.rotate_image,
1417
            make_image(dtype=dtype, device=device),
Philip Meier's avatar
Philip Meier committed
1418
1419
1420
1421
1422
1423
1424
1425
1426
            **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"],
    )
1427
    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
Philip Meier's avatar
Philip Meier committed
1428
1429
    @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8])
    @pytest.mark.parametrize("device", cpu_and_cuda())
1430
    def test_kernel_bounding_boxes(self, param, value, format, dtype, device):
Philip Meier's avatar
Philip Meier committed
1431
1432
1433
1434
        kwargs = {param: value}
        if param != "angle":
            kwargs["angle"] = self._MINIMAL_AFFINE_KWARGS["angle"]

1435
        bounding_boxes = make_bounding_boxes(format=format, dtype=dtype, device=device)
Philip Meier's avatar
Philip Meier committed
1436
1437

        check_kernel(
1438
1439
            F.rotate_bounding_boxes,
            bounding_boxes,
Philip Meier's avatar
Philip Meier committed
1440
            format=format,
Philip Meier's avatar
Philip Meier committed
1441
            canvas_size=bounding_boxes.canvas_size,
Philip Meier's avatar
Philip Meier committed
1442
1443
1444
            **kwargs,
        )

1445
1446
1447
    @pytest.mark.parametrize("make_mask", [make_segmentation_mask, make_detection_mask])
    def test_kernel_mask(self, make_mask):
        check_kernel(F.rotate_mask, make_mask(), **self._MINIMAL_AFFINE_KWARGS)
Philip Meier's avatar
Philip Meier committed
1448
1449

    def test_kernel_video(self):
1450
        check_kernel(F.rotate_video, make_video(), **self._MINIMAL_AFFINE_KWARGS)
Philip Meier's avatar
Philip Meier committed
1451
1452

    @pytest.mark.parametrize(
Philip Meier's avatar
Philip Meier committed
1453
        "make_input",
1454
        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
Philip Meier's avatar
Philip Meier committed
1455
    )
Nicolas Hug's avatar
Nicolas Hug committed
1456
1457
    def test_functional(self, make_input):
        check_functional(F.rotate, make_input(), **self._MINIMAL_AFFINE_KWARGS)
Philip Meier's avatar
Philip Meier committed
1458
1459

    @pytest.mark.parametrize(
1460
        ("kernel", "input_type"),
Philip Meier's avatar
Philip Meier committed
1461
        [
1462
1463
            (F.rotate_image, torch.Tensor),
            (F._rotate_image_pil, PIL.Image.Image),
1464
1465
1466
1467
            (F.rotate_image, tv_tensors.Image),
            (F.rotate_bounding_boxes, tv_tensors.BoundingBoxes),
            (F.rotate_mask, tv_tensors.Mask),
            (F.rotate_video, tv_tensors.Video),
Philip Meier's avatar
Philip Meier committed
1468
1469
        ],
    )
Nicolas Hug's avatar
Nicolas Hug committed
1470
1471
    def test_functional_signature(self, kernel, input_type):
        check_functional_kernel_signature_match(F.rotate, kernel=kernel, input_type=input_type)
Philip Meier's avatar
Philip Meier committed
1472
1473

    @pytest.mark.parametrize(
1474
        "make_input",
1475
        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
Philip Meier's avatar
Philip Meier committed
1476
1477
    )
    @pytest.mark.parametrize("device", cpu_and_cuda())
1478
1479
    def test_transform(self, make_input, device):
        check_transform(
1480
            transforms.RandomRotation(**self._CORRECTNESS_TRANSFORM_AFFINE_RANGES), make_input(device=device)
1481
        )
Philip Meier's avatar
Philip Meier committed
1482
1483
1484
1485
1486
1487
1488
1489
1490

    @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):
1491
        image = make_image(dtype=torch.uint8, device="cpu")
Philip Meier's avatar
Philip Meier committed
1492
1493
1494
1495

        fill = adapt_fill(fill, dtype=torch.uint8)

        actual = F.rotate(image, angle=angle, center=center, interpolation=interpolation, expand=expand, fill=fill)
1496
        expected = F.to_image(
Philip Meier's avatar
Philip Meier committed
1497
            F.rotate(
1498
                F.to_pil_image(image), angle=angle, center=center, interpolation=interpolation, expand=expand, fill=fill
Philip Meier's avatar
Philip Meier committed
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
            )
        )

        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):
1513
        image = make_image(dtype=torch.uint8, device="cpu")
Philip Meier's avatar
Philip Meier committed
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528

        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)
1529
        expected = F.to_image(transform(F.to_pil_image(image)))
Philip Meier's avatar
Philip Meier committed
1530
1531
1532
1533

        mae = (actual.float() - expected.float()).abs().mean()
        assert mae < 1 if interpolation is transforms.InterpolationMode.NEAREST else 6

1534
1535
1536
1537
1538
    def _compute_output_canvas_size(self, *, expand, canvas_size, affine_matrix):
        if not expand:
            return canvas_size, (0.0, 0.0)

        input_height, input_width = canvas_size
Philip Meier's avatar
Philip Meier committed
1539

1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
        input_image_frame = np.array(
            [
                [0.0, 0.0, 1.0],
                [0.0, input_height, 1.0],
                [input_width, input_height, 1.0],
                [input_width, 0.0, 1.0],
            ],
            dtype=np.float64,
        )
        output_image_frame = np.matmul(input_image_frame, affine_matrix.astype(input_image_frame.dtype).T)

        recenter_x = float(np.min(output_image_frame[:, 0]))
        recenter_y = float(np.min(output_image_frame[:, 1]))

        output_width = int(np.max(output_image_frame[:, 0]) - recenter_x)
        output_height = int(np.max(output_image_frame[:, 1]) - recenter_y)

        return (output_height, output_width), (recenter_x, recenter_y)

    def _recenter_bounding_boxes_after_expand(self, bounding_boxes, *, recenter_xy):
        x, y = recenter_xy
1561
        if bounding_boxes.format is tv_tensors.BoundingBoxFormat.XYXY:
1562
1563
1564
            translate = [x, y, x, y]
        else:
            translate = [x, y, 0.0, 0.0]
1565
        return tv_tensors.wrap(
1566
1567
1568
1569
            (bounding_boxes.to(torch.float64) - torch.tensor(translate)).to(bounding_boxes.dtype), like=bounding_boxes
        )

    def _reference_rotate_bounding_boxes(self, bounding_boxes, *, angle, expand, center):
Philip Meier's avatar
Philip Meier committed
1570
        if center is None:
Philip Meier's avatar
Philip Meier committed
1571
            center = [s * 0.5 for s in bounding_boxes.canvas_size[::-1]]
1572
        cx, cy = center
Philip Meier's avatar
Philip Meier committed
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582

        a = np.cos(angle * np.pi / 180.0)
        b = np.sin(angle * np.pi / 180.0)
        affine_matrix = np.array(
            [
                [a, b, cx - cx * a - b * cy],
                [-b, a, cy + cx * b - a * cy],
            ],
        )

1583
1584
1585
1586
1587
        new_canvas_size, recenter_xy = self._compute_output_canvas_size(
            expand=expand, canvas_size=bounding_boxes.canvas_size, affine_matrix=affine_matrix
        )

        output = reference_affine_bounding_boxes_helper(
1588
            bounding_boxes,
Philip Meier's avatar
Philip Meier committed
1589
            affine_matrix=affine_matrix,
1590
1591
            new_canvas_size=new_canvas_size,
            clamp=False,
Philip Meier's avatar
Philip Meier committed
1592
1593
        )

1594
1595
1596
        return F.clamp_bounding_boxes(self._recenter_bounding_boxes_after_expand(output, recenter_xy=recenter_xy)).to(
            bounding_boxes
        )
Philip Meier's avatar
Philip Meier committed
1597

1598
    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
Philip Meier's avatar
Philip Meier committed
1599
    @pytest.mark.parametrize("angle", _CORRECTNESS_AFFINE_KWARGS["angle"])
1600
    @pytest.mark.parametrize("expand", [False, True])
Philip Meier's avatar
Philip Meier committed
1601
    @pytest.mark.parametrize("center", _CORRECTNESS_AFFINE_KWARGS["center"])
1602
    def test_functional_bounding_boxes_correctness(self, format, angle, expand, center):
1603
        bounding_boxes = make_bounding_boxes(format=format)
Philip Meier's avatar
Philip Meier committed
1604

1605
1606
        actual = F.rotate(bounding_boxes, angle=angle, expand=expand, center=center)
        expected = self._reference_rotate_bounding_boxes(bounding_boxes, angle=angle, expand=expand, center=center)
Philip Meier's avatar
Philip Meier committed
1607
1608

        torch.testing.assert_close(actual, expected)
1609
        torch.testing.assert_close(F.get_size(actual), F.get_size(expected), atol=2 if expand else 0, rtol=0)
Philip Meier's avatar
Philip Meier committed
1610

1611
    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
1612
    @pytest.mark.parametrize("expand", [False, True])
Philip Meier's avatar
Philip Meier committed
1613
1614
    @pytest.mark.parametrize("center", _CORRECTNESS_AFFINE_KWARGS["center"])
    @pytest.mark.parametrize("seed", list(range(5)))
1615
    def test_transform_bounding_boxes_correctness(self, format, expand, center, seed):
1616
        bounding_boxes = make_bounding_boxes(format=format)
Philip Meier's avatar
Philip Meier committed
1617
1618
1619
1620

        transform = transforms.RandomRotation(**self._CORRECTNESS_TRANSFORM_AFFINE_RANGES, expand=expand, center=center)

        torch.manual_seed(seed)
1621
        params = transform._get_params([bounding_boxes])
Philip Meier's avatar
Philip Meier committed
1622
1623

        torch.manual_seed(seed)
1624
        actual = transform(bounding_boxes)
Philip Meier's avatar
Philip Meier committed
1625

1626
        expected = self._reference_rotate_bounding_boxes(bounding_boxes, **params, expand=expand, center=center)
Philip Meier's avatar
Philip Meier committed
1627
1628

        torch.testing.assert_close(actual, expected)
1629
        torch.testing.assert_close(F.get_size(actual), F.get_size(expected), atol=2 if expand else 0, rtol=0)
Philip Meier's avatar
Philip Meier committed
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665

    @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")
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726


class TestCompose:
    class BuiltinTransform(transforms.Transform):
        def _transform(self, inpt, params):
            return inpt

    class PackedInputTransform(nn.Module):
        def forward(self, sample):
            assert len(sample) == 2
            return sample

    class UnpackedInputTransform(nn.Module):
        def forward(self, image, label):
            return image, label

    @pytest.mark.parametrize(
        "transform_clss",
        [
            [BuiltinTransform],
            [PackedInputTransform],
            [UnpackedInputTransform],
            [BuiltinTransform, BuiltinTransform],
            [PackedInputTransform, PackedInputTransform],
            [UnpackedInputTransform, UnpackedInputTransform],
            [BuiltinTransform, PackedInputTransform, BuiltinTransform],
            [BuiltinTransform, UnpackedInputTransform, BuiltinTransform],
            [PackedInputTransform, BuiltinTransform, PackedInputTransform],
            [UnpackedInputTransform, BuiltinTransform, UnpackedInputTransform],
        ],
    )
    @pytest.mark.parametrize("unpack", [True, False])
    def test_packed_unpacked(self, transform_clss, unpack):
        needs_packed_inputs = any(issubclass(cls, self.PackedInputTransform) for cls in transform_clss)
        needs_unpacked_inputs = any(issubclass(cls, self.UnpackedInputTransform) for cls in transform_clss)
        assert not (needs_packed_inputs and needs_unpacked_inputs)

        transform = transforms.Compose([cls() for cls in transform_clss])

        image = make_image()
        label = 3
        packed_input = (image, label)

        def call_transform():
            if unpack:
                return transform(*packed_input)
            else:
                return transform(packed_input)

        if needs_unpacked_inputs and not unpack:
            with pytest.raises(TypeError, match="missing 1 required positional argument"):
                call_transform()
        elif needs_packed_inputs and unpack:
            with pytest.raises(TypeError, match="takes 2 positional arguments but 3 were given"):
                call_transform()
        else:
            output = call_transform()

            assert isinstance(output, tuple) and len(output) == 2
            assert output[0] is image
            assert output[1] is label
1727
1728
1729
1730
1731
1732


class TestToDtype:
    @pytest.mark.parametrize(
        ("kernel", "make_input"),
        [
1733
1734
            (F.to_dtype_image, make_image_tensor),
            (F.to_dtype_image, make_image),
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
            (F.to_dtype_video, make_video),
        ],
    )
    @pytest.mark.parametrize("input_dtype", [torch.float32, torch.float64, torch.uint8])
    @pytest.mark.parametrize("output_dtype", [torch.float32, torch.float64, torch.uint8])
    @pytest.mark.parametrize("device", cpu_and_cuda())
    @pytest.mark.parametrize("scale", (True, False))
    def test_kernel(self, kernel, make_input, input_dtype, output_dtype, device, scale):
        check_kernel(
            kernel,
            make_input(dtype=input_dtype, device=device),
            expect_same_dtype=input_dtype is output_dtype,
            dtype=output_dtype,
            scale=scale,
        )

Philip Meier's avatar
Philip Meier committed
1751
    @pytest.mark.parametrize("make_input", [make_image_tensor, make_image, make_video])
1752
1753
1754
1755
    @pytest.mark.parametrize("input_dtype", [torch.float32, torch.float64, torch.uint8])
    @pytest.mark.parametrize("output_dtype", [torch.float32, torch.float64, torch.uint8])
    @pytest.mark.parametrize("device", cpu_and_cuda())
    @pytest.mark.parametrize("scale", (True, False))
Nicolas Hug's avatar
Nicolas Hug committed
1756
1757
    def test_functional(self, make_input, input_dtype, output_dtype, device, scale):
        check_functional(
1758
1759
1760
1761
1762
1763
1764
1765
            F.to_dtype,
            make_input(dtype=input_dtype, device=device),
            dtype=output_dtype,
            scale=scale,
        )

    @pytest.mark.parametrize(
        "make_input",
1766
        [make_image_tensor, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
    )
    @pytest.mark.parametrize("input_dtype", [torch.float32, torch.float64, torch.uint8])
    @pytest.mark.parametrize("output_dtype", [torch.float32, torch.float64, torch.uint8])
    @pytest.mark.parametrize("device", cpu_and_cuda())
    @pytest.mark.parametrize("scale", (True, False))
    @pytest.mark.parametrize("as_dict", (True, False))
    def test_transform(self, make_input, input_dtype, output_dtype, device, scale, as_dict):
        input = make_input(dtype=input_dtype, device=device)
        if as_dict:
            output_dtype = {type(input): output_dtype}
1777
        check_transform(transforms.ToDtype(dtype=output_dtype, scale=scale), input)
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840

    def reference_convert_dtype_image_tensor(self, image, dtype=torch.float, scale=False):
        input_dtype = image.dtype
        output_dtype = dtype

        if not scale:
            return image.to(dtype)

        if output_dtype == input_dtype:
            return image

        def fn(value):
            if input_dtype.is_floating_point:
                if output_dtype.is_floating_point:
                    return value
                else:
                    return round(decimal.Decimal(value) * torch.iinfo(output_dtype).max)
            else:
                input_max_value = torch.iinfo(input_dtype).max

                if output_dtype.is_floating_point:
                    return float(decimal.Decimal(value) / input_max_value)
                else:
                    output_max_value = torch.iinfo(output_dtype).max

                    if input_max_value > output_max_value:
                        factor = (input_max_value + 1) // (output_max_value + 1)
                        return value / factor
                    else:
                        factor = (output_max_value + 1) // (input_max_value + 1)
                        return value * factor

        return torch.tensor(tree_map(fn, image.tolist()), dtype=dtype, device=image.device)

    @pytest.mark.parametrize("input_dtype", [torch.float32, torch.float64, torch.uint8])
    @pytest.mark.parametrize("output_dtype", [torch.float32, torch.float64, torch.uint8])
    @pytest.mark.parametrize("device", cpu_and_cuda())
    @pytest.mark.parametrize("scale", (True, False))
    def test_image_correctness(self, input_dtype, output_dtype, device, scale):
        if input_dtype.is_floating_point and output_dtype == torch.int64:
            pytest.xfail("float to int64 conversion is not supported")

        input = make_image(dtype=input_dtype, device=device)

        out = F.to_dtype(input, dtype=output_dtype, scale=scale)
        expected = self.reference_convert_dtype_image_tensor(input, dtype=output_dtype, scale=scale)

        if input_dtype.is_floating_point and not output_dtype.is_floating_point and scale:
            torch.testing.assert_close(out, expected, atol=1, rtol=0)
        else:
            torch.testing.assert_close(out, expected)

    def was_scaled(self, inpt):
        # this assumes the target dtype is float
        return inpt.max() <= 1

    def make_inpt_with_bbox_and_mask(self, make_input):
        H, W = 10, 10
        inpt_dtype = torch.uint8
        bbox_dtype = torch.float32
        mask_dtype = torch.bool
        sample = {
            "inpt": make_input(size=(H, W), dtype=inpt_dtype),
1841
            "bbox": make_bounding_boxes(canvas_size=(H, W), dtype=bbox_dtype),
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
            "mask": make_detection_mask(size=(H, W), dtype=mask_dtype),
        }

        return sample, inpt_dtype, bbox_dtype, mask_dtype

    @pytest.mark.parametrize("make_input", (make_image_tensor, make_image, make_video))
    @pytest.mark.parametrize("scale", (True, False))
    def test_dtype_not_a_dict(self, make_input, scale):
        # assert only inpt gets transformed when dtype isn't a dict

        sample, inpt_dtype, bbox_dtype, mask_dtype = self.make_inpt_with_bbox_and_mask(make_input)
        out = transforms.ToDtype(dtype=torch.float32, scale=scale)(sample)

        assert out["inpt"].dtype != inpt_dtype
        assert out["inpt"].dtype == torch.float32
        if scale:
            assert self.was_scaled(out["inpt"])
        else:
            assert not self.was_scaled(out["inpt"])
        assert out["bbox"].dtype == bbox_dtype
        assert out["mask"].dtype == mask_dtype

    @pytest.mark.parametrize("make_input", (make_image_tensor, make_image, make_video))
    def test_others_catch_all_and_none(self, make_input):
        # make sure "others" works as a catch-all and that None means no conversion

        sample, inpt_dtype, bbox_dtype, mask_dtype = self.make_inpt_with_bbox_and_mask(make_input)
1869
        out = transforms.ToDtype(dtype={tv_tensors.Mask: torch.int64, "others": None})(sample)
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
        assert out["inpt"].dtype == inpt_dtype
        assert out["bbox"].dtype == bbox_dtype
        assert out["mask"].dtype != mask_dtype
        assert out["mask"].dtype == torch.int64

    @pytest.mark.parametrize("make_input", (make_image_tensor, make_image, make_video))
    def test_typical_use_case(self, make_input):
        # Typical use-case: want to convert dtype and scale for inpt and just dtype for masks.
        # This just makes sure we now have a decent API for this

        sample, inpt_dtype, bbox_dtype, mask_dtype = self.make_inpt_with_bbox_and_mask(make_input)
        out = transforms.ToDtype(
1882
            dtype={type(sample["inpt"]): torch.float32, tv_tensors.Mask: torch.int64, "others": None}, scale=True
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
        )(sample)
        assert out["inpt"].dtype != inpt_dtype
        assert out["inpt"].dtype == torch.float32
        assert self.was_scaled(out["inpt"])
        assert out["bbox"].dtype == bbox_dtype
        assert out["mask"].dtype != mask_dtype
        assert out["mask"].dtype == torch.int64

    @pytest.mark.parametrize("make_input", (make_image_tensor, make_image, make_video))
    def test_errors_warnings(self, make_input):
        sample, inpt_dtype, bbox_dtype, mask_dtype = self.make_inpt_with_bbox_and_mask(make_input)

        with pytest.raises(ValueError, match="No dtype was specified for"):
1896
            out = transforms.ToDtype(dtype={tv_tensors.Mask: torch.float32})(sample)
1897
        with pytest.warns(UserWarning, match=re.escape("plain `torch.Tensor` will *not* be transformed")):
1898
            transforms.ToDtype(dtype={torch.Tensor: torch.float32, tv_tensors.Image: torch.float32})
1899
1900
1901
1902
1903
        with pytest.warns(UserWarning, match="no scaling will be done"):
            out = transforms.ToDtype(dtype={"others": None}, scale=True)(sample)
        assert out["inpt"].dtype == inpt_dtype
        assert out["bbox"].dtype == bbox_dtype
        assert out["mask"].dtype == mask_dtype
1904
1905


1906
1907
1908
1909
1910
1911
1912
class TestAdjustBrightness:
    _CORRECTNESS_BRIGHTNESS_FACTORS = [0.5, 0.0, 1.0, 5.0]
    _DEFAULT_BRIGHTNESS_FACTOR = _CORRECTNESS_BRIGHTNESS_FACTORS[0]

    @pytest.mark.parametrize(
        ("kernel", "make_input"),
        [
1913
            (F.adjust_brightness_image, make_image),
1914
1915
1916
1917
1918
1919
1920
1921
            (F.adjust_brightness_video, make_video),
        ],
    )
    @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8])
    @pytest.mark.parametrize("device", cpu_and_cuda())
    def test_kernel(self, kernel, make_input, dtype, device):
        check_kernel(kernel, make_input(dtype=dtype, device=device), brightness_factor=self._DEFAULT_BRIGHTNESS_FACTOR)

Philip Meier's avatar
Philip Meier committed
1922
    @pytest.mark.parametrize("make_input", [make_image_tensor, make_image_pil, make_image, make_video])
Nicolas Hug's avatar
Nicolas Hug committed
1923
1924
    def test_functional(self, make_input):
        check_functional(F.adjust_brightness, make_input(), brightness_factor=self._DEFAULT_BRIGHTNESS_FACTOR)
1925
1926
1927
1928

    @pytest.mark.parametrize(
        ("kernel", "input_type"),
        [
1929
1930
            (F.adjust_brightness_image, torch.Tensor),
            (F._adjust_brightness_image_pil, PIL.Image.Image),
1931
1932
            (F.adjust_brightness_image, tv_tensors.Image),
            (F.adjust_brightness_video, tv_tensors.Video),
1933
1934
        ],
    )
Nicolas Hug's avatar
Nicolas Hug committed
1935
1936
    def test_functional_signature(self, kernel, input_type):
        check_functional_kernel_signature_match(F.adjust_brightness, kernel=kernel, input_type=input_type)
1937
1938
1939
1940
1941
1942

    @pytest.mark.parametrize("brightness_factor", _CORRECTNESS_BRIGHTNESS_FACTORS)
    def test_image_correctness(self, brightness_factor):
        image = make_image(dtype=torch.uint8, device="cpu")

        actual = F.adjust_brightness(image, brightness_factor=brightness_factor)
1943
        expected = F.to_image(F.adjust_brightness(F.to_pil_image(image), brightness_factor=brightness_factor))
1944
1945
1946
1947

        torch.testing.assert_close(actual, expected)


1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
class TestCutMixMixUp:
    class DummyDataset:
        def __init__(self, size, num_classes):
            self.size = size
            self.num_classes = num_classes
            assert size < num_classes

        def __getitem__(self, idx):
            img = torch.rand(3, 100, 100)
            label = idx  # This ensures all labels in a batch are unique and makes testing easier
            return img, label

        def __len__(self):
            return self.size

Nicolas Hug's avatar
Nicolas Hug committed
1963
    @pytest.mark.parametrize("T", [transforms.CutMix, transforms.MixUp])
1964
1965
1966
1967
1968
1969
1970
    def test_supported_input_structure(self, T):

        batch_size = 32
        num_classes = 100

        dataset = self.DummyDataset(size=batch_size, num_classes=num_classes)

1971
        cutmix_mixup = T(num_classes=num_classes)
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
2007
2008
2009
2010
2011
2012

        dl = DataLoader(dataset, batch_size=batch_size)

        # Input sanity checks
        img, target = next(iter(dl))
        input_img_size = img.shape[-3:]
        assert isinstance(img, torch.Tensor) and isinstance(target, torch.Tensor)
        assert target.shape == (batch_size,)

        def check_output(img, target):
            assert img.shape == (batch_size, *input_img_size)
            assert target.shape == (batch_size, num_classes)
            torch.testing.assert_close(target.sum(axis=-1), torch.ones(batch_size))
            num_non_zero_labels = (target != 0).sum(axis=-1)
            assert (num_non_zero_labels == 2).all()

        # After Dataloader, as unpacked input
        img, target = next(iter(dl))
        assert target.shape == (batch_size,)
        img, target = cutmix_mixup(img, target)
        check_output(img, target)

        # After Dataloader, as packed input
        packed_from_dl = next(iter(dl))
        assert isinstance(packed_from_dl, list)
        img, target = cutmix_mixup(packed_from_dl)
        check_output(img, target)

        # As collation function. We expect default_collate to be used by users.
        def collate_fn_1(batch):
            return cutmix_mixup(default_collate(batch))

        def collate_fn_2(batch):
            return cutmix_mixup(*default_collate(batch))

        for collate_fn in (collate_fn_1, collate_fn_2):
            dl = DataLoader(dataset, batch_size=batch_size, collate_fn=collate_fn)
            img, target = next(iter(dl))
            check_output(img, target)

    @needs_cuda
Nicolas Hug's avatar
Nicolas Hug committed
2013
    @pytest.mark.parametrize("T", [transforms.CutMix, transforms.MixUp])
2014
2015
2016
2017
2018
2019
2020
2021
2022
2023
2024
    def test_cpu_vs_gpu(self, T):
        num_classes = 10
        batch_size = 3
        H, W = 12, 12

        imgs = torch.rand(batch_size, 3, H, W)
        labels = torch.randint(0, num_classes, (batch_size,))
        cutmix_mixup = T(alpha=0.5, num_classes=num_classes)

        _check_kernel_cuda_vs_cpu(cutmix_mixup, imgs, labels, rtol=None, atol=None)

Nicolas Hug's avatar
Nicolas Hug committed
2025
    @pytest.mark.parametrize("T", [transforms.CutMix, transforms.MixUp])
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
    def test_error(self, T):

        num_classes = 10
        batch_size = 9

        imgs = torch.rand(batch_size, 3, 12, 12)
        cutmix_mixup = T(alpha=0.5, num_classes=num_classes)

        for input_with_bad_type in (
            F.to_pil_image(imgs[0]),
2036
2037
            tv_tensors.Mask(torch.rand(12, 12)),
            tv_tensors.BoundingBoxes(torch.rand(2, 4), format="XYXY", canvas_size=12),
2038
2039
2040
2041
2042
2043
2044
2045
2046
2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
        ):
            with pytest.raises(ValueError, match="does not support PIL images, "):
                cutmix_mixup(input_with_bad_type)

        with pytest.raises(ValueError, match="Could not infer where the labels are"):
            cutmix_mixup({"img": imgs, "Nothing_else": 3})

        with pytest.raises(ValueError, match="labels tensor should be of shape"):
            # Note: the error message isn't ideal, but that's because the label heuristic found the img as the label
            # It's OK, it's an edge-case. The important thing is that this fails loudly instead of passing silently
            cutmix_mixup(imgs)

        with pytest.raises(ValueError, match="When using the default labels_getter"):
            cutmix_mixup(imgs, "not_a_tensor")

        with pytest.raises(ValueError, match="labels tensor should be of shape"):
            cutmix_mixup(imgs, torch.randint(0, 2, size=(2, 3)))

        with pytest.raises(ValueError, match="Expected a batched input with 4 dims"):
            cutmix_mixup(imgs[None, None], torch.randint(0, num_classes, size=(batch_size,)))

        with pytest.raises(ValueError, match="does not match the batch size of the labels"):
            cutmix_mixup(imgs, torch.randint(0, num_classes, size=(batch_size + 1,)))

        with pytest.raises(ValueError, match="labels tensor should be of shape"):
            # The purpose of this check is more about documenting the current
            # behaviour of what happens on a Compose(), rather than actually
            # asserting the expected behaviour. We may support Compose() in the
            # future, e.g. for 2 consecutive CutMix?
            labels = torch.randint(0, num_classes, size=(batch_size,))
            transforms.Compose([cutmix_mixup, cutmix_mixup])(imgs, labels)


@pytest.mark.parametrize("key", ("labels", "LABELS", "LaBeL", "SOME_WEIRD_KEY_THAT_HAS_LABeL_IN_IT"))
@pytest.mark.parametrize("sample_type", (tuple, list, dict))
def test_labels_getter_default_heuristic(key, sample_type):
    labels = torch.arange(10)
    sample = {key: labels, "another_key": "whatever"}
    if sample_type is not dict:
        sample = sample_type((None, sample, "whatever_again"))
    assert transforms._utils._find_labels_default_heuristic(sample) is labels

    if key.lower() != "labels":
        # If "labels" is in the dict (case-insensitive),
        # it takes precedence over other keys which would otherwise be a match
        d = {key: "something_else", "labels": labels}
        assert transforms._utils._find_labels_default_heuristic(d) is labels
2085
2086
2087
2088
2089
2090


class TestShapeGetters:
    @pytest.mark.parametrize(
        ("kernel", "make_input"),
        [
2091
2092
2093
            (F.get_dimensions_image, make_image_tensor),
            (F._get_dimensions_image_pil, make_image_pil),
            (F.get_dimensions_image, make_image),
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
            (F.get_dimensions_video, make_video),
        ],
    )
    def test_get_dimensions(self, kernel, make_input):
        size = (10, 10)
        color_space, num_channels = "RGB", 3

        input = make_input(size, color_space=color_space)

        assert kernel(input) == F.get_dimensions(input) == [num_channels, *size]

    @pytest.mark.parametrize(
        ("kernel", "make_input"),
        [
2108
2109
2110
            (F.get_num_channels_image, make_image_tensor),
            (F._get_num_channels_image_pil, make_image_pil),
            (F.get_num_channels_image, make_image),
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
            (F.get_num_channels_video, make_video),
        ],
    )
    def test_get_num_channels(self, kernel, make_input):
        color_space, num_channels = "RGB", 3

        input = make_input(color_space=color_space)

        assert kernel(input) == F.get_num_channels(input) == num_channels

    @pytest.mark.parametrize(
        ("kernel", "make_input"),
        [
2124
2125
2126
            (F.get_size_image, make_image_tensor),
            (F._get_size_image_pil, make_image_pil),
            (F.get_size_image, make_image),
2127
            (F.get_size_bounding_boxes, make_bounding_boxes),
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
            (F.get_size_mask, make_detection_mask),
            (F.get_size_mask, make_segmentation_mask),
            (F.get_size_video, make_video),
        ],
    )
    def test_get_size(self, kernel, make_input):
        size = (10, 10)

        input = make_input(size)

        assert kernel(input) == F.get_size(input) == list(size)

    @pytest.mark.parametrize(
        ("kernel", "make_input"),
        [
            (F.get_num_frames_video, make_video_tensor),
            (F.get_num_frames_video, make_video),
        ],
    )
    def test_get_num_frames(self, kernel, make_input):
        num_frames = 4

        input = make_input(num_frames=num_frames)

        assert kernel(input) == F.get_num_frames(input) == num_frames

    @pytest.mark.parametrize(
Nicolas Hug's avatar
Nicolas Hug committed
2155
        ("functional", "make_input"),
2156
        [
2157
            (F.get_dimensions, make_bounding_boxes),
2158
2159
            (F.get_dimensions, make_detection_mask),
            (F.get_dimensions, make_segmentation_mask),
2160
            (F.get_num_channels, make_bounding_boxes),
2161
2162
2163
2164
            (F.get_num_channels, make_detection_mask),
            (F.get_num_channels, make_segmentation_mask),
            (F.get_num_frames, make_image_pil),
            (F.get_num_frames, make_image),
2165
            (F.get_num_frames, make_bounding_boxes),
2166
2167
2168
2169
            (F.get_num_frames, make_detection_mask),
            (F.get_num_frames, make_segmentation_mask),
        ],
    )
Nicolas Hug's avatar
Nicolas Hug committed
2170
    def test_unsupported_types(self, functional, make_input):
2171
2172
2173
        input = make_input()

        with pytest.raises(TypeError, match=re.escape(str(type(input)))):
Nicolas Hug's avatar
Nicolas Hug committed
2174
            functional(input)
2175
2176
2177


class TestRegisterKernel:
Nicolas Hug's avatar
Nicolas Hug committed
2178
2179
    @pytest.mark.parametrize("functional", (F.resize, "resize"))
    def test_register_kernel(self, functional):
2180
        class CustomTVTensor(tv_tensors.TVTensor):
2181
2182
2183
2184
            pass

        kernel_was_called = False

2185
        @F.register_kernel(functional, CustomTVTensor)
2186
2187
2188
2189
2190
2191
2192
        def new_resize(dp, *args, **kwargs):
            nonlocal kernel_was_called
            kernel_was_called = True
            return dp

        t = transforms.Resize(size=(224, 224), antialias=True)

2193
        my_dp = CustomTVTensor(torch.rand(3, 10, 10))
2194
2195
2196
2197
2198
2199
        out = t(my_dp)
        assert out is my_dp
        assert kernel_was_called

        # Sanity check to make sure we didn't override the kernel of other types
        t(torch.rand(3, 10, 10)).shape == (3, 224, 224)
2200
        t(tv_tensors.Image(torch.rand(3, 10, 10))).shape == (3, 224, 224)
2201

2202
    def test_errors(self):
Nicolas Hug's avatar
Nicolas Hug committed
2203
        with pytest.raises(ValueError, match="Could not find functional with name"):
2204
            F.register_kernel("bad_name", tv_tensors.Image)
2205

Nicolas Hug's avatar
Nicolas Hug committed
2206
        with pytest.raises(ValueError, match="Kernels can only be registered on functionals"):
2207
            F.register_kernel(tv_tensors.Image, F.resize)
2208
2209
2210
2211

        with pytest.raises(ValueError, match="Kernels can only be registered for subclasses"):
            F.register_kernel(F.resize, object)

2212
2213
        with pytest.raises(ValueError, match="cannot be registered for the builtin tv_tensor classes"):
            F.register_kernel(F.resize, tv_tensors.Image)(F.resize_image)
2214

2215
        class CustomTVTensor(tv_tensors.TVTensor):
2216
2217
            pass

2218
        def resize_custom_tv_tensor():
2219
2220
            pass

2221
        F.register_kernel(F.resize, CustomTVTensor)(resize_custom_tv_tensor)
2222
2223

        with pytest.raises(ValueError, match="already has a kernel registered for type"):
2224
            F.register_kernel(F.resize, CustomTVTensor)(resize_custom_tv_tensor)
2225

2226
2227

class TestGetKernel:
Nicolas Hug's avatar
Nicolas Hug committed
2228
    # We are using F.resize as functional and the kernels below as proxy. Any other functional / kernels combination
2229
2230
    # would also be fine
    KERNELS = {
2231
2232
        torch.Tensor: F.resize_image,
        PIL.Image.Image: F._resize_image_pil,
2233
2234
2235
2236
        tv_tensors.Image: F.resize_image,
        tv_tensors.BoundingBoxes: F.resize_bounding_boxes,
        tv_tensors.Mask: F.resize_mask,
        tv_tensors.Video: F.resize_video,
2237
2238
    }

2239
2240
2241
2242
    @pytest.mark.parametrize("input_type", [str, int, object])
    def test_unsupported_types(self, input_type):
        with pytest.raises(TypeError, match="supports inputs of type"):
            _get_kernel(F.resize, input_type)
2243
2244
2245

    def test_exact_match(self):
        # We cannot use F.resize together with self.KERNELS mapping here directly here, since this is only the
Nicolas Hug's avatar
Nicolas Hug committed
2246
        # ideal wrapping. Practically, we have an intermediate wrapper layer. Thus, we create a new resize functional
2247
2248
2249
2250
2251
        # here, register the kernels without wrapper, and check the exact matching afterwards.
        def resize_with_pure_kernels():
            pass

        for input_type, kernel in self.KERNELS.items():
2252
            _register_kernel_internal(resize_with_pure_kernels, input_type, tv_tensor_wrapper=False)(kernel)
2253
2254
2255

            assert _get_kernel(resize_with_pure_kernels, input_type) is kernel

2256
    def test_builtin_tv_tensor_subclass(self):
2257
        # We cannot use F.resize together with self.KERNELS mapping here directly here, since this is only the
Nicolas Hug's avatar
Nicolas Hug committed
2258
        # ideal wrapping. Practically, we have an intermediate wrapper layer. Thus, we create a new resize functional
2259
        # here, register the kernels without wrapper, and check if subclasses of our builtin tv_tensors get dispatched
2260
2261
2262
2263
        # to the kernel of the corresponding superclass
        def resize_with_pure_kernels():
            pass

2264
        class MyImage(tv_tensors.Image):
2265
2266
            pass

2267
        class MyBoundingBoxes(tv_tensors.BoundingBoxes):
2268
2269
            pass

2270
        class MyMask(tv_tensors.Mask):
2271
2272
            pass

2273
        class MyVideo(tv_tensors.Video):
2274
2275
            pass

2276
        for custom_tv_tensor_subclass in [
2277
2278
2279
2280
2281
            MyImage,
            MyBoundingBoxes,
            MyMask,
            MyVideo,
        ]:
2282
2283
2284
2285
            builtin_tv_tensor_class = custom_tv_tensor_subclass.__mro__[1]
            builtin_tv_tensor_kernel = self.KERNELS[builtin_tv_tensor_class]
            _register_kernel_internal(resize_with_pure_kernels, builtin_tv_tensor_class, tv_tensor_wrapper=False)(
                builtin_tv_tensor_kernel
2286
2287
            )

2288
            assert _get_kernel(resize_with_pure_kernels, custom_tv_tensor_subclass) is builtin_tv_tensor_kernel
2289

2290
2291
    def test_tv_tensor_subclass(self):
        class MyTVTensor(tv_tensors.TVTensor):
2292
2293
            pass

2294
        with pytest.raises(TypeError, match="supports inputs of type"):
2295
            _get_kernel(F.resize, MyTVTensor)
2296

2297
        def resize_my_tv_tensor():
2298
2299
            pass

2300
        _register_kernel_internal(F.resize, MyTVTensor, tv_tensor_wrapper=False)(resize_my_tv_tensor)
2301

2302
        assert _get_kernel(F.resize, MyTVTensor) is resize_my_tv_tensor
2303

2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
    def test_pil_image_subclass(self):
        opened_image = PIL.Image.open(Path(__file__).parent / "assets" / "encode_jpeg" / "grace_hopper_517x606.jpg")
        loaded_image = opened_image.convert("RGB")

        # check the assumptions
        assert isinstance(opened_image, PIL.Image.Image)
        assert type(opened_image) is not PIL.Image.Image

        assert type(loaded_image) is PIL.Image.Image

        size = [17, 11]
        for image in [opened_image, loaded_image]:
            kernel = _get_kernel(F.resize, type(image))

            output = kernel(image, size=size)

            assert F.get_size(output) == size

2322
2323
2324
2325
2326
2327
2328

class TestPermuteChannels:
    _DEFAULT_PERMUTATION = [2, 0, 1]

    @pytest.mark.parametrize(
        ("kernel", "make_input"),
        [
2329
            (F.permute_channels_image, make_image_tensor),
2330
2331
            # FIXME
            # check_kernel does not support PIL kernel, but it should
2332
            (F.permute_channels_image, make_image),
2333
2334
2335
2336
2337
2338
2339
2340
            (F.permute_channels_video, make_video),
        ],
    )
    @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8])
    @pytest.mark.parametrize("device", cpu_and_cuda())
    def test_kernel(self, kernel, make_input, dtype, device):
        check_kernel(kernel, make_input(dtype=dtype, device=device), permutation=self._DEFAULT_PERMUTATION)

Nicolas Hug's avatar
Nicolas Hug committed
2341
    @pytest.mark.parametrize("make_input", [make_image_tensor, make_image_pil, make_image, make_video])
Nicolas Hug's avatar
Nicolas Hug committed
2342
2343
    def test_functional(self, make_input):
        check_functional(F.permute_channels, make_input(), permutation=self._DEFAULT_PERMUTATION)
2344
2345
2346
2347

    @pytest.mark.parametrize(
        ("kernel", "input_type"),
        [
2348
2349
            (F.permute_channels_image, torch.Tensor),
            (F._permute_channels_image_pil, PIL.Image.Image),
2350
2351
            (F.permute_channels_image, tv_tensors.Image),
            (F.permute_channels_video, tv_tensors.Video),
2352
2353
        ],
    )
Nicolas Hug's avatar
Nicolas Hug committed
2354
2355
    def test_functional_signature(self, kernel, input_type):
        check_functional_kernel_signature_match(F.permute_channels, kernel=kernel, input_type=input_type)
2356
2357
2358
2359

    def reference_image_correctness(self, image, permutation):
        channel_images = image.split(1, dim=-3)
        permuted_channel_images = [channel_images[channel_idx] for channel_idx in permutation]
2360
        return tv_tensors.Image(torch.concat(permuted_channel_images, dim=-3))
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370

    @pytest.mark.parametrize("permutation", [[2, 0, 1], [1, 2, 0], [2, 0, 1], [0, 1, 2]])
    @pytest.mark.parametrize("batch_dims", [(), (2,), (2, 1)])
    def test_image_correctness(self, permutation, batch_dims):
        image = make_image(batch_dims=batch_dims)

        actual = F.permute_channels(image, permutation=permutation)
        expected = self.reference_image_correctness(image, permutation=permutation)

        torch.testing.assert_close(actual, expected)
Philip Meier's avatar
Philip Meier committed
2371
2372
2373
2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387
2388


class TestElastic:
    def _make_displacement(self, inpt):
        return torch.rand(
            1,
            *F.get_size(inpt),
            2,
            dtype=torch.float32,
            device=inpt.device if isinstance(inpt, torch.Tensor) else "cpu",
        )

    @param_value_parametrization(
        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())
2389
    def test_kernel_image(self, param, value, dtype, device):
Philip Meier's avatar
Philip Meier committed
2390
2391
2392
        image = make_image_tensor(dtype=dtype, device=device)

        check_kernel(
Philip Meier's avatar
Philip Meier committed
2393
            F.elastic_image,
Philip Meier's avatar
Philip Meier committed
2394
2395
2396
2397
2398
2399
            image,
            displacement=self._make_displacement(image),
            **{param: value},
            check_scripted_vs_eager=not (param == "fill" and isinstance(value, (int, float))),
        )

2400
    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
Philip Meier's avatar
Philip Meier committed
2401
2402
2403
    @pytest.mark.parametrize("dtype", [torch.float32, torch.int64])
    @pytest.mark.parametrize("device", cpu_and_cuda())
    def test_kernel_bounding_boxes(self, format, dtype, device):
2404
        bounding_boxes = make_bounding_boxes(format=format, dtype=dtype, device=device)
Philip Meier's avatar
Philip Meier committed
2405
2406
2407
2408
2409
2410
2411
2412
2413
2414
2415
2416
2417
2418
2419
2420
2421
2422
2423
2424

        check_kernel(
            F.elastic_bounding_boxes,
            bounding_boxes,
            format=bounding_boxes.format,
            canvas_size=bounding_boxes.canvas_size,
            displacement=self._make_displacement(bounding_boxes),
        )

    @pytest.mark.parametrize("make_mask", [make_segmentation_mask, make_detection_mask])
    def test_kernel_mask(self, make_mask):
        mask = make_mask()
        check_kernel(F.elastic_mask, mask, displacement=self._make_displacement(mask))

    def test_kernel_video(self):
        video = make_video()
        check_kernel(F.elastic_video, video, displacement=self._make_displacement(video))

    @pytest.mark.parametrize(
        "make_input",
2425
        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
Philip Meier's avatar
Philip Meier committed
2426
2427
2428
2429
2430
2431
2432
2433
    )
    def test_functional(self, make_input):
        input = make_input()
        check_functional(F.elastic, input, displacement=self._make_displacement(input))

    @pytest.mark.parametrize(
        ("kernel", "input_type"),
        [
Philip Meier's avatar
Philip Meier committed
2434
2435
            (F.elastic_image, torch.Tensor),
            (F._elastic_image_pil, PIL.Image.Image),
2436
2437
2438
2439
            (F.elastic_image, tv_tensors.Image),
            (F.elastic_bounding_boxes, tv_tensors.BoundingBoxes),
            (F.elastic_mask, tv_tensors.Mask),
            (F.elastic_video, tv_tensors.Video),
Philip Meier's avatar
Philip Meier committed
2440
2441
2442
2443
2444
2445
2446
        ],
    )
    def test_functional_signature(self, kernel, input_type):
        check_functional_kernel_signature_match(F.elastic, kernel=kernel, input_type=input_type)

    @pytest.mark.parametrize(
        "make_input",
2447
        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
Philip Meier's avatar
Philip Meier committed
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
    )
    def test_displacement_error(self, make_input):
        input = make_input()

        with pytest.raises(TypeError, match="displacement should be a Tensor"):
            F.elastic(input, displacement=None)

        with pytest.raises(ValueError, match="displacement shape should be"):
            F.elastic(input, displacement=torch.rand(F.get_size(input)))

    @pytest.mark.parametrize(
        "make_input",
2460
        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
Philip Meier's avatar
Philip Meier committed
2461
2462
2463
2464
2465
    )
    # ElasticTransform needs larger images to avoid the needed internal padding being larger than the actual image
    @pytest.mark.parametrize("size", [(163, 163), (72, 333), (313, 95)])
    @pytest.mark.parametrize("device", cpu_and_cuda())
    def test_transform(self, make_input, size, device):
2466
2467
2468
2469
2470
2471
        check_transform(
            transforms.ElasticTransform(),
            make_input(size, device=device),
            # We updated gaussian blur kernel generation with a faster and numerically more stable version
            check_v1_compatibility=dict(rtol=0, atol=1),
        )
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481


class TestToPureTensor:
    def test_correctness(self):
        input = {
            "img": make_image(),
            "img_tensor": make_image_tensor(),
            "img_pil": make_image_pil(),
            "mask": make_detection_mask(),
            "video": make_video(),
2482
            "bbox": make_bounding_boxes(),
2483
2484
2485
2486
2487
2488
            "str": "str",
        }

        out = transforms.ToPureTensor()(input)

        for input_value, out_value in zip(input.values(), out.values()):
2489
2490
            if isinstance(input_value, tv_tensors.TVTensor):
                assert isinstance(out_value, torch.Tensor) and not isinstance(out_value, tv_tensors.TVTensor)
2491
2492
            else:
                assert isinstance(out_value, type(input_value))
2493
2494
2495
2496
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522
2523
2524
2525
2526
2527
2528
2529
2530
2531
2532
2533
2534
2535
2536
2537
2538
2539
2540
2541
2542
2543
2544
2545
2546
2547
2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
2558
2559
2560
2561
2562
2563
2564
2565
2566
2567
2568
2569
2570
2571
2572
2573
2574
2575
2576
2577
2578
2579
2580
2581
2582
2583
2584
2585
2586
2587
2588
2589
2590
2591
2592
2593
2594
2595
2596
2597
2598
2599
2600
2601
2602
2603
2604
2605
2606
2607
2608
2609
2610
2611
2612
2613
2614
2615
2616
2617
2618
2619
2620
2621
2622
2623
2624
2625
2626
2627
2628
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
2644
2645
2646
2647
2648
2649
2650
2651
2652
2653
2654
2655
2656
2657
2658
2659
2660
2661
2662
2663
2664
2665
2666
2667
2668
2669
2670
2671
2672
2673
2674
2675
2676
2677
2678
2679
2680
2681
2682
2683
2684
2685
2686
2687
2688
2689
2690
2691
2692
2693
2694
2695
2696
2697
2698
2699
2700
2701
2702
2703
2704
2705
2706
2707
2708
2709
2710
2711
2712
2713
2714
2715
2716


class TestCrop:
    INPUT_SIZE = (21, 11)

    CORRECTNESS_CROP_KWARGS = [
        # center
        dict(top=5, left=5, height=10, width=5),
        # larger than input, i.e. pad
        dict(top=-5, left=-5, height=30, width=20),
        # sides: left, right, top, bottom
        dict(top=-5, left=-5, height=30, width=10),
        dict(top=-5, left=5, height=30, width=10),
        dict(top=-5, left=-5, height=20, width=20),
        dict(top=5, left=-5, height=20, width=20),
        # corners: top-left, top-right, bottom-left, bottom-right
        dict(top=-5, left=-5, height=20, width=10),
        dict(top=-5, left=5, height=20, width=10),
        dict(top=5, left=-5, height=20, width=10),
        dict(top=5, left=5, height=20, width=10),
    ]
    MINIMAL_CROP_KWARGS = CORRECTNESS_CROP_KWARGS[0]

    @pytest.mark.parametrize("kwargs", CORRECTNESS_CROP_KWARGS)
    @pytest.mark.parametrize("dtype", [torch.uint8, torch.float32])
    @pytest.mark.parametrize("device", cpu_and_cuda())
    def test_kernel_image(self, kwargs, dtype, device):
        check_kernel(F.crop_image, make_image(self.INPUT_SIZE, dtype=dtype, device=device), **kwargs)

    @pytest.mark.parametrize("kwargs", CORRECTNESS_CROP_KWARGS)
    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
    @pytest.mark.parametrize("dtype", [torch.float32, torch.int64])
    @pytest.mark.parametrize("device", cpu_and_cuda())
    def test_kernel_bounding_box(self, kwargs, format, dtype, device):
        bounding_boxes = make_bounding_boxes(self.INPUT_SIZE, format=format, dtype=dtype, device=device)
        check_kernel(F.crop_bounding_boxes, bounding_boxes, format=format, **kwargs)

    @pytest.mark.parametrize("make_mask", [make_segmentation_mask, make_detection_mask])
    def test_kernel_mask(self, make_mask):
        check_kernel(F.crop_mask, make_mask(self.INPUT_SIZE), **self.MINIMAL_CROP_KWARGS)

    def test_kernel_video(self):
        check_kernel(F.crop_video, make_video(self.INPUT_SIZE), **self.MINIMAL_CROP_KWARGS)

    @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.crop, make_input(self.INPUT_SIZE), **self.MINIMAL_CROP_KWARGS)

    @pytest.mark.parametrize(
        ("kernel", "input_type"),
        [
            (F.crop_image, torch.Tensor),
            (F._crop_image_pil, PIL.Image.Image),
            (F.crop_image, tv_tensors.Image),
            (F.crop_bounding_boxes, tv_tensors.BoundingBoxes),
            (F.crop_mask, tv_tensors.Mask),
            (F.crop_video, tv_tensors.Video),
        ],
    )
    def test_functional_signature(self, kernel, input_type):
        check_functional_kernel_signature_match(F.crop, kernel=kernel, input_type=input_type)

    @pytest.mark.parametrize("kwargs", CORRECTNESS_CROP_KWARGS)
    def test_functional_image_correctness(self, kwargs):
        image = make_image(self.INPUT_SIZE, dtype=torch.uint8, device="cpu")

        actual = F.crop(image, **kwargs)
        expected = F.to_image(F.crop(F.to_pil_image(image), **kwargs))

        assert_equal(actual, expected)

    @param_value_parametrization(
        size=[(10, 5), (25, 15), (25, 5), (10, 15)],
        fill=EXHAUSTIVE_TYPE_FILLS,
    )
    @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, param, value, make_input):
        input = make_input(self.INPUT_SIZE)

        kwargs = {param: value}
        if param == "fill":
            # 1. size is required
            # 2. the fill parameter only has an affect if we need padding
            kwargs["size"] = [s + 4 for s in self.INPUT_SIZE]

            if isinstance(input, tv_tensors.Mask) and isinstance(value, (tuple, list)):
                pytest.skip("F.pad_mask doesn't support non-scalar fill.")

        check_transform(
            transforms.RandomCrop(**kwargs, pad_if_needed=True),
            input,
            check_v1_compatibility=param != "fill" or isinstance(value, (int, float)),
        )

    @pytest.mark.parametrize("padding", [1, (1, 1), (1, 1, 1, 1)])
    def test_transform_padding(self, padding):
        inpt = make_image(self.INPUT_SIZE)

        output_size = [s + 2 for s in F.get_size(inpt)]
        transform = transforms.RandomCrop(output_size, padding=padding)

        output = transform(inpt)

        assert F.get_size(output) == output_size

    @pytest.mark.parametrize("padding", [None, 1, (1, 1), (1, 1, 1, 1)])
    def test_transform_insufficient_padding(self, padding):
        inpt = make_image(self.INPUT_SIZE)

        output_size = [s + 3 for s in F.get_size(inpt)]
        transform = transforms.RandomCrop(output_size, padding=padding)

        with pytest.raises(ValueError, match="larger than (padded )?input image size"):
            transform(inpt)

    def test_transform_pad_if_needed(self):
        inpt = make_image(self.INPUT_SIZE)

        output_size = [s * 2 for s in F.get_size(inpt)]
        transform = transforms.RandomCrop(output_size, pad_if_needed=True)

        output = transform(inpt)

        assert F.get_size(output) == output_size

    @param_value_parametrization(
        size=[(10, 5), (25, 15), (25, 5), (10, 15)],
        fill=CORRECTNESS_FILLS,
        padding_mode=["constant", "edge", "reflect", "symmetric"],
    )
    @pytest.mark.parametrize("seed", list(range(5)))
    def test_transform_image_correctness(self, param, value, seed):
        kwargs = {param: value}
        if param != "size":
            # 1. size is required
            # 2. the fill / padding_mode parameters only have an affect if we need padding
            kwargs["size"] = [s + 4 for s in self.INPUT_SIZE]
        if param == "fill":
            kwargs["fill"] = adapt_fill(kwargs["fill"], dtype=torch.uint8)

        transform = transforms.RandomCrop(pad_if_needed=True, **kwargs)

        image = make_image(self.INPUT_SIZE)

        with freeze_rng_state():
            torch.manual_seed(seed)
            actual = transform(image)

            torch.manual_seed(seed)
            expected = F.to_image(transform(F.to_pil_image(image)))

        assert_equal(actual, expected)

    def _reference_crop_bounding_boxes(self, bounding_boxes, *, top, left, height, width):
        affine_matrix = np.array(
            [
                [1, 0, -left],
                [0, 1, -top],
            ],
        )
        return reference_affine_bounding_boxes_helper(
            bounding_boxes, affine_matrix=affine_matrix, new_canvas_size=(height, width)
        )

    @pytest.mark.parametrize("kwargs", CORRECTNESS_CROP_KWARGS)
    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
    @pytest.mark.parametrize("dtype", [torch.float32, torch.int64])
    @pytest.mark.parametrize("device", cpu_and_cuda())
    def test_functional_bounding_box_correctness(self, kwargs, format, dtype, device):
        bounding_boxes = make_bounding_boxes(self.INPUT_SIZE, format=format, dtype=dtype, device=device)

        actual = F.crop(bounding_boxes, **kwargs)
        expected = self._reference_crop_bounding_boxes(bounding_boxes, **kwargs)

        assert_equal(actual, expected, atol=1, rtol=0)
        assert_equal(F.get_size(actual), F.get_size(expected))

    @pytest.mark.parametrize("output_size", [(17, 11), (11, 17), (11, 11)])
    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
    @pytest.mark.parametrize("dtype", [torch.float32, torch.int64])
    @pytest.mark.parametrize("device", cpu_and_cuda())
    @pytest.mark.parametrize("seed", list(range(5)))
    def test_transform_bounding_boxes_correctness(self, output_size, format, dtype, device, seed):
        input_size = [s * 2 for s in output_size]
        bounding_boxes = make_bounding_boxes(input_size, format=format, dtype=dtype, device=device)

        transform = transforms.RandomCrop(output_size)

        with freeze_rng_state():
            torch.manual_seed(seed)
            params = transform._get_params([bounding_boxes])
            assert not params.pop("needs_pad")
            del params["padding"]
            assert params.pop("needs_crop")

            torch.manual_seed(seed)
            actual = transform(bounding_boxes)

        expected = self._reference_crop_bounding_boxes(bounding_boxes, **params)

        assert_equal(actual, expected)
        assert_equal(F.get_size(actual), F.get_size(expected))

    def test_errors(self):
        with pytest.raises(ValueError, match="Please provide only two dimensions"):
            transforms.RandomCrop([10, 12, 14])

        with pytest.raises(TypeError, match="Got inappropriate padding arg"):
            transforms.RandomCrop([10, 12], padding="abc")

        with pytest.raises(ValueError, match="Padding must be an int or a 1, 2, or 4"):
            transforms.RandomCrop([10, 12], padding=[-0.7, 0, 0.7])

        with pytest.raises(TypeError, match="Got inappropriate fill arg"):
            transforms.RandomCrop([10, 12], padding=1, fill="abc")

        with pytest.raises(ValueError, match="Padding mode should be either"):
            transforms.RandomCrop([10, 12], padding=1, padding_mode="abc")
2717
2718
2719
2720
2721
2722
2723
2724
2725
2726
2727
2728
2729
2730
2731
2732
2733
2734
2735
2736
2737
2738
2739
2740
2741
2742
2743
2744
2745
2746
2747
2748
2749
2750
2751
2752
2753
2754
2755
2756
2757
2758
2759
2760
2761
2762
2763
2764
2765
2766
2767
2768
2769
2770
2771
2772
2773
2774
2775
2776
2777
2778
2779
2780
2781
2782
2783
2784
2785
2786
2787
2788
2789
2790
2791
2792
2793
2794
2795
2796
2797
2798
2799
2800
2801
2802
2803
2804
2805
2806
2807
2808
2809
2810
2811
2812
2813
2814
2815
2816
2817
2818
2819
2820
2821
2822
2823
2824
2825
2826
2827
2828
2829
2830
2831
2832
2833
2834
2835
2836
2837
2838
2839
2840
2841
2842
2843
2844
2845
2846
2847
2848
2849
2850
2851
2852
2853
2854
2855
2856
2857


class TestErase:
    INPUT_SIZE = (17, 11)
    FUNCTIONAL_KWARGS = dict(
        zip("ijhwv", [2, 2, 10, 8, torch.tensor(0.0, dtype=torch.float32, device="cpu").reshape(-1, 1, 1)])
    )

    @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8])
    @pytest.mark.parametrize("device", cpu_and_cuda())
    def test_kernel_image(self, dtype, device):
        check_kernel(F.erase_image, make_image(self.INPUT_SIZE, dtype=dtype, device=device), **self.FUNCTIONAL_KWARGS)

    @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8])
    @pytest.mark.parametrize("device", cpu_and_cuda())
    def test_kernel_image_inplace(self, dtype, device):
        input = make_image(self.INPUT_SIZE, dtype=dtype, device=device)
        input_version = input._version

        output_out_of_place = F.erase_image(input, **self.FUNCTIONAL_KWARGS)
        assert output_out_of_place.data_ptr() != input.data_ptr()
        assert output_out_of_place is not input

        output_inplace = F.erase_image(input, **self.FUNCTIONAL_KWARGS, inplace=True)
        assert output_inplace.data_ptr() == input.data_ptr()
        assert output_inplace._version > input_version
        assert output_inplace is input

        assert_equal(output_inplace, output_out_of_place)

    def test_kernel_video(self):
        check_kernel(F.erase_video, make_video(self.INPUT_SIZE), **self.FUNCTIONAL_KWARGS)

    @pytest.mark.parametrize(
        "make_input",
        [make_image_tensor, make_image_pil, make_image, make_video],
    )
    def test_functional(self, make_input):
        check_functional(F.erase, make_input(), **self.FUNCTIONAL_KWARGS)

    @pytest.mark.parametrize(
        ("kernel", "input_type"),
        [
            (F.erase_image, torch.Tensor),
            (F._erase_image_pil, PIL.Image.Image),
            (F.erase_image, tv_tensors.Image),
            (F.erase_video, tv_tensors.Video),
        ],
    )
    def test_functional_signature(self, kernel, input_type):
        check_functional_kernel_signature_match(F.erase, kernel=kernel, input_type=input_type)

    @pytest.mark.parametrize(
        "make_input",
        [make_image_tensor, make_image_pil, make_image, make_video],
    )
    @pytest.mark.parametrize("device", cpu_and_cuda())
    def test_transform(self, make_input, device):
        check_transform(transforms.RandomErasing(p=1), make_input(device=device))

    def _reference_erase_image(self, image, *, i, j, h, w, v):
        mask = torch.zeros_like(image, dtype=torch.bool)
        mask[..., i : i + h, j : j + w] = True

        # The broadcasting and type casting logic is handled automagically in the kernel through indexing
        value = torch.broadcast_to(v, (*image.shape[:-2], h, w)).to(image)

        erased_image = torch.empty_like(image)
        erased_image[mask] = value.flatten()
        erased_image[~mask] = image[~mask]

        return erased_image

    @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8])
    @pytest.mark.parametrize("device", cpu_and_cuda())
    def test_functional_image_correctness(self, dtype, device):
        image = make_image(dtype=dtype, device=device)

        actual = F.erase(image, **self.FUNCTIONAL_KWARGS)
        expected = self._reference_erase_image(image, **self.FUNCTIONAL_KWARGS)

        assert_equal(actual, expected)

    @param_value_parametrization(
        scale=[(0.1, 0.2), [0.0, 1.0]],
        ratio=[(0.3, 0.7), [0.1, 5.0]],
        value=[0, 0.5, (0, 1, 0), [-0.2, 0.0, 1.3], "random"],
    )
    @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8])
    @pytest.mark.parametrize("device", cpu_and_cuda())
    @pytest.mark.parametrize("seed", list(range(5)))
    def test_transform_image_correctness(self, param, value, dtype, device, seed):
        transform = transforms.RandomErasing(**{param: value}, p=1)

        image = make_image(dtype=dtype, device=device)

        with freeze_rng_state():
            torch.manual_seed(seed)
            # This emulates the random apply check that happens before _get_params is called
            torch.rand(1)
            params = transform._get_params([image])

            torch.manual_seed(seed)
            actual = transform(image)

        expected = self._reference_erase_image(image, **params)

        assert_equal(actual, expected)

    def test_transform_errors(self):
        with pytest.raises(TypeError, match="Argument value should be either a number or str or a sequence"):
            transforms.RandomErasing(value={})

        with pytest.raises(ValueError, match="If value is str, it should be 'random'"):
            transforms.RandomErasing(value="abc")

        with pytest.raises(TypeError, match="Scale should be a sequence"):
            transforms.RandomErasing(scale=123)

        with pytest.raises(TypeError, match="Ratio should be a sequence"):
            transforms.RandomErasing(ratio=123)

        with pytest.raises(ValueError, match="Scale should be between 0 and 1"):
            transforms.RandomErasing(scale=[-1, 2])

        transform = transforms.RandomErasing(value=[1, 2, 3, 4])

        with pytest.raises(ValueError, match="If value is a sequence, it should have either a single value"):
            transform._get_params([make_image()])

    @pytest.mark.parametrize("make_input", [make_bounding_boxes, make_detection_mask])
    def test_transform_passthrough(self, make_input):
        transform = transforms.RandomErasing(p=1)

        input = make_input(self.INPUT_SIZE)

        with pytest.warns(UserWarning, match="currently passing through inputs of type"):
            # RandomErasing requires an image or video to be present
            _, output = transform(make_image(self.INPUT_SIZE), input)

        assert output is input
2858
2859
2860
2861
2862
2863
2864
2865
2866
2867
2868
2869
2870
2871
2872
2873
2874
2875
2876
2877
2878
2879
2880
2881
2882
2883
2884
2885
2886
2887
2888
2889
2890
2891
2892
2893
2894
2895
2896
2897
2898
2899
2900


class TestGaussianBlur:
    @pytest.mark.parametrize(
        "make_input",
        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
    )
    @pytest.mark.parametrize("device", cpu_and_cuda())
    @pytest.mark.parametrize("sigma", [5, (0.5, 2)])
    def test_transform(self, make_input, device, sigma):
        check_transform(transforms.GaussianBlur(kernel_size=3, sigma=sigma), make_input(device=device))

    def test_assertions(self):
        with pytest.raises(ValueError, match="Kernel size should be a tuple/list of two integers"):
            transforms.GaussianBlur([10, 12, 14])

        with pytest.raises(ValueError, match="Kernel size value should be an odd and positive number"):
            transforms.GaussianBlur(4)

        with pytest.raises(ValueError, match="If sigma is a sequence its length should be 1 or 2. Got 3"):
            transforms.GaussianBlur(3, sigma=[1, 2, 3])

        with pytest.raises(ValueError, match="sigma values should be positive and of the form"):
            transforms.GaussianBlur(3, sigma=-1.0)

        with pytest.raises(ValueError, match="sigma values should be positive and of the form"):
            transforms.GaussianBlur(3, sigma=[2.0, 1.0])

        with pytest.raises(TypeError, match="sigma should be a number or a sequence of numbers"):
            transforms.GaussianBlur(3, sigma={})

    @pytest.mark.parametrize("sigma", [10.0, [10.0, 12.0], (10, 12.0), [10]])
    def test__get_params(self, sigma):
        transform = transforms.GaussianBlur(3, sigma=sigma)
        params = transform._get_params([])

        if isinstance(sigma, float):
            assert params["sigma"][0] == params["sigma"][1] == sigma
        elif isinstance(sigma, list) and len(sigma) == 1:
            assert params["sigma"][0] == params["sigma"][1] == sigma[0]
        else:
            assert sigma[0] <= params["sigma"][0] <= sigma[1]
            assert sigma[0] <= params["sigma"][1] <= sigma[1]