test_transforms_tensor.py 34.5 KB
Newer Older
1
import os
2
import sys
3
4

import numpy as np
5
import pytest
6
import torch
7
import torchvision.transforms._pil_constants as _pil_constants
Nicolas Hug's avatar
Nicolas Hug committed
8
9
10
11
12
13
14
15
from common_utils import (
    get_tmp_dir,
    int_dtypes,
    float_dtypes,
    _create_data,
    _create_data_batch,
    _assert_equal_tensor_to_pil,
    _assert_approx_equal_tensor_to_pil,
16
    cpu_and_gpu,
17
    assert_equal,
Nicolas Hug's avatar
Nicolas Hug committed
18
)
19
20
21
from torchvision import transforms as T
from torchvision.transforms import InterpolationMode
from torchvision.transforms import functional as F
22
from torchvision.transforms.autoaugment import _apply_op
23

24
NEAREST, BILINEAR, BICUBIC = InterpolationMode.NEAREST, InterpolationMode.BILINEAR, InterpolationMode.BICUBIC
25
26


27
28
29
30
31
32
def _test_transform_vs_scripted(transform, s_transform, tensor, msg=None):
    torch.manual_seed(12)
    out1 = transform(tensor)
    torch.manual_seed(12)
    out2 = s_transform(tensor)
    assert_equal(out1, out2, msg=msg)
33

34

35
36
37
def _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors, msg=None):
    torch.manual_seed(12)
    transformed_batch = transform(batch_tensors)
38

39
40
    for i in range(len(batch_tensors)):
        img_tensor = batch_tensors[i, ...]
41
        torch.manual_seed(12)
42
43
        transformed_img = transform(img_tensor)
        assert_equal(transformed_img, transformed_batch[i, ...], msg=msg)
44

45
46
47
    torch.manual_seed(12)
    s_transformed_batch = s_transform(batch_tensors)
    assert_equal(transformed_batch, s_transformed_batch, msg=msg)
48
49


50
def _test_functional_op(f, device, channels=3, fn_kwargs=None, test_exact_match=True, **match_kwargs):
51
    fn_kwargs = fn_kwargs or {}
52

53
    tensor, pil_img = _create_data(height=10, width=10, channels=channels, device=device)
54
55
56
57
58
59
    transformed_tensor = f(tensor, **fn_kwargs)
    transformed_pil_img = f(pil_img, **fn_kwargs)
    if test_exact_match:
        _assert_equal_tensor_to_pil(transformed_tensor, transformed_pil_img, **match_kwargs)
    else:
        _assert_approx_equal_tensor_to_pil(transformed_tensor, transformed_pil_img, **match_kwargs)
vfdev's avatar
vfdev committed
60
61


vfdev's avatar
vfdev committed
62
def _test_class_op(transform_cls, device, channels=3, meth_kwargs=None, test_exact_match=True, **match_kwargs):
63
    meth_kwargs = meth_kwargs or {}
64

65
    # test for class interface
vfdev's avatar
vfdev committed
66
    f = transform_cls(**meth_kwargs)
67
    scripted_fn = torch.jit.script(f)
68

69
    tensor, pil_img = _create_data(26, 34, channels, device=device)
70
71
72
73
74
75
76
77
78
    # set seed to reproduce the same transformation for tensor and PIL image
    torch.manual_seed(12)
    transformed_tensor = f(tensor)
    torch.manual_seed(12)
    transformed_pil_img = f(pil_img)
    if test_exact_match:
        _assert_equal_tensor_to_pil(transformed_tensor, transformed_pil_img, **match_kwargs)
    else:
        _assert_approx_equal_tensor_to_pil(transformed_tensor.float(), transformed_pil_img, **match_kwargs)
79

80
81
82
83
    torch.manual_seed(12)
    transformed_tensor_script = scripted_fn(tensor)
    assert_equal(transformed_tensor, transformed_tensor_script)

84
    batch_tensors = _create_data_batch(height=23, width=34, channels=channels, num_samples=4, device=device)
85
86
87
    _test_transform_vs_scripted_on_batch(f, scripted_fn, batch_tensors)

    with get_tmp_dir() as tmp_dir:
vfdev's avatar
vfdev committed
88
        scripted_fn.save(os.path.join(tmp_dir, f"t_{transform_cls.__name__}.pt"))
89

90

91
92
93
def _test_op(func, method, device, channels=3, fn_kwargs=None, meth_kwargs=None, test_exact_match=True, **match_kwargs):
    _test_functional_op(func, device, channels, fn_kwargs, test_exact_match=test_exact_match, **match_kwargs)
    _test_class_op(method, device, channels, meth_kwargs, test_exact_match=test_exact_match, **match_kwargs)
94
95


96
@pytest.mark.parametrize("device", cpu_and_gpu())
97
@pytest.mark.parametrize(
98
99
    "func,method,fn_kwargs,match_kwargs",
    [
100
101
102
103
104
105
        (F.hflip, T.RandomHorizontalFlip, None, {}),
        (F.vflip, T.RandomVerticalFlip, None, {}),
        (F.invert, T.RandomInvert, None, {}),
        (F.posterize, T.RandomPosterize, {"bits": 4}, {}),
        (F.solarize, T.RandomSolarize, {"threshold": 192.0}, {}),
        (F.adjust_sharpness, T.RandomAdjustSharpness, {"sharpness_factor": 2.0}, {}),
106
107
108
109
110
111
112
113
        (
            F.autocontrast,
            T.RandomAutocontrast,
            None,
            {"test_exact_match": False, "agg_method": "max", "tol": (1 + 1e-5), "allowed_percentage_diff": 0.05},
        ),
        (F.equalize, T.RandomEqualize, None, {}),
    ],
114
)
115
@pytest.mark.parametrize("channels", [1, 3])
116
117
def test_random(func, method, device, channels, fn_kwargs, match_kwargs):
    _test_op(func, method, device, channels, fn_kwargs, fn_kwargs, **match_kwargs)
118

119

120
@pytest.mark.parametrize("seed", range(10))
121
122
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize("channels", [1, 3])
123
class TestColorJitter:
124
125
126
127
    @pytest.fixture(autouse=True)
    def set_random_seed(self, seed):
        torch.random.manual_seed(seed)

128
    @pytest.mark.parametrize("brightness", [0.1, 0.5, 1.0, 1.34, (0.3, 0.7), [0.4, 0.5]])
129
    def test_color_jitter_brightness(self, brightness, device, channels):
130
131
132
        tol = 1.0 + 1e-10
        meth_kwargs = {"brightness": brightness}
        _test_class_op(
133
134
135
136
137
138
139
            T.ColorJitter,
            meth_kwargs=meth_kwargs,
            test_exact_match=False,
            device=device,
            tol=tol,
            agg_method="max",
            channels=channels,
140
141
        )

142
    @pytest.mark.parametrize("contrast", [0.2, 0.5, 1.0, 1.5, (0.3, 0.7), [0.4, 0.5]])
143
    def test_color_jitter_contrast(self, contrast, device, channels):
144
145
146
        tol = 1.0 + 1e-10
        meth_kwargs = {"contrast": contrast}
        _test_class_op(
147
148
149
150
151
152
153
            T.ColorJitter,
            meth_kwargs=meth_kwargs,
            test_exact_match=False,
            device=device,
            tol=tol,
            agg_method="max",
            channels=channels,
154
155
        )

156
    @pytest.mark.parametrize("saturation", [0.5, 0.75, 1.0, 1.25, (0.3, 0.7), [0.3, 0.4]])
157
    def test_color_jitter_saturation(self, saturation, device, channels):
158
159
160
        tol = 1.0 + 1e-10
        meth_kwargs = {"saturation": saturation}
        _test_class_op(
161
162
163
164
165
166
167
            T.ColorJitter,
            meth_kwargs=meth_kwargs,
            test_exact_match=False,
            device=device,
            tol=tol,
            agg_method="max",
            channels=channels,
168
169
        )

170
    @pytest.mark.parametrize("hue", [0.2, 0.5, (-0.2, 0.3), [-0.4, 0.5]])
171
    def test_color_jitter_hue(self, hue, device, channels):
172
173
        meth_kwargs = {"hue": hue}
        _test_class_op(
174
175
176
177
178
179
180
            T.ColorJitter,
            meth_kwargs=meth_kwargs,
            test_exact_match=False,
            device=device,
            tol=16.1,
            agg_method="max",
            channels=channels,
181
182
        )

183
    def test_color_jitter_all(self, device, channels):
184
185
186
        # All 4 parameters together
        meth_kwargs = {"brightness": 0.2, "contrast": 0.2, "saturation": 0.2, "hue": 0.2}
        _test_class_op(
187
188
189
190
191
192
193
            T.ColorJitter,
            meth_kwargs=meth_kwargs,
            test_exact_match=False,
            device=device,
            tol=12.1,
            agg_method="max",
            channels=channels,
194
195
196
        )


197
198
199
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize("m", ["constant", "edge", "reflect", "symmetric"])
@pytest.mark.parametrize("mul", [1, -1])
200
201
202
203
def test_pad(m, mul, device):
    fill = 127 if m == "constant" else 0

    # Test functional.pad (PIL and Tensor) with padding as single int
204
    _test_functional_op(F.pad, fn_kwargs={"padding": mul * 2, "fill": fill, "padding_mode": m}, device=device)
205
    # Test functional.pad and transforms.Pad with padding as [int, ]
206
207
208
209
210
211
212
213
    fn_kwargs = meth_kwargs = {
        "padding": [
            mul * 2,
        ],
        "fill": fill,
        "padding_mode": m,
    }
    _test_op(F.pad, T.Pad, device=device, fn_kwargs=fn_kwargs, meth_kwargs=meth_kwargs)
214
215
    # Test functional.pad and transforms.Pad with padding as list
    fn_kwargs = meth_kwargs = {"padding": [mul * 4, 4], "fill": fill, "padding_mode": m}
216
    _test_op(F.pad, T.Pad, device=device, fn_kwargs=fn_kwargs, meth_kwargs=meth_kwargs)
217
218
    # Test functional.pad and transforms.Pad with padding as tuple
    fn_kwargs = meth_kwargs = {"padding": (mul * 2, 2, 2, mul * 2), "fill": fill, "padding_mode": m}
219
    _test_op(F.pad, T.Pad, device=device, fn_kwargs=fn_kwargs, meth_kwargs=meth_kwargs)
220
221


222
@pytest.mark.parametrize("device", cpu_and_gpu())
223
224
225
def test_crop(device):
    fn_kwargs = {"top": 2, "left": 3, "height": 4, "width": 5}
    # Test transforms.RandomCrop with size and padding as tuple
226
227
228
229
230
231
    meth_kwargs = {
        "size": (4, 5),
        "padding": (4, 4),
        "pad_if_needed": True,
    }
    _test_op(F.crop, T.RandomCrop, device=device, fn_kwargs=fn_kwargs, meth_kwargs=meth_kwargs)
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249

    # Test transforms.functional.crop including outside the image area
    fn_kwargs = {"top": -2, "left": 3, "height": 4, "width": 5}  # top
    _test_functional_op(F.crop, fn_kwargs=fn_kwargs, device=device)

    fn_kwargs = {"top": 1, "left": -3, "height": 4, "width": 5}  # left
    _test_functional_op(F.crop, fn_kwargs=fn_kwargs, device=device)

    fn_kwargs = {"top": 7, "left": 3, "height": 4, "width": 5}  # bottom
    _test_functional_op(F.crop, fn_kwargs=fn_kwargs, device=device)

    fn_kwargs = {"top": 3, "left": 8, "height": 4, "width": 5}  # right
    _test_functional_op(F.crop, fn_kwargs=fn_kwargs, device=device)

    fn_kwargs = {"top": -3, "left": -3, "height": 15, "width": 15}  # all
    _test_functional_op(F.crop, fn_kwargs=fn_kwargs, device=device)


250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize(
    "padding_config",
    [
        {"padding_mode": "constant", "fill": 0},
        {"padding_mode": "constant", "fill": 10},
        {"padding_mode": "constant", "fill": 20},
        {"padding_mode": "edge"},
        {"padding_mode": "reflect"},
    ],
)
@pytest.mark.parametrize(
    "size",
    [
        5,
        [
            5,
        ],
        [6, 6],
    ],
)
271
272
273
def test_crop_pad(size, padding_config, device):
    config = dict(padding_config)
    config["size"] = size
274
    _test_class_op(T.RandomCrop, device, meth_kwargs=config)
275
276


277
@pytest.mark.parametrize("device", cpu_and_gpu())
278
def test_center_crop(device, tmpdir):
279
    fn_kwargs = {"output_size": (4, 5)}
280
281
282
283
    meth_kwargs = {
        "size": (4, 5),
    }
    _test_op(F.center_crop, T.CenterCrop, device=device, fn_kwargs=fn_kwargs, meth_kwargs=meth_kwargs)
284
    fn_kwargs = {"output_size": (5,)}
285
    meth_kwargs = {"size": (5,)}
286
    _test_op(F.center_crop, T.CenterCrop, device=device, fn_kwargs=fn_kwargs, meth_kwargs=meth_kwargs)
287
288
289
290
291
292
293
    tensor = torch.randint(0, 256, (3, 10, 10), dtype=torch.uint8, device=device)
    # Test torchscript of transforms.CenterCrop with size as int
    f = T.CenterCrop(size=5)
    scripted_fn = torch.jit.script(f)
    scripted_fn(tensor)

    # Test torchscript of transforms.CenterCrop with size as [int, ]
294
295
296
297
298
    f = T.CenterCrop(
        size=[
            5,
        ]
    )
299
300
301
302
303
304
305
306
    scripted_fn = torch.jit.script(f)
    scripted_fn(tensor)

    # Test torchscript of transforms.CenterCrop with size as tuple
    f = T.CenterCrop(size=(6, 6))
    scripted_fn = torch.jit.script(f)
    scripted_fn(tensor)

307
    scripted_fn.save(os.path.join(tmpdir, "t_center_crop.pt"))
308
309


310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize(
    "fn, method, out_length",
    [
        # test_five_crop
        (F.five_crop, T.FiveCrop, 5),
        # test_ten_crop
        (F.ten_crop, T.TenCrop, 10),
    ],
)
@pytest.mark.parametrize(
    "size",
    [
        (5,),
        [
            5,
        ],
        (4, 5),
        [4, 5],
    ],
)
331
def test_x_crop(fn, method, out_length, size, device):
332
    meth_kwargs = fn_kwargs = {"size": size}
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
    scripted_fn = torch.jit.script(fn)

    tensor, pil_img = _create_data(height=20, width=20, device=device)
    transformed_t_list = fn(tensor, **fn_kwargs)
    transformed_p_list = fn(pil_img, **fn_kwargs)
    assert len(transformed_t_list) == len(transformed_p_list)
    assert len(transformed_t_list) == out_length
    for transformed_tensor, transformed_pil_img in zip(transformed_t_list, transformed_p_list):
        _assert_equal_tensor_to_pil(transformed_tensor, transformed_pil_img)

    transformed_t_list_script = scripted_fn(tensor.detach().clone(), **fn_kwargs)
    assert len(transformed_t_list) == len(transformed_t_list_script)
    assert len(transformed_t_list_script) == out_length
    for transformed_tensor, transformed_tensor_script in zip(transformed_t_list, transformed_t_list_script):
        assert_equal(transformed_tensor, transformed_tensor_script)

    # test for class interface
    fn = method(**meth_kwargs)
    scripted_fn = torch.jit.script(fn)
    output = scripted_fn(tensor)
    assert len(output) == len(transformed_t_list_script)

    # test on batch of tensors
    batch_tensors = _create_data_batch(height=23, width=34, channels=3, num_samples=4, device=device)
    torch.manual_seed(12)
    transformed_batch_list = fn(batch_tensors)

    for i in range(len(batch_tensors)):
        img_tensor = batch_tensors[i, ...]
        torch.manual_seed(12)
        transformed_img_list = fn(img_tensor)
        for transformed_img, transformed_batch in zip(transformed_img_list, transformed_batch_list):
            assert_equal(transformed_img, transformed_batch[i, ...])


368
@pytest.mark.parametrize("method", ["FiveCrop", "TenCrop"])
369
def test_x_crop_save(method, tmpdir):
370
371
372
373
374
    fn = getattr(T, method)(
        size=[
            5,
        ]
    )
375
    scripted_fn = torch.jit.script(fn)
376
    scripted_fn.save(os.path.join(tmpdir, f"t_op_list_{method}.pt"))
377
378
379


class TestResize:
380
    @pytest.mark.parametrize("size", [32, 34, 35, 36, 38])
381
382
383
384
385
386
387
388
389
390
391
    def test_resize_int(self, size):
        # TODO: Minimal check for bug-fix, improve this later
        x = torch.rand(3, 32, 46)
        t = T.Resize(size=size)
        y = t(x)
        # If size is an int, smaller edge of the image will be matched to this number.
        # i.e, if height > width, then image will be rescaled to (size * height / width, size).
        assert isinstance(y, torch.Tensor)
        assert y.shape[1] == size
        assert y.shape[2] == int(size * 46 / 32)

392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
    @pytest.mark.parametrize("device", cpu_and_gpu())
    @pytest.mark.parametrize("dt", [None, torch.float32, torch.float64])
    @pytest.mark.parametrize(
        "size",
        [
            [
                32,
            ],
            [32, 32],
            (32, 32),
            [34, 35],
        ],
    )
    @pytest.mark.parametrize("max_size", [None, 35, 1000])
    @pytest.mark.parametrize("interpolation", [BILINEAR, BICUBIC, NEAREST])
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
    def test_resize_scripted(self, dt, size, max_size, interpolation, device):
        tensor, _ = _create_data(height=34, width=36, device=device)
        batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device)

        if dt is not None:
            # This is a trivial cast to float of uint8 data to test all cases
            tensor = tensor.to(dt)
        if max_size is not None and len(size) != 1:
            pytest.xfail("with max_size, size must be a sequence with 2 elements")

        transform = T.Resize(size=size, interpolation=interpolation, max_size=max_size)
        s_transform = torch.jit.script(transform)
        _test_transform_vs_scripted(transform, s_transform, tensor)
        _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors)

422
    def test_resize_save(self, tmpdir):
423
424
425
426
427
        transform = T.Resize(
            size=[
                32,
            ]
        )
428
        s_transform = torch.jit.script(transform)
429
        s_transform.save(os.path.join(tmpdir, "t_resize.pt"))
430

431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
    @pytest.mark.parametrize("device", cpu_and_gpu())
    @pytest.mark.parametrize("scale", [(0.7, 1.2), [0.7, 1.2]])
    @pytest.mark.parametrize("ratio", [(0.75, 1.333), [0.75, 1.333]])
    @pytest.mark.parametrize(
        "size",
        [
            (32,),
            [
                44,
            ],
            [
                32,
            ],
            [32, 32],
            (32, 32),
            [44, 55],
        ],
    )
    @pytest.mark.parametrize("interpolation", [NEAREST, BILINEAR, BICUBIC])
450
451
452
453
454
455
    @pytest.mark.parametrize("antialias", [None, True, False])
    def test_resized_crop(self, scale, ratio, size, interpolation, antialias, device):

        if antialias and interpolation == NEAREST:
            pytest.skip("Can not resize if interpolation mode is NEAREST and antialias=True")

456
457
        tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=device)
        batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device)
458
459
460
        transform = T.RandomResizedCrop(
            size=size, scale=scale, ratio=ratio, interpolation=interpolation, antialias=antialias
        )
461
462
463
464
        s_transform = torch.jit.script(transform)
        _test_transform_vs_scripted(transform, s_transform, tensor)
        _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors)

465
    def test_resized_crop_save(self, tmpdir):
466
467
468
469
470
        transform = T.RandomResizedCrop(
            size=[
                32,
            ]
        )
471
        s_transform = torch.jit.script(transform)
472
        s_transform.save(os.path.join(tmpdir, "t_resized_crop.pt"))
473
474


475
476
477
478
479
480
481
482
483
484
def _test_random_affine_helper(device, **kwargs):
    tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=device)
    batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device)
    transform = T.RandomAffine(**kwargs)
    s_transform = torch.jit.script(transform)

    _test_transform_vs_scripted(transform, s_transform, tensor)
    _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors)


485
@pytest.mark.parametrize("device", cpu_and_gpu())
486
def test_random_affine(device, tmpdir):
487
488
    transform = T.RandomAffine(degrees=45.0)
    s_transform = torch.jit.script(transform)
489
    s_transform.save(os.path.join(tmpdir, "t_random_affine.pt"))
490
491


492
493
494
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize("interpolation", [NEAREST, BILINEAR])
@pytest.mark.parametrize("shear", [15, 10.0, (5.0, 10.0), [-15, 15], [-10.0, 10.0, -11.0, 11.0]])
495
496
497
498
def test_random_affine_shear(device, interpolation, shear):
    _test_random_affine_helper(device, degrees=0.0, interpolation=interpolation, shear=shear)


499
500
501
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize("interpolation", [NEAREST, BILINEAR])
@pytest.mark.parametrize("scale", [(0.7, 1.2), [0.7, 1.2]])
502
503
504
505
def test_random_affine_scale(device, interpolation, scale):
    _test_random_affine_helper(device, degrees=0.0, interpolation=interpolation, scale=scale)


506
507
508
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize("interpolation", [NEAREST, BILINEAR])
@pytest.mark.parametrize("translate", [(0.1, 0.2), [0.2, 0.1]])
509
510
511
512
def test_random_affine_translate(device, interpolation, translate):
    _test_random_affine_helper(device, degrees=0.0, interpolation=interpolation, translate=translate)


513
514
515
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize("interpolation", [NEAREST, BILINEAR])
@pytest.mark.parametrize("degrees", [45, 35.0, (-45, 45), [-90.0, 90.0]])
516
517
518
519
def test_random_affine_degrees(device, interpolation, degrees):
    _test_random_affine_helper(device, degrees=degrees, interpolation=interpolation)


520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize("interpolation", [NEAREST, BILINEAR])
@pytest.mark.parametrize(
    "fill",
    [
        85,
        (10, -10, 10),
        0.7,
        [0.0, 0.0, 0.0],
        [
            1,
        ],
        1,
    ],
)
535
536
537
538
def test_random_affine_fill(device, interpolation, fill):
    _test_random_affine_helper(device, degrees=0.0, interpolation=interpolation, fill=fill)


539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize("center", [(0, 0), [10, 10], None, (56, 44)])
@pytest.mark.parametrize("expand", [True, False])
@pytest.mark.parametrize("degrees", [45, 35.0, (-45, 45), [-90.0, 90.0]])
@pytest.mark.parametrize("interpolation", [NEAREST, BILINEAR])
@pytest.mark.parametrize(
    "fill",
    [
        85,
        (10, -10, 10),
        0.7,
        [0.0, 0.0, 0.0],
        [
            1,
        ],
        1,
    ],
)
557
558
559
560
def test_random_rotate(device, center, expand, degrees, interpolation, fill):
    tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=device)
    batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device)

561
    transform = T.RandomRotation(degrees=degrees, interpolation=interpolation, expand=expand, center=center, fill=fill)
562
563
564
565
566
567
    s_transform = torch.jit.script(transform)

    _test_transform_vs_scripted(transform, s_transform, tensor)
    _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors)


568
def test_random_rotate_save(tmpdir):
569
570
    transform = T.RandomRotation(degrees=45.0)
    s_transform = torch.jit.script(transform)
571
    s_transform.save(os.path.join(tmpdir, "t_random_rotate.pt"))
572
573


574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize("distortion_scale", np.linspace(0.1, 1.0, num=20))
@pytest.mark.parametrize("interpolation", [NEAREST, BILINEAR])
@pytest.mark.parametrize(
    "fill",
    [
        85,
        (10, -10, 10),
        0.7,
        [0.0, 0.0, 0.0],
        [
            1,
        ],
        1,
    ],
)
590
591
592
593
def test_random_perspective(device, distortion_scale, interpolation, fill):
    tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=device)
    batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device)

594
    transform = T.RandomPerspective(distortion_scale=distortion_scale, interpolation=interpolation, fill=fill)
595
596
597
598
599
600
    s_transform = torch.jit.script(transform)

    _test_transform_vs_scripted(transform, s_transform, tensor)
    _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors)


601
def test_random_perspective_save(tmpdir):
602
603
    transform = T.RandomPerspective()
    s_transform = torch.jit.script(transform)
604
    s_transform.save(os.path.join(tmpdir, "t_perspective.pt"))
605
606


607
608
609
610
611
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize(
    "Klass, meth_kwargs",
    [(T.Grayscale, {"num_output_channels": 1}), (T.Grayscale, {"num_output_channels": 3}), (T.RandomGrayscale, {})],
)
612
613
def test_to_grayscale(device, Klass, meth_kwargs):
    tol = 1.0 + 1e-10
614
    _test_class_op(Klass, meth_kwargs=meth_kwargs, test_exact_match=False, device=device, tol=tol, agg_method="max")
615
616


617
618
619
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize("in_dtype", int_dtypes() + float_dtypes())
@pytest.mark.parametrize("out_dtype", int_dtypes() + float_dtypes())
620
621
622
623
624
625
626
627
628
629
def test_convert_image_dtype(device, in_dtype, out_dtype):
    tensor, _ = _create_data(26, 34, device=device)
    batch_tensors = torch.rand(4, 3, 44, 56, device=device)

    in_tensor = tensor.to(in_dtype)
    in_batch_tensors = batch_tensors.to(in_dtype)

    fn = T.ConvertImageDtype(dtype=out_dtype)
    scripted_fn = torch.jit.script(fn)

630
631
632
    if (in_dtype == torch.float32 and out_dtype in (torch.int32, torch.int64)) or (
        in_dtype == torch.float64 and out_dtype == torch.int64
    ):
633
634
635
636
637
638
639
640
641
642
        with pytest.raises(RuntimeError, match=r"cannot be performed safely"):
            _test_transform_vs_scripted(fn, scripted_fn, in_tensor)
        with pytest.raises(RuntimeError, match=r"cannot be performed safely"):
            _test_transform_vs_scripted_on_batch(fn, scripted_fn, in_batch_tensors)
        return

    _test_transform_vs_scripted(fn, scripted_fn, in_tensor)
    _test_transform_vs_scripted_on_batch(fn, scripted_fn, in_batch_tensors)


643
def test_convert_image_dtype_save(tmpdir):
644
645
    fn = T.ConvertImageDtype(dtype=torch.uint8)
    scripted_fn = torch.jit.script(fn)
646
    scripted_fn.save(os.path.join(tmpdir, "t_convert_dtype.pt"))
647
648


649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize("policy", [policy for policy in T.AutoAugmentPolicy])
@pytest.mark.parametrize(
    "fill",
    [
        None,
        85,
        (10, -10, 10),
        0.7,
        [0.0, 0.0, 0.0],
        [
            1,
        ],
        1,
    ],
)
665
666
667
668
669
670
671
672
673
674
675
def test_autoaugment(device, policy, fill):
    tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=device)
    batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device)

    transform = T.AutoAugment(policy=policy, fill=fill)
    s_transform = torch.jit.script(transform)
    for _ in range(25):
        _test_transform_vs_scripted(transform, s_transform, tensor)
        _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors)


676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize("num_ops", [1, 2, 3])
@pytest.mark.parametrize("magnitude", [7, 9, 11])
@pytest.mark.parametrize(
    "fill",
    [
        None,
        85,
        (10, -10, 10),
        0.7,
        [0.0, 0.0, 0.0],
        [
            1,
        ],
        1,
    ],
)
693
694
695
696
697
698
699
700
701
702
703
def test_randaugment(device, num_ops, magnitude, fill):
    tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=device)
    batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device)

    transform = T.RandAugment(num_ops=num_ops, magnitude=magnitude, fill=fill)
    s_transform = torch.jit.script(transform)
    for _ in range(25):
        _test_transform_vs_scripted(transform, s_transform, tensor)
        _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors)


704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize(
    "fill",
    [
        None,
        85,
        (10, -10, 10),
        0.7,
        [0.0, 0.0, 0.0],
        [
            1,
        ],
        1,
    ],
)
719
720
721
722
723
724
725
726
727
728
729
def test_trivialaugmentwide(device, fill):
    tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=device)
    batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device)

    transform = T.TrivialAugmentWide(fill=fill)
    s_transform = torch.jit.script(transform)
    for _ in range(25):
        _test_transform_vs_scripted(transform, s_transform, tensor)
        _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors)


730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize(
    "fill",
    [
        None,
        85,
        (10, -10, 10),
        0.7,
        [0.0, 0.0, 0.0],
        [
            1,
        ],
        1,
    ],
)
def test_augmix(device, fill):
    tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=device)
    batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device)

    class DeterministicAugMix(T.AugMix):
        def _sample_dirichlet(self, params: torch.Tensor) -> torch.Tensor:
            # patch the method to ensure that the order of rand calls doesn't affect the outcome
            return params.softmax(dim=-1)

    transform = DeterministicAugMix(fill=fill)
    s_transform = torch.jit.script(transform)
    for _ in range(25):
        _test_transform_vs_scripted(transform, s_transform, tensor)
        _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors)


@pytest.mark.parametrize("augmentation", [T.AutoAugment, T.RandAugment, T.TrivialAugmentWide, T.AugMix])
762
763
def test_autoaugment_save(augmentation, tmpdir):
    transform = augmentation()
764
    s_transform = torch.jit.script(transform)
765
    s_transform.save(os.path.join(tmpdir, "t_autoaugment.pt"))
766
767


768
769
770
771
772
773
774
775
776
777
778
779
780
@pytest.mark.parametrize("interpolation", [F.InterpolationMode.NEAREST, F.InterpolationMode.BILINEAR])
@pytest.mark.parametrize("mode", ["X", "Y"])
def test_autoaugment__op_apply_shear(interpolation, mode):
    # We check that torchvision's implementation of shear is equivalent
    # to official CIFAR10 autoaugment implementation:
    # https://github.com/tensorflow/models/blob/885fda091c46c59d6c7bb5c7e760935eacc229da/research/autoaugment/augmentation_transforms.py#L273-L290
    image_size = 32

    def shear(pil_img, level, mode, resample):
        if mode == "X":
            matrix = (1, level, 0, 0, 1, 0)
        elif mode == "Y":
            matrix = (1, 0, 0, level, 1, 0)
781
        return pil_img.transform((image_size, image_size), _pil_constants.AFFINE, matrix, resample=resample)
782
783
784
785

    t_img, pil_img = _create_data(image_size, image_size)

    resample_pil = {
786
787
        F.InterpolationMode.NEAREST: _pil_constants.NEAREST,
        F.InterpolationMode.BILINEAR: _pil_constants.BILINEAR,
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
    }[interpolation]

    level = 0.3
    expected_out = shear(pil_img, level, mode=mode, resample=resample_pil)

    # Check pil output vs expected pil
    out = _apply_op(pil_img, op_name=f"Shear{mode}", magnitude=level, interpolation=interpolation, fill=0)
    assert out == expected_out

    if interpolation == F.InterpolationMode.BILINEAR:
        # We skip bilinear mode for tensors as
        # affine transformation results are not exactly the same
        # between tensors and pil images
        # MAE as around 1.40
        # Max Abs error can be 163 or 170
        return

    # Check tensor output vs expected pil
    out = _apply_op(t_img, op_name=f"Shear{mode}", magnitude=level, interpolation=interpolation, fill=0)
    _assert_approx_equal_tensor_to_pil(out, expected_out)


810
@pytest.mark.parametrize("device", cpu_and_gpu())
811
@pytest.mark.parametrize(
812
813
    "config",
    [{"value": 0.2}, {"value": "random"}, {"value": (0.2, 0.2, 0.2)}, {"value": "random", "ratio": (0.1, 0.2)}],
814
815
816
817
818
819
820
821
822
823
824
)
def test_random_erasing(device, config):
    tensor, _ = _create_data(24, 32, channels=3, device=device)
    batch_tensors = torch.rand(4, 3, 44, 56, device=device)

    fn = T.RandomErasing(**config)
    scripted_fn = torch.jit.script(fn)
    _test_transform_vs_scripted(fn, scripted_fn, tensor)
    _test_transform_vs_scripted_on_batch(fn, scripted_fn, batch_tensors)


825
def test_random_erasing_save(tmpdir):
826
827
    fn = T.RandomErasing(value=0.2)
    scripted_fn = torch.jit.script(fn)
828
    scripted_fn.save(os.path.join(tmpdir, "t_random_erasing.pt"))
829
830
831
832
833
834
835
836
837
838


def test_random_erasing_with_invalid_data():
    img = torch.rand(3, 60, 60)
    # Test Set 0: invalid value
    random_erasing = T.RandomErasing(value=(0.1, 0.2, 0.3, 0.4), p=1.0)
    with pytest.raises(ValueError, match="If value is a sequence, it should have either a single value or 3"):
        random_erasing(img)


839
@pytest.mark.parametrize("device", cpu_and_gpu())
840
def test_normalize(device, tmpdir):
841
842
843
844
845
846
847
848
849
850
851
852
853
854
    fn = T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    tensor, _ = _create_data(26, 34, device=device)

    with pytest.raises(TypeError, match="Input tensor should be a float tensor"):
        fn(tensor)

    batch_tensors = torch.rand(4, 3, 44, 56, device=device)
    tensor = tensor.to(dtype=torch.float32) / 255.0
    # test for class interface
    scripted_fn = torch.jit.script(fn)

    _test_transform_vs_scripted(fn, scripted_fn, tensor)
    _test_transform_vs_scripted_on_batch(fn, scripted_fn, batch_tensors)

855
    scripted_fn.save(os.path.join(tmpdir, "t_norm.pt"))
856
857


858
@pytest.mark.parametrize("device", cpu_and_gpu())
859
def test_linear_transformation(device, tmpdir):
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
    c, h, w = 3, 24, 32

    tensor, _ = _create_data(h, w, channels=c, device=device)

    matrix = torch.rand(c * h * w, c * h * w, device=device)
    mean_vector = torch.rand(c * h * w, device=device)

    fn = T.LinearTransformation(matrix, mean_vector)
    scripted_fn = torch.jit.script(fn)

    _test_transform_vs_scripted(fn, scripted_fn, tensor)

    batch_tensors = torch.rand(4, c, h, w, device=device)
    # We skip some tests from _test_transform_vs_scripted_on_batch as
    # results for scripted and non-scripted transformations are not exactly the same
    torch.manual_seed(12)
    transformed_batch = fn(batch_tensors)
    torch.manual_seed(12)
    s_transformed_batch = scripted_fn(batch_tensors)
    assert_equal(transformed_batch, s_transformed_batch)

881
    scripted_fn.save(os.path.join(tmpdir, "t_norm.pt"))
882
883


884
@pytest.mark.parametrize("device", cpu_and_gpu())
885
886
887
def test_compose(device):
    tensor, _ = _create_data(26, 34, device=device)
    tensor = tensor.to(dtype=torch.float32) / 255.0
888
889
890
891
892
893
    transforms = T.Compose(
        [
            T.CenterCrop(10),
            T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        ]
    )
894
895
896
897
898
899
900
    s_transforms = torch.nn.Sequential(*transforms.transforms)

    scripted_fn = torch.jit.script(s_transforms)
    torch.manual_seed(12)
    transformed_tensor = transforms(tensor)
    torch.manual_seed(12)
    transformed_tensor_script = scripted_fn(tensor)
901
    assert_equal(transformed_tensor, transformed_tensor_script, msg=f"{transforms}")
902

903
904
905
906
907
    t = T.Compose(
        [
            lambda x: x,
        ]
    )
908
    with pytest.raises(RuntimeError, match="cannot call a value of type 'Tensor'"):
909
910
911
        torch.jit.script(t)


912
@pytest.mark.parametrize("device", cpu_and_gpu())
913
914
915
916
def test_random_apply(device):
    tensor, _ = _create_data(26, 34, device=device)
    tensor = tensor.to(dtype=torch.float32) / 255.0

917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
    transforms = T.RandomApply(
        [
            T.RandomHorizontalFlip(),
            T.ColorJitter(),
        ],
        p=0.4,
    )
    s_transforms = T.RandomApply(
        torch.nn.ModuleList(
            [
                T.RandomHorizontalFlip(),
                T.ColorJitter(),
            ]
        ),
        p=0.4,
    )
933
934
935
936
937
938

    scripted_fn = torch.jit.script(s_transforms)
    torch.manual_seed(12)
    transformed_tensor = transforms(tensor)
    torch.manual_seed(12)
    transformed_tensor_script = scripted_fn(tensor)
939
    assert_equal(transformed_tensor, transformed_tensor_script, msg=f"{transforms}")
940
941
942
943

    if device == "cpu":
        # Can't check this twice, otherwise
        # "Can't redefine method: forward on class: __torch__.torchvision.transforms.transforms.RandomApply"
944
945
946
947
948
949
        transforms = T.RandomApply(
            [
                T.ColorJitter(),
            ],
            p=0.3,
        )
950
951
952
953
        with pytest.raises(RuntimeError, match="Module 'RandomApply' has no attribute 'transforms'"):
            torch.jit.script(transforms)


954
955
956
957
958
959
960
961
962
963
964
965
966
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize(
    "meth_kwargs",
    [
        {"kernel_size": 3, "sigma": 0.75},
        {"kernel_size": 23, "sigma": [0.1, 2.0]},
        {"kernel_size": 23, "sigma": (0.1, 2.0)},
        {"kernel_size": [3, 3], "sigma": (1.0, 1.0)},
        {"kernel_size": (3, 3), "sigma": (0.1, 2.0)},
        {"kernel_size": [23], "sigma": 0.75},
    ],
)
@pytest.mark.parametrize("channels", [1, 3])
967
def test_gaussian_blur(device, channels, meth_kwargs):
968
969
970
971
972
973
974
975
976
977
978
    if all(
        [
            device == "cuda",
            channels == 1,
            meth_kwargs["kernel_size"] in [23, [23]],
            torch.version.cuda == "11.3",
            sys.platform in ("win32", "cygwin"),
        ]
    ):
        pytest.skip("Fails on Windows, see https://github.com/pytorch/vision/issues/5464")

979
    tol = 1.0 + 1e-10
980
    torch.manual_seed(12)
981
    _test_class_op(
982
983
984
985
986
987
988
        T.GaussianBlur,
        meth_kwargs=meth_kwargs,
        channels=channels,
        test_exact_match=False,
        device=device,
        agg_method="max",
        tol=tol,
989
    )