test_transforms_v2_consistency.py 45.5 KB
Newer Older
1
2
import importlib.machinery
import importlib.util
3
import inspect
4
import random
5
import re
6
from pathlib import Path
7

8
import numpy as np
9
import PIL.Image
10
import pytest
11
12

import torch
13
import torchvision.transforms.v2 as v2_transforms
14
from common_utils import assert_close, assert_equal, set_rng_seed
15
from torch import nn
16
from torchvision import transforms as legacy_transforms, tv_tensors
17
from torchvision._utils import sequence_to_str
18

19
from torchvision.transforms import functional as legacy_F
20
from torchvision.transforms.v2 import functional as prototype_F
Nicolas Hug's avatar
Nicolas Hug committed
21
from torchvision.transforms.v2._utils import _get_fill, query_size
22
from torchvision.transforms.v2.functional import to_pil_image
23
24
25
26
27
28
29
30
from transforms_v2_legacy_utils import (
    ArgsKwargs,
    make_bounding_boxes,
    make_detection_mask,
    make_image,
    make_images,
    make_segmentation_mask,
)
31

32
DEFAULT_MAKE_IMAGES_KWARGS = dict(color_spaces=["RGB"], extra_dims=[(4,)])
33
34


Nicolas Hug's avatar
Nicolas Hug committed
35
36
37
38
39
40
@pytest.fixture(autouse=True)
def fix_rng_seed():
    set_rng_seed(0)
    yield


41
42
43
44
45
46
47
48
49
class NotScriptableArgsKwargs(ArgsKwargs):
    """
    This class is used to mark parameters that render the transform non-scriptable. They still work in eager mode and
    thus will be tested there, but will be skipped by the JIT tests.
    """

    pass


50
51
class ConsistencyConfig:
    def __init__(
52
53
54
        self,
        prototype_cls,
        legacy_cls,
55
56
        # If no args_kwargs is passed, only the signature will be checked
        args_kwargs=(),
57
58
59
        make_images_kwargs=None,
        supports_pil=True,
        removed_params=(),
60
        closeness_kwargs=None,
61
62
63
    ):
        self.prototype_cls = prototype_cls
        self.legacy_cls = legacy_cls
64
        self.args_kwargs = args_kwargs
65
66
        self.make_images_kwargs = make_images_kwargs or DEFAULT_MAKE_IMAGES_KWARGS
        self.supports_pil = supports_pil
67
        self.removed_params = removed_params
68
        self.closeness_kwargs = closeness_kwargs or dict(rtol=0, atol=0)
69
70


71
72
73
74
# These are here since both the prototype and legacy transform need to be constructed with the same random parameters
LINEAR_TRANSFORMATION_MEAN = torch.rand(36)
LINEAR_TRANSFORMATION_MATRIX = torch.rand([LINEAR_TRANSFORMATION_MEAN.numel()] * 2)

75
76
CONSISTENCY_CONFIGS = [
    ConsistencyConfig(
77
        v2_transforms.Normalize,
78
79
80
81
82
83
84
85
        legacy_transforms.Normalize,
        [
            ArgsKwargs(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
        ],
        supports_pil=False,
        make_images_kwargs=dict(DEFAULT_MAKE_IMAGES_KWARGS, dtypes=[torch.float]),
    ),
    ConsistencyConfig(
86
        v2_transforms.CenterCrop,
87
88
89
90
91
92
        legacy_transforms.CenterCrop,
        [
            ArgsKwargs(18),
            ArgsKwargs((18, 13)),
        ],
    ),
93
    ConsistencyConfig(
94
        v2_transforms.FiveCrop,
95
96
97
98
99
100
101
102
        legacy_transforms.FiveCrop,
        [
            ArgsKwargs(18),
            ArgsKwargs((18, 13)),
        ],
        make_images_kwargs=dict(DEFAULT_MAKE_IMAGES_KWARGS, sizes=[(20, 19)]),
    ),
    ConsistencyConfig(
103
        v2_transforms.TenCrop,
104
105
106
107
        legacy_transforms.TenCrop,
        [
            ArgsKwargs(18),
            ArgsKwargs((18, 13)),
108
            ArgsKwargs(18, vertical_flip=True),
109
110
111
112
        ],
        make_images_kwargs=dict(DEFAULT_MAKE_IMAGES_KWARGS, sizes=[(20, 19)]),
    ),
    ConsistencyConfig(
113
        v2_transforms.Pad,
114
115
        legacy_transforms.Pad,
        [
116
            NotScriptableArgsKwargs(3),
117
118
119
            ArgsKwargs([3]),
            ArgsKwargs([2, 3]),
            ArgsKwargs([3, 2, 1, 4]),
120
121
122
123
124
            NotScriptableArgsKwargs(5, fill=1, padding_mode="constant"),
            ArgsKwargs([5], fill=1, padding_mode="constant"),
            NotScriptableArgsKwargs(5, padding_mode="edge"),
            NotScriptableArgsKwargs(5, padding_mode="reflect"),
            NotScriptableArgsKwargs(5, padding_mode="symmetric"),
125
126
        ],
    ),
127
128
    *[
        ConsistencyConfig(
129
            v2_transforms.LinearTransformation,
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
            legacy_transforms.LinearTransformation,
            [
                ArgsKwargs(LINEAR_TRANSFORMATION_MATRIX.to(matrix_dtype), LINEAR_TRANSFORMATION_MEAN.to(matrix_dtype)),
            ],
            # Make sure that the product of the height, width and number of channels matches the number of elements in
            # `LINEAR_TRANSFORMATION_MEAN`. For example 2 * 6 * 3 == 4 * 3 * 3 == 36.
            make_images_kwargs=dict(
                DEFAULT_MAKE_IMAGES_KWARGS, sizes=[(2, 6), (4, 3)], color_spaces=["RGB"], dtypes=[image_dtype]
            ),
            supports_pil=False,
        )
        for matrix_dtype, image_dtype in [
            (torch.float32, torch.float32),
            (torch.float64, torch.float64),
            (torch.float32, torch.uint8),
            (torch.float64, torch.float32),
            (torch.float32, torch.float64),
        ]
    ],
149
    ConsistencyConfig(
150
        v2_transforms.Grayscale,
151
152
153
154
155
        legacy_transforms.Grayscale,
        [
            ArgsKwargs(num_output_channels=1),
            ArgsKwargs(num_output_channels=3),
        ],
156
        make_images_kwargs=dict(DEFAULT_MAKE_IMAGES_KWARGS, color_spaces=["RGB", "GRAY"]),
157
158
        # Use default tolerances of `torch.testing.assert_close`
        closeness_kwargs=dict(rtol=None, atol=None),
159
    ),
160
    ConsistencyConfig(
161
        v2_transforms.ToPILImage,
162
        legacy_transforms.ToPILImage,
163
        [NotScriptableArgsKwargs()],
164
165
        make_images_kwargs=dict(
            color_spaces=[
166
167
168
169
                "GRAY",
                "GRAY_ALPHA",
                "RGB",
                "RGBA",
170
171
172
173
174
175
            ],
            extra_dims=[()],
        ),
        supports_pil=False,
    ),
    ConsistencyConfig(
176
        v2_transforms.Lambda,
177
178
        legacy_transforms.Lambda,
        [
179
            NotScriptableArgsKwargs(lambda image: image / 2),
180
181
182
183
184
        ],
        # Technically, this also supports PIL, but it is overkill to write a function here that supports tensor and PIL
        # images given that the transform does nothing but call it anyway.
        supports_pil=False,
    ),
185
    ConsistencyConfig(
186
        v2_transforms.RandomEqualize,
187
188
189
190
191
192
193
194
        legacy_transforms.RandomEqualize,
        [
            ArgsKwargs(p=0),
            ArgsKwargs(p=1),
        ],
        make_images_kwargs=dict(DEFAULT_MAKE_IMAGES_KWARGS, dtypes=[torch.uint8]),
    ),
    ConsistencyConfig(
195
        v2_transforms.RandomInvert,
196
197
198
199
200
201
202
        legacy_transforms.RandomInvert,
        [
            ArgsKwargs(p=0),
            ArgsKwargs(p=1),
        ],
    ),
    ConsistencyConfig(
203
        v2_transforms.RandomPosterize,
204
205
206
207
208
209
210
211
212
        legacy_transforms.RandomPosterize,
        [
            ArgsKwargs(p=0, bits=5),
            ArgsKwargs(p=1, bits=1),
            ArgsKwargs(p=1, bits=3),
        ],
        make_images_kwargs=dict(DEFAULT_MAKE_IMAGES_KWARGS, dtypes=[torch.uint8]),
    ),
    ConsistencyConfig(
213
        v2_transforms.RandomSolarize,
214
215
216
217
218
219
220
        legacy_transforms.RandomSolarize,
        [
            ArgsKwargs(p=0, threshold=0.5),
            ArgsKwargs(p=1, threshold=0.3),
            ArgsKwargs(p=1, threshold=0.99),
        ],
    ),
221
222
    *[
        ConsistencyConfig(
223
            v2_transforms.RandomAutocontrast,
224
225
226
227
228
229
230
231
232
233
            legacy_transforms.RandomAutocontrast,
            [
                ArgsKwargs(p=0),
                ArgsKwargs(p=1),
            ],
            make_images_kwargs=dict(DEFAULT_MAKE_IMAGES_KWARGS, dtypes=[dt]),
            closeness_kwargs=ckw,
        )
        for dt, ckw in [(torch.uint8, dict(atol=1, rtol=0)), (torch.float32, dict(rtol=None, atol=None))]
    ],
234
    ConsistencyConfig(
235
        v2_transforms.RandomAdjustSharpness,
236
237
238
        legacy_transforms.RandomAdjustSharpness,
        [
            ArgsKwargs(p=0, sharpness_factor=0.5),
239
            ArgsKwargs(p=1, sharpness_factor=0.2),
240
241
            ArgsKwargs(p=1, sharpness_factor=0.99),
        ],
242
        closeness_kwargs={"atol": 1e-6, "rtol": 1e-6},
243
244
    ),
    ConsistencyConfig(
245
        v2_transforms.RandomGrayscale,
246
247
248
249
250
        legacy_transforms.RandomGrayscale,
        [
            ArgsKwargs(p=0),
            ArgsKwargs(p=1),
        ],
251
252
253
        make_images_kwargs=dict(DEFAULT_MAKE_IMAGES_KWARGS, color_spaces=["RGB", "GRAY"]),
        # Use default tolerances of `torch.testing.assert_close`
        closeness_kwargs=dict(rtol=None, atol=None),
254
255
    ),
    ConsistencyConfig(
256
        v2_transforms.RandomResizedCrop,
257
258
259
260
261
        legacy_transforms.RandomResizedCrop,
        [
            ArgsKwargs(16),
            ArgsKwargs(17, scale=(0.3, 0.7)),
            ArgsKwargs(25, ratio=(0.5, 1.5)),
262
            ArgsKwargs((31, 28), interpolation=v2_transforms.InterpolationMode.NEAREST),
263
            ArgsKwargs((31, 28), interpolation=PIL.Image.NEAREST),
264
265
266
            ArgsKwargs((29, 32), antialias=False),
            ArgsKwargs((28, 31), antialias=True),
        ],
267
268
        # atol=1 due to Resize v2 is using native uint8 interpolate path for bilinear and nearest modes
        closeness_kwargs=dict(rtol=0, atol=1),
269
    ),
270
271
272
273
274
275
276
277
278
    ConsistencyConfig(
        v2_transforms.RandomResizedCrop,
        legacy_transforms.RandomResizedCrop,
        [
            ArgsKwargs((33, 26), interpolation=v2_transforms.InterpolationMode.BICUBIC, antialias=True),
            ArgsKwargs((33, 26), interpolation=PIL.Image.BICUBIC, antialias=True),
        ],
        closeness_kwargs=dict(rtol=0, atol=21),
    ),
279
    ConsistencyConfig(
280
        v2_transforms.ColorJitter,
281
282
283
284
285
286
287
288
289
290
291
        legacy_transforms.ColorJitter,
        [
            ArgsKwargs(),
            ArgsKwargs(brightness=0.1),
            ArgsKwargs(brightness=(0.2, 0.3)),
            ArgsKwargs(contrast=0.4),
            ArgsKwargs(contrast=(0.5, 0.6)),
            ArgsKwargs(saturation=0.7),
            ArgsKwargs(saturation=(0.8, 0.9)),
            ArgsKwargs(hue=0.3),
            ArgsKwargs(hue=(-0.1, 0.2)),
292
            ArgsKwargs(brightness=0.1, contrast=0.4, saturation=0.5, hue=0.3),
293
        ],
294
        closeness_kwargs={"atol": 1e-5, "rtol": 1e-5},
295
296
    ),
    ConsistencyConfig(
297
        v2_transforms.GaussianBlur,
298
299
300
301
302
303
304
        legacy_transforms.GaussianBlur,
        [
            ArgsKwargs(kernel_size=3),
            ArgsKwargs(kernel_size=(1, 5)),
            ArgsKwargs(kernel_size=3, sigma=0.7),
            ArgsKwargs(kernel_size=5, sigma=(0.3, 1.4)),
        ],
305
        closeness_kwargs={"rtol": 1e-5, "atol": 1e-5},
306
307
    ),
    ConsistencyConfig(
308
        v2_transforms.RandomPerspective,
309
310
311
312
313
        legacy_transforms.RandomPerspective,
        [
            ArgsKwargs(p=0),
            ArgsKwargs(p=1),
            ArgsKwargs(p=1, distortion_scale=0.3),
314
            ArgsKwargs(p=1, distortion_scale=0.2, interpolation=v2_transforms.InterpolationMode.NEAREST),
315
            ArgsKwargs(p=1, distortion_scale=0.2, interpolation=PIL.Image.NEAREST),
316
317
318
            ArgsKwargs(p=1, distortion_scale=0.1, fill=1),
            ArgsKwargs(p=1, distortion_scale=0.4, fill=(1, 2, 3)),
        ],
319
        closeness_kwargs={"atol": None, "rtol": None},
320
    ),
321
    ConsistencyConfig(
322
        v2_transforms.PILToTensor,
323
324
325
        legacy_transforms.PILToTensor,
    ),
    ConsistencyConfig(
326
        v2_transforms.ToTensor,
327
328
329
        legacy_transforms.ToTensor,
    ),
    ConsistencyConfig(
330
        v2_transforms.Compose,
331
332
333
        legacy_transforms.Compose,
    ),
    ConsistencyConfig(
334
        v2_transforms.RandomApply,
335
336
337
        legacy_transforms.RandomApply,
    ),
    ConsistencyConfig(
338
        v2_transforms.RandomChoice,
339
340
341
        legacy_transforms.RandomChoice,
    ),
    ConsistencyConfig(
342
        v2_transforms.RandomOrder,
343
344
345
        legacy_transforms.RandomOrder,
    ),
    ConsistencyConfig(
346
        v2_transforms.AugMix,
347
348
349
        legacy_transforms.AugMix,
    ),
    ConsistencyConfig(
350
        v2_transforms.AutoAugment,
351
352
353
        legacy_transforms.AutoAugment,
    ),
    ConsistencyConfig(
354
        v2_transforms.RandAugment,
355
356
357
        legacy_transforms.RandAugment,
    ),
    ConsistencyConfig(
358
        v2_transforms.TrivialAugmentWide,
359
360
        legacy_transforms.TrivialAugmentWide,
    ),
361
362
363
]


364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
@pytest.mark.parametrize("config", CONSISTENCY_CONFIGS, ids=lambda config: config.legacy_cls.__name__)
def test_signature_consistency(config):
    legacy_params = dict(inspect.signature(config.legacy_cls).parameters)
    prototype_params = dict(inspect.signature(config.prototype_cls).parameters)

    for param in config.removed_params:
        legacy_params.pop(param, None)

    missing = legacy_params.keys() - prototype_params.keys()
    if missing:
        raise AssertionError(
            f"The prototype transform does not support the parameters "
            f"{sequence_to_str(sorted(missing), separate_last='and ')}, but the legacy transform does. "
            f"If that is intentional, e.g. pending deprecation, please add the parameters to the `removed_params` on "
            f"the `ConsistencyConfig`."
        )

    extra = prototype_params.keys() - legacy_params.keys()
382
383
384
385
386
387
    extra_without_default = {
        param
        for param in extra
        if prototype_params[param].default is inspect.Parameter.empty
        and prototype_params[param].kind not in {inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD}
    }
388
389
    if extra_without_default:
        raise AssertionError(
390
391
392
            f"The prototype transform requires the parameters "
            f"{sequence_to_str(sorted(extra_without_default), separate_last='and ')}, but the legacy transform does "
            f"not. Please add a default value."
393
394
        )

395
396
397
398
399
400
    legacy_signature = list(legacy_params.keys())
    # Since we made sure that we don't have any extra parameters without default above, we clamp the prototype signature
    # to the same number of parameters as the legacy one
    prototype_signature = list(prototype_params.keys())[: len(legacy_signature)]

    assert prototype_signature == legacy_signature
401
402


403
404
405
def check_call_consistency(
    prototype_transform, legacy_transform, images=None, supports_pil=True, closeness_kwargs=None
):
406
407
    if images is None:
        images = make_images(**DEFAULT_MAKE_IMAGES_KWARGS)
408

409
410
    closeness_kwargs = closeness_kwargs or dict()

411
412
    for image in images:
        image_repr = f"[{tuple(image.shape)}, {str(image.dtype).rsplit('.')[-1]}]"
413
414
415

        image_tensor = torch.Tensor(image)
        try:
416
            torch.manual_seed(0)
417
            output_legacy_tensor = legacy_transform(image_tensor)
418
419
        except Exception as exc:
            raise pytest.UsageError(
420
                f"Transforming a tensor image {image_repr} failed in the legacy transform with the "
421
                f"error above. This means that you need to specify the parameters passed to `make_images` through the "
422
423
424
425
                "`make_images_kwargs` of the `ConsistencyConfig`."
            ) from exc

        try:
426
            torch.manual_seed(0)
427
            output_prototype_tensor = prototype_transform(image_tensor)
428
429
        except Exception as exc:
            raise AssertionError(
430
                f"Transforming a tensor image with shape {image_repr} failed in the prototype transform with "
431
                f"the error above. This means there is a consistency bug either in `_get_params` or in the "
432
                f"`is_pure_tensor` path in `_transform`."
433
434
            ) from exc

435
        assert_close(
436
437
438
            output_prototype_tensor,
            output_legacy_tensor,
            msg=lambda msg: f"Tensor image consistency check failed with: \n\n{msg}",
439
            **closeness_kwargs,
440
441
442
        )

        try:
443
            torch.manual_seed(0)
444
            output_prototype_image = prototype_transform(image)
445
446
        except Exception as exc:
            raise AssertionError(
447
                f"Transforming a image tv_tensor with shape {image_repr} failed in the prototype transform with "
448
                f"the error above. This means there is a consistency bug either in `_get_params` or in the "
449
                f"`tv_tensors.Image` path in `_transform`."
450
451
            ) from exc

452
        assert_close(
453
            output_prototype_image,
454
            output_prototype_tensor,
455
            msg=lambda msg: f"Output for tv_tensor and tensor images is not equal: \n\n{msg}",
456
            **closeness_kwargs,
457
458
        )

459
        if image.ndim == 3 and supports_pil:
460
            image_pil = to_pil_image(image)
461

462
            try:
463
                torch.manual_seed(0)
464
                output_legacy_pil = legacy_transform(image_pil)
465
466
            except Exception as exc:
                raise pytest.UsageError(
467
                    f"Transforming a PIL image with shape {image_repr} failed in the legacy transform with the "
468
469
470
471
472
                    f"error above. If this transform does not support PIL images, set `supports_pil=False` on the "
                    "`ConsistencyConfig`. "
                ) from exc

            try:
473
                torch.manual_seed(0)
474
                output_prototype_pil = prototype_transform(image_pil)
475
476
            except Exception as exc:
                raise AssertionError(
477
                    f"Transforming a PIL image with shape {image_repr} failed in the prototype transform with "
478
479
480
481
                    f"the error above. This means there is a consistency bug either in `_get_params` or in the "
                    f"`PIL.Image.Image` path in `_transform`."
                ) from exc

482
            assert_close(
483
484
                output_prototype_pil,
                output_legacy_pil,
485
                msg=lambda msg: f"PIL image consistency check failed with: \n\n{msg}",
486
                **closeness_kwargs,
487
            )
488
489


490
@pytest.mark.parametrize(
491
492
    ("config", "args_kwargs"),
    [
493
494
495
        pytest.param(
            config, args_kwargs, id=f"{config.legacy_cls.__name__}-{idx:0{len(str(len(config.args_kwargs)))}d}"
        )
496
        for config in CONSISTENCY_CONFIGS
497
        for idx, args_kwargs in enumerate(config.args_kwargs)
498
    ],
499
)
500
@pytest.mark.filterwarnings("ignore")
501
def test_call_consistency(config, args_kwargs):
502
503
504
    args, kwargs = args_kwargs

    try:
505
        legacy_transform = config.legacy_cls(*args, **kwargs)
506
507
508
509
510
511
512
    except Exception as exc:
        raise pytest.UsageError(
            f"Initializing the legacy transform failed with the error above. "
            f"Please correct the `ArgsKwargs({args_kwargs})` in the `ConsistencyConfig`."
        ) from exc

    try:
513
        prototype_transform = config.prototype_cls(*args, **kwargs)
514
515
516
517
518
519
    except Exception as exc:
        raise AssertionError(
            "Initializing the prototype transform failed with the error above. "
            "This means there is a consistency bug in the constructor."
        ) from exc

520
521
522
523
524
    check_call_consistency(
        prototype_transform,
        legacy_transform,
        images=make_images(**config.make_images_kwargs),
        supports_pil=config.supports_pil,
525
        closeness_kwargs=config.closeness_kwargs,
526
527
528
    )


529
530
531
532
533
534
535
536
537
get_params_parametrization = pytest.mark.parametrize(
    ("config", "get_params_args_kwargs"),
    [
        pytest.param(
            next(config for config in CONSISTENCY_CONFIGS if config.prototype_cls is transform_cls),
            get_params_args_kwargs,
            id=transform_cls.__name__,
        )
        for transform_cls, get_params_args_kwargs in [
538
539
540
541
542
            (v2_transforms.RandomResizedCrop, ArgsKwargs(make_image(), scale=[0.3, 0.7], ratio=[0.5, 1.5])),
            (v2_transforms.ColorJitter, ArgsKwargs(brightness=None, contrast=None, saturation=None, hue=None)),
            (v2_transforms.GaussianBlur, ArgsKwargs(0.3, 1.4)),
            (v2_transforms.RandomPerspective, ArgsKwargs(23, 17, 0.5)),
            (v2_transforms.AutoAugment, ArgsKwargs(5)),
543
544
        ]
    ],
545
)
546
547


548
@get_params_parametrization
549
def test_get_params_alias(config, get_params_args_kwargs):
550
551
    assert config.prototype_cls.get_params is config.legacy_cls.get_params

552
553
554
555
556
    if not config.args_kwargs:
        return
    args, kwargs = config.args_kwargs[0]
    legacy_transform = config.legacy_cls(*args, **kwargs)
    prototype_transform = config.prototype_cls(*args, **kwargs)
557

558
559
560
    assert prototype_transform.get_params is legacy_transform.get_params


561
@get_params_parametrization
562
563
564
565
566
567
568
569
570
def test_get_params_jit(config, get_params_args_kwargs):
    get_params_args, get_params_kwargs = get_params_args_kwargs

    torch.jit.script(config.prototype_cls.get_params)(*get_params_args, **get_params_kwargs)

    if not config.args_kwargs:
        return
    args, kwargs = config.args_kwargs[0]
    transform = config.prototype_cls(*args, **kwargs)
571

572
    torch.jit.script(transform.get_params)(*get_params_args, **get_params_kwargs)
573
574


575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
@pytest.mark.parametrize(
    ("config", "args_kwargs"),
    [
        pytest.param(
            config, args_kwargs, id=f"{config.legacy_cls.__name__}-{idx:0{len(str(len(config.args_kwargs)))}d}"
        )
        for config in CONSISTENCY_CONFIGS
        for idx, args_kwargs in enumerate(config.args_kwargs)
        if not isinstance(args_kwargs, NotScriptableArgsKwargs)
    ],
)
def test_jit_consistency(config, args_kwargs):
    args, kwargs = args_kwargs

    prototype_transform_eager = config.prototype_cls(*args, **kwargs)
    legacy_transform_eager = config.legacy_cls(*args, **kwargs)

    legacy_transform_scripted = torch.jit.script(legacy_transform_eager)
    prototype_transform_scripted = torch.jit.script(prototype_transform_eager)

    for image in make_images(**config.make_images_kwargs):
        image = image.as_subclass(torch.Tensor)

        torch.manual_seed(0)
        output_legacy_scripted = legacy_transform_scripted(image)

        torch.manual_seed(0)
        output_prototype_scripted = prototype_transform_scripted(image)

        assert_close(output_prototype_scripted, output_legacy_scripted, **config.closeness_kwargs)


607
608
609
610
611
612
613
614
615
616
class TestContainerTransforms:
    """
    Since we are testing containers here, we also need some transforms to wrap. Thus, testing a container transform for
    consistency automatically tests the wrapped transforms consistency.

    Instead of complicated mocking or creating custom transforms just for these tests, here we use deterministic ones
    that were already tested for consistency above.
    """

    def test_compose(self):
617
        prototype_transform = v2_transforms.Compose(
618
            [
619
620
                v2_transforms.Resize(256),
                v2_transforms.CenterCrop(224),
621
622
623
624
625
626
627
628
629
            ]
        )
        legacy_transform = legacy_transforms.Compose(
            [
                legacy_transforms.Resize(256),
                legacy_transforms.CenterCrop(224),
            ]
        )

630
631
        # atol=1 due to Resize v2 is using native uint8 interpolate path for bilinear and nearest modes
        check_call_consistency(prototype_transform, legacy_transform, closeness_kwargs=dict(rtol=0, atol=1))
632
633

    @pytest.mark.parametrize("p", [0, 0.1, 0.5, 0.9, 1])
634
635
    @pytest.mark.parametrize("sequence_type", [list, nn.ModuleList])
    def test_random_apply(self, p, sequence_type):
636
        prototype_transform = v2_transforms.RandomApply(
637
638
            sequence_type(
                [
639
640
                    v2_transforms.Resize(256),
                    v2_transforms.CenterCrop(224),
641
642
                ]
            ),
643
644
645
            p=p,
        )
        legacy_transform = legacy_transforms.RandomApply(
646
647
648
649
650
651
            sequence_type(
                [
                    legacy_transforms.Resize(256),
                    legacy_transforms.CenterCrop(224),
                ]
            ),
652
653
654
            p=p,
        )

655
656
        # atol=1 due to Resize v2 is using native uint8 interpolate path for bilinear and nearest modes
        check_call_consistency(prototype_transform, legacy_transform, closeness_kwargs=dict(rtol=0, atol=1))
657

658
659
660
661
662
        if sequence_type is nn.ModuleList:
            # quick and dirty test that it is jit-scriptable
            scripted = torch.jit.script(prototype_transform)
            scripted(torch.rand(1, 3, 300, 300))

663
    # We can't test other values for `p` since the random parameter generation is different
664
665
    @pytest.mark.parametrize("probabilities", [(0, 1), (1, 0)])
    def test_random_choice(self, probabilities):
666
        prototype_transform = v2_transforms.RandomChoice(
667
            [
668
                v2_transforms.Resize(256),
669
670
                legacy_transforms.CenterCrop(224),
            ],
671
            p=probabilities,
672
673
674
675
676
677
        )
        legacy_transform = legacy_transforms.RandomChoice(
            [
                legacy_transforms.Resize(256),
                legacy_transforms.CenterCrop(224),
            ],
678
            p=probabilities,
679
680
        )

681
682
        # atol=1 due to Resize v2 is using native uint8 interpolate path for bilinear and nearest modes
        check_call_consistency(prototype_transform, legacy_transform, closeness_kwargs=dict(rtol=0, atol=1))
683
684


685
686
class TestToTensorTransforms:
    def test_pil_to_tensor(self):
687
        prototype_transform = v2_transforms.PILToTensor()
688
689
        legacy_transform = legacy_transforms.PILToTensor()

690
        for image in make_images(extra_dims=[()]):
691
            image_pil = to_pil_image(image)
692
693
694
695

            assert_equal(prototype_transform(image_pil), legacy_transform(image_pil))

    def test_to_tensor(self):
696
        with pytest.warns(UserWarning, match=re.escape("The transform `ToTensor()` is deprecated")):
697
            prototype_transform = v2_transforms.ToTensor()
698
699
        legacy_transform = legacy_transforms.ToTensor()

700
        for image in make_images(extra_dims=[()]):
701
            image_pil = to_pil_image(image)
702
703
704
705
            image_numpy = np.array(image_pil)

            assert_equal(prototype_transform(image_pil), legacy_transform(image_pil))
            assert_equal(prototype_transform(image_numpy), legacy_transform(image_numpy))
706
707
708
709
710
711
712
713


class TestAATransforms:
    @pytest.mark.parametrize(
        "inpt",
        [
            torch.randint(0, 256, size=(1, 3, 256, 256), dtype=torch.uint8),
            PIL.Image.new("RGB", (256, 256), 123),
714
            tv_tensors.Image(torch.randint(0, 256, size=(1, 3, 256, 256), dtype=torch.uint8)),
715
716
717
718
        ],
    )
    @pytest.mark.parametrize(
        "interpolation",
719
        [
720
721
            v2_transforms.InterpolationMode.NEAREST,
            v2_transforms.InterpolationMode.BILINEAR,
722
723
            PIL.Image.NEAREST,
        ],
724
725
726
    )
    def test_randaug(self, inpt, interpolation, mocker):
        t_ref = legacy_transforms.RandAugment(interpolation=interpolation, num_ops=1)
727
        t = v2_transforms.RandAugment(interpolation=interpolation, num_ops=1)
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748

        le = len(t._AUGMENTATION_SPACE)
        keys = list(t._AUGMENTATION_SPACE.keys())
        randint_values = []
        for i in range(le):
            # Stable API, op_index random call
            randint_values.append(i)
            # Stable API, if signed there is another random call
            if t._AUGMENTATION_SPACE[keys[i]][1]:
                randint_values.append(0)
            # New API, _get_random_item
            randint_values.append(i)
        randint_values = iter(randint_values)

        mocker.patch("torch.randint", side_effect=lambda *arg, **kwargs: torch.tensor(next(randint_values)))
        mocker.patch("torch.rand", return_value=1.0)

        for i in range(le):
            expected_output = t_ref(inpt)
            output = t(inpt)

749
            assert_close(expected_output, output, atol=1, rtol=0.1)
750

751
752
753
754
755
756
757
    @pytest.mark.parametrize(
        "interpolation",
        [
            v2_transforms.InterpolationMode.NEAREST,
            v2_transforms.InterpolationMode.BILINEAR,
        ],
    )
758
759
    @pytest.mark.parametrize("fill", [None, 85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1], 1])
    def test_randaug_jit(self, interpolation, fill):
760
        inpt = torch.randint(0, 256, size=(1, 3, 256, 256), dtype=torch.uint8)
761
762
        t_ref = legacy_transforms.RandAugment(interpolation=interpolation, num_ops=1, fill=fill)
        t = v2_transforms.RandAugment(interpolation=interpolation, num_ops=1, fill=fill)
763
764
765
766
767
768
769
770
771
772
773
774

        tt_ref = torch.jit.script(t_ref)
        tt = torch.jit.script(t)

        torch.manual_seed(12)
        expected_output = tt_ref(inpt)

        torch.manual_seed(12)
        scripted_output = tt(inpt)

        assert_equal(scripted_output, expected_output)

775
776
777
778
779
    @pytest.mark.parametrize(
        "inpt",
        [
            torch.randint(0, 256, size=(1, 3, 256, 256), dtype=torch.uint8),
            PIL.Image.new("RGB", (256, 256), 123),
780
            tv_tensors.Image(torch.randint(0, 256, size=(1, 3, 256, 256), dtype=torch.uint8)),
781
782
783
784
        ],
    )
    @pytest.mark.parametrize(
        "interpolation",
785
        [
786
787
            v2_transforms.InterpolationMode.NEAREST,
            v2_transforms.InterpolationMode.BILINEAR,
788
789
            PIL.Image.NEAREST,
        ],
790
791
792
    )
    def test_trivial_aug(self, inpt, interpolation, mocker):
        t_ref = legacy_transforms.TrivialAugmentWide(interpolation=interpolation)
793
        t = v2_transforms.TrivialAugmentWide(interpolation=interpolation)
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

        le = len(t._AUGMENTATION_SPACE)
        keys = list(t._AUGMENTATION_SPACE.keys())
        randint_values = []
        for i in range(le):
            # Stable API, op_index random call
            randint_values.append(i)
            key = keys[i]
            # Stable API, random magnitude
            aug_op = t._AUGMENTATION_SPACE[key]
            magnitudes = aug_op[0](2, 0, 0)
            if magnitudes is not None:
                randint_values.append(5)
            # Stable API, if signed there is another random call
            if aug_op[1]:
                randint_values.append(0)
            # New API, _get_random_item
            randint_values.append(i)
            # New API, random magnitude
            if magnitudes is not None:
                randint_values.append(5)

        randint_values = iter(randint_values)

        mocker.patch("torch.randint", side_effect=lambda *arg, **kwargs: torch.tensor(next(randint_values)))
        mocker.patch("torch.rand", return_value=1.0)

        for _ in range(le):
            expected_output = t_ref(inpt)
            output = t(inpt)

825
            assert_close(expected_output, output, atol=1, rtol=0.1)
826

827
828
829
830
831
832
833
    @pytest.mark.parametrize(
        "interpolation",
        [
            v2_transforms.InterpolationMode.NEAREST,
            v2_transforms.InterpolationMode.BILINEAR,
        ],
    )
834
835
    @pytest.mark.parametrize("fill", [None, 85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1], 1])
    def test_trivial_aug_jit(self, interpolation, fill):
836
        inpt = torch.randint(0, 256, size=(1, 3, 256, 256), dtype=torch.uint8)
837
838
        t_ref = legacy_transforms.TrivialAugmentWide(interpolation=interpolation, fill=fill)
        t = v2_transforms.TrivialAugmentWide(interpolation=interpolation, fill=fill)
839
840
841
842
843
844
845
846
847
848
849
850

        tt_ref = torch.jit.script(t_ref)
        tt = torch.jit.script(t)

        torch.manual_seed(12)
        expected_output = tt_ref(inpt)

        torch.manual_seed(12)
        scripted_output = tt(inpt)

        assert_equal(scripted_output, expected_output)

851
852
853
854
855
    @pytest.mark.parametrize(
        "inpt",
        [
            torch.randint(0, 256, size=(1, 3, 256, 256), dtype=torch.uint8),
            PIL.Image.new("RGB", (256, 256), 123),
856
            tv_tensors.Image(torch.randint(0, 256, size=(1, 3, 256, 256), dtype=torch.uint8)),
857
858
859
860
        ],
    )
    @pytest.mark.parametrize(
        "interpolation",
861
        [
862
863
            v2_transforms.InterpolationMode.NEAREST,
            v2_transforms.InterpolationMode.BILINEAR,
864
865
            PIL.Image.NEAREST,
        ],
866
867
868
869
    )
    def test_augmix(self, inpt, interpolation, mocker):
        t_ref = legacy_transforms.AugMix(interpolation=interpolation, mixture_width=1, chain_depth=1)
        t_ref._sample_dirichlet = lambda t: t.softmax(dim=-1)
870
        t = v2_transforms.AugMix(interpolation=interpolation, mixture_width=1, chain_depth=1)
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
        t._sample_dirichlet = lambda t: t.softmax(dim=-1)

        le = len(t._AUGMENTATION_SPACE)
        keys = list(t._AUGMENTATION_SPACE.keys())
        randint_values = []
        for i in range(le):
            # Stable API, op_index random call
            randint_values.append(i)
            key = keys[i]
            # Stable API, random magnitude
            aug_op = t._AUGMENTATION_SPACE[key]
            magnitudes = aug_op[0](2, 0, 0)
            if magnitudes is not None:
                randint_values.append(5)
            # Stable API, if signed there is another random call
            if aug_op[1]:
                randint_values.append(0)
            # New API, _get_random_item
            randint_values.append(i)
            # New API, random magnitude
            if magnitudes is not None:
                randint_values.append(5)

        randint_values = iter(randint_values)

        mocker.patch("torch.randint", side_effect=lambda *arg, **kwargs: torch.tensor(next(randint_values)))
        mocker.patch("torch.rand", return_value=1.0)

        expected_output = t_ref(inpt)
        output = t(inpt)

        assert_equal(expected_output, output)

904
905
906
907
908
909
910
    @pytest.mark.parametrize(
        "interpolation",
        [
            v2_transforms.InterpolationMode.NEAREST,
            v2_transforms.InterpolationMode.BILINEAR,
        ],
    )
911
912
    @pytest.mark.parametrize("fill", [None, 85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1], 1])
    def test_augmix_jit(self, interpolation, fill):
913
914
        inpt = torch.randint(0, 256, size=(1, 3, 256, 256), dtype=torch.uint8)

915
916
        t_ref = legacy_transforms.AugMix(interpolation=interpolation, mixture_width=1, chain_depth=1, fill=fill)
        t = v2_transforms.AugMix(interpolation=interpolation, mixture_width=1, chain_depth=1, fill=fill)
917
918
919
920
921
922
923
924
925
926
927
928

        tt_ref = torch.jit.script(t_ref)
        tt = torch.jit.script(t)

        torch.manual_seed(12)
        expected_output = tt_ref(inpt)

        torch.manual_seed(12)
        scripted_output = tt(inpt)

        assert_equal(scripted_output, expected_output)

929
930
931
932
933
    @pytest.mark.parametrize(
        "inpt",
        [
            torch.randint(0, 256, size=(1, 3, 256, 256), dtype=torch.uint8),
            PIL.Image.new("RGB", (256, 256), 123),
934
            tv_tensors.Image(torch.randint(0, 256, size=(1, 3, 256, 256), dtype=torch.uint8)),
935
936
937
938
        ],
    )
    @pytest.mark.parametrize(
        "interpolation",
939
        [
940
941
            v2_transforms.InterpolationMode.NEAREST,
            v2_transforms.InterpolationMode.BILINEAR,
942
943
            PIL.Image.NEAREST,
        ],
944
945
946
947
    )
    def test_aa(self, inpt, interpolation):
        aa_policy = legacy_transforms.AutoAugmentPolicy("imagenet")
        t_ref = legacy_transforms.AutoAugment(aa_policy, interpolation=interpolation)
948
        t = v2_transforms.AutoAugment(aa_policy, interpolation=interpolation)
949
950
951
952
953
954
955
956

        torch.manual_seed(12)
        expected_output = t_ref(inpt)

        torch.manual_seed(12)
        output = t(inpt)

        assert_equal(expected_output, output)
957

958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
    @pytest.mark.parametrize(
        "interpolation",
        [
            v2_transforms.InterpolationMode.NEAREST,
            v2_transforms.InterpolationMode.BILINEAR,
        ],
    )
    def test_aa_jit(self, interpolation):
        inpt = torch.randint(0, 256, size=(1, 3, 256, 256), dtype=torch.uint8)
        aa_policy = legacy_transforms.AutoAugmentPolicy("imagenet")
        t_ref = legacy_transforms.AutoAugment(aa_policy, interpolation=interpolation)
        t = v2_transforms.AutoAugment(aa_policy, interpolation=interpolation)

        tt_ref = torch.jit.script(t_ref)
        tt = torch.jit.script(t)

        torch.manual_seed(12)
        expected_output = tt_ref(inpt)

        torch.manual_seed(12)
        scripted_output = tt(inpt)

        assert_equal(scripted_output, expected_output)

982

983
def import_transforms_from_references(reference):
984
985
986
987
988
989
990
991
992
993
    HERE = Path(__file__).parent
    PROJECT_ROOT = HERE.parent

    loader = importlib.machinery.SourceFileLoader(
        "transforms", str(PROJECT_ROOT / "references" / reference / "transforms.py")
    )
    spec = importlib.util.spec_from_loader("transforms", loader)
    module = importlib.util.module_from_spec(spec)
    loader.exec_module(module)
    return module
994
995
996


det_transforms = import_transforms_from_references("detection")
997
998
999


class TestRefDetTransforms:
1000
    def make_tv_tensors(self, with_mask=True):
1001
1002
1003
        size = (600, 800)
        num_objects = 22

1004
1005
1006
        def make_label(extra_dims, categories):
            return torch.randint(categories, extra_dims, dtype=torch.int64)

1007
        pil_image = to_pil_image(make_image(size=size, color_space="RGB"))
1008
        target = {
1009
            "boxes": make_bounding_boxes(canvas_size=size, format="XYXY", batch_dims=(num_objects,), dtype=torch.float),
1010
1011
1012
1013
1014
1015
1016
            "labels": make_label(extra_dims=(num_objects,), categories=80),
        }
        if with_mask:
            target["masks"] = make_detection_mask(size=size, num_objects=num_objects, dtype=torch.long)

        yield (pil_image, target)

1017
        tensor_image = torch.Tensor(make_image(size=size, color_space="RGB", dtype=torch.float32))
1018
        target = {
1019
            "boxes": make_bounding_boxes(canvas_size=size, format="XYXY", batch_dims=(num_objects,), dtype=torch.float),
1020
1021
1022
1023
1024
1025
1026
            "labels": make_label(extra_dims=(num_objects,), categories=80),
        }
        if with_mask:
            target["masks"] = make_detection_mask(size=size, num_objects=num_objects, dtype=torch.long)

        yield (tensor_image, target)

1027
        tv_tensor_image = make_image(size=size, color_space="RGB", dtype=torch.float32)
1028
        target = {
1029
            "boxes": make_bounding_boxes(canvas_size=size, format="XYXY", batch_dims=(num_objects,), dtype=torch.float),
1030
1031
1032
1033
1034
            "labels": make_label(extra_dims=(num_objects,), categories=80),
        }
        if with_mask:
            target["masks"] = make_detection_mask(size=size, num_objects=num_objects, dtype=torch.long)

1035
        yield (tv_tensor_image, target)
1036
1037
1038
1039

    @pytest.mark.parametrize(
        "t_ref, t, data_kwargs",
        [
1040
            (det_transforms.RandomHorizontalFlip(p=1.0), v2_transforms.RandomHorizontalFlip(p=1.0), {}),
1041
1042
1043
1044
1045
            (
                det_transforms.RandomIoUCrop(),
                v2_transforms.Compose(
                    [
                        v2_transforms.RandomIoUCrop(),
1046
                        v2_transforms.SanitizeBoundingBoxes(labels_getter=lambda sample: sample[1]["labels"]),
1047
1048
1049
1050
                    ]
                ),
                {"with_mask": False},
            ),
1051
            (det_transforms.RandomZoomOut(), v2_transforms.RandomZoomOut(), {"with_mask": False}),
1052
            (det_transforms.ScaleJitter((1024, 1024)), v2_transforms.ScaleJitter((1024, 1024), antialias=True), {}),
1053
1054
1055
1056
            (
                det_transforms.RandomShortestSize(
                    min_size=(480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800), max_size=1333
                ),
1057
                v2_transforms.RandomShortestSize(
1058
1059
1060
1061
1062
1063
1064
                    min_size=(480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800), max_size=1333
                ),
                {},
            ),
        ],
    )
    def test_transform(self, t_ref, t, data_kwargs):
1065
        for dp in self.make_tv_tensors(**data_kwargs):
1066
1067
1068
1069
1070
1071
1072
1073
1074

            # We should use prototype transform first as reference transform performs inplace target update
            torch.manual_seed(12)
            output = t(dp)

            torch.manual_seed(12)
            expected_output = t_ref(*dp)

            assert_equal(expected_output, output)
1075
1076
1077
1078
1079
1080
1081
1082
1083


seg_transforms = import_transforms_from_references("segmentation")


# We need this transform for two reasons:
# 1. transforms.RandomCrop uses a different scheme to pad images and masks of insufficient size than its name
#    counterpart in the detection references. Thus, we cannot use it with `pad_if_needed=True`
# 2. transforms.Pad only supports a fixed padding, but the segmentation datasets don't have a fixed image size.
1084
class PadIfSmaller(v2_transforms.Transform):
1085
1086
1087
    def __init__(self, size, fill=0):
        super().__init__()
        self.size = size
1088
        self.fill = v2_transforms._geometry._setup_fill_arg(fill)
1089
1090

    def _get_params(self, sample):
Philip Meier's avatar
Philip Meier committed
1091
        height, width = query_size(sample)
1092
1093
1094
1095
1096
1097
1098
1099
        padding = [0, 0, max(self.size - width, 0), max(self.size - height, 0)]
        needs_padding = any(padding)
        return dict(padding=padding, needs_padding=needs_padding)

    def _transform(self, inpt, params):
        if not params["needs_padding"]:
            return inpt

1100
        fill = _get_fill(self.fill, type(inpt))
1101
        return prototype_F.pad(inpt, padding=params["padding"], fill=fill)
1102
1103
1104


class TestRefSegTransforms:
1105
    def make_tv_tensors(self, supports_pil=True, image_dtype=torch.uint8):
1106
        size = (256, 460)
1107
1108
1109
1110
        num_categories = 21

        conv_fns = []
        if supports_pil:
1111
            conv_fns.append(to_pil_image)
1112
1113
1114
        conv_fns.extend([torch.Tensor, lambda x: x])

        for conv_fn in conv_fns:
1115
1116
            tv_tensor_image = make_image(size=size, color_space="RGB", dtype=image_dtype)
            tv_tensor_mask = make_segmentation_mask(size=size, num_categories=num_categories, dtype=torch.uint8)
1117

1118
            dp = (conv_fn(tv_tensor_image), tv_tensor_mask)
1119
            dp_ref = (
1120
1121
                to_pil_image(tv_tensor_image) if supports_pil else tv_tensor_image.as_subclass(torch.Tensor),
                to_pil_image(tv_tensor_mask),
1122
1123
1124
1125
1126
1127
1128
1129
1130
            )

            yield dp, dp_ref

    def set_seed(self, seed=12):
        torch.manual_seed(seed)
        random.seed(seed)

    def check(self, t, t_ref, data_kwargs=None):
1131
        for dp, dp_ref in self.make_tv_tensors(**data_kwargs or dict()):
1132
1133

            self.set_seed()
1134
            actual = actual_image, actual_mask = t(dp)
1135
1136

            self.set_seed()
1137
1138
1139
1140
1141
            expected_image, expected_mask = t_ref(*dp_ref)
            if isinstance(actual_image, torch.Tensor) and not isinstance(expected_image, torch.Tensor):
                expected_image = legacy_F.pil_to_tensor(expected_image)
            expected_mask = legacy_F.pil_to_tensor(expected_mask).squeeze(0)
            expected = (expected_image, expected_mask)
1142

1143
            assert_equal(actual, expected)
1144
1145
1146
1147
1148
1149

    @pytest.mark.parametrize(
        ("t_ref", "t", "data_kwargs"),
        [
            (
                seg_transforms.RandomHorizontalFlip(flip_prob=1.0),
1150
                v2_transforms.RandomHorizontalFlip(p=1.0),
1151
1152
1153
1154
                dict(),
            ),
            (
                seg_transforms.RandomHorizontalFlip(flip_prob=0.0),
1155
                v2_transforms.RandomHorizontalFlip(p=0.0),
1156
1157
1158
1159
                dict(),
            ),
            (
                seg_transforms.RandomCrop(size=480),
1160
                v2_transforms.Compose(
1161
                    [
1162
                        PadIfSmaller(size=480, fill={tv_tensors.Mask: 255, "others": 0}),
1163
                        v2_transforms.RandomCrop(size=480),
1164
1165
1166
1167
1168
1169
                    ]
                ),
                dict(),
            ),
            (
                seg_transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
1170
                v2_transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
1171
1172
1173
1174
1175
1176
1177
                dict(supports_pil=False, image_dtype=torch.float),
            ),
        ],
    )
    def test_common(self, t_ref, t, data_kwargs):
        self.check(t, t_ref, data_kwargs)

1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189

@pytest.mark.parametrize(
    ("legacy_dispatcher", "name_only_params"),
    [
        (legacy_F.get_dimensions, {}),
        (legacy_F.get_image_size, {}),
        (legacy_F.get_image_num_channels, {}),
        (legacy_F.to_tensor, {}),
        (legacy_F.pil_to_tensor, {}),
        (legacy_F.convert_image_dtype, {}),
        (legacy_F.to_pil_image, {}),
        (legacy_F.normalize, {}),
1190
        (legacy_F.resize, {"interpolation"}),
1191
1192
1193
        (legacy_F.pad, {"padding", "fill"}),
        (legacy_F.crop, {}),
        (legacy_F.center_crop, {}),
1194
        (legacy_F.resized_crop, {"interpolation"}),
1195
        (legacy_F.hflip, {}),
1196
        (legacy_F.perspective, {"startpoints", "endpoints", "fill", "interpolation"}),
1197
1198
1199
1200
1201
1202
1203
1204
        (legacy_F.vflip, {}),
        (legacy_F.five_crop, {}),
        (legacy_F.ten_crop, {}),
        (legacy_F.adjust_brightness, {}),
        (legacy_F.adjust_contrast, {}),
        (legacy_F.adjust_saturation, {}),
        (legacy_F.adjust_hue, {}),
        (legacy_F.adjust_gamma, {}),
1205
1206
        (legacy_F.rotate, {"center", "fill", "interpolation"}),
        (legacy_F.affine, {"angle", "translate", "center", "fill", "interpolation"}),
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
        (legacy_F.to_grayscale, {}),
        (legacy_F.rgb_to_grayscale, {}),
        (legacy_F.to_tensor, {}),
        (legacy_F.erase, {}),
        (legacy_F.gaussian_blur, {}),
        (legacy_F.invert, {}),
        (legacy_F.posterize, {}),
        (legacy_F.solarize, {}),
        (legacy_F.adjust_sharpness, {}),
        (legacy_F.autocontrast, {}),
        (legacy_F.equalize, {}),
1218
        (legacy_F.elastic_transform, {"fill", "interpolation"}),
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
    ],
)
def test_dispatcher_signature_consistency(legacy_dispatcher, name_only_params):
    legacy_signature = inspect.signature(legacy_dispatcher)
    legacy_params = list(legacy_signature.parameters.values())[1:]

    try:
        prototype_dispatcher = getattr(prototype_F, legacy_dispatcher.__name__)
    except AttributeError:
        raise AssertionError(
            f"Legacy dispatcher `F.{legacy_dispatcher.__name__}` has no prototype equivalent"
        ) from None

    prototype_signature = inspect.signature(prototype_dispatcher)
    prototype_params = list(prototype_signature.parameters.values())[1:]

    # Some dispatchers got extra parameters. This makes sure they have a default argument and thus are BC. We don't
    # need to check if parameters were added in the middle rather than at the end, since that will be caught by the
    # regular check below.
    prototype_params, new_prototype_params = (
        prototype_params[: len(legacy_params)],
        prototype_params[len(legacy_params) :],
    )
    for param in new_prototype_params:
        assert param.default is not param.empty

    # Some annotations were changed mostly to supersets of what was there before. Plus, some legacy dispatchers had no
    # annotations. In these cases we simply drop the annotation and default argument from the comparison
    for prototype_param, legacy_param in zip(prototype_params, legacy_params):
        if legacy_param.name in name_only_params:
            prototype_param._annotation = prototype_param._default = inspect.Parameter.empty
            legacy_param._annotation = legacy_param._default = inspect.Parameter.empty
        elif legacy_param.annotation is inspect.Parameter.empty:
            prototype_param._annotation = inspect.Parameter.empty

    assert prototype_params == legacy_params