test_transforms_tensor.py 31 KB
Newer Older
1
import os
2
3

import numpy as np
4
import pytest
5
import torch
Nicolas Hug's avatar
Nicolas Hug committed
6
7
8
9
10
11
12
13
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,
14
    cpu_and_gpu,
15
    assert_equal,
Nicolas Hug's avatar
Nicolas Hug committed
16
)
17
18
19
from torchvision import transforms as T
from torchvision.transforms import InterpolationMode
from torchvision.transforms import functional as F
20

21
NEAREST, BILINEAR, BICUBIC = InterpolationMode.NEAREST, InterpolationMode.BILINEAR, InterpolationMode.BICUBIC
22
23


24
25
26
27
28
29
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)
30

31

32
33
34
def _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors, msg=None):
    torch.manual_seed(12)
    transformed_batch = transform(batch_tensors)
35

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

42
43
44
    torch.manual_seed(12)
    s_transformed_batch = s_transform(batch_tensors)
    assert_equal(transformed_batch, s_transformed_batch, msg=msg)
45
46


47
def _test_functional_op(f, device, channels=3, fn_kwargs=None, test_exact_match=True, **match_kwargs):
48
    fn_kwargs = fn_kwargs or {}
49

50
    tensor, pil_img = _create_data(height=10, width=10, channels=channels, device=device)
51
52
53
54
55
56
    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
57
58


59
def _test_class_op(method, device, channels=3, meth_kwargs=None, test_exact_match=True, **match_kwargs):
60
61
    # TODO: change the name: it's not a method, it's a class.
    meth_kwargs = meth_kwargs or {}
62

63
64
65
    # test for class interface
    f = method(**meth_kwargs)
    scripted_fn = torch.jit.script(f)
66

67
    tensor, pil_img = _create_data(26, 34, channels, device=device)
68
69
70
71
72
73
74
75
76
    # 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)
77

78
79
80
81
    torch.manual_seed(12)
    transformed_tensor_script = scripted_fn(tensor)
    assert_equal(transformed_tensor, transformed_tensor_script)

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

    with get_tmp_dir() as tmp_dir:
        scripted_fn.save(os.path.join(tmp_dir, f"t_{method.__name__}.pt"))
87

88

89
90
91
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)
92
93


94
@pytest.mark.parametrize("device", cpu_and_gpu())
95
@pytest.mark.parametrize(
96
97
    "func,method,fn_kwargs,match_kwargs",
    [
98
99
100
101
102
103
        (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}, {}),
104
105
106
107
108
109
110
111
        (
            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, {}),
    ],
112
)
113
@pytest.mark.parametrize("channels", [1, 3])
114
115
def test_random(func, method, device, channels, fn_kwargs, match_kwargs):
    _test_op(func, method, device, channels, fn_kwargs, fn_kwargs, **match_kwargs)
116

117

118
119
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize("channels", [1, 3])
120
class TestColorJitter:
121
    @pytest.mark.parametrize("brightness", [0.1, 0.5, 1.0, 1.34, (0.3, 0.7), [0.4, 0.5]])
122
    def test_color_jitter_brightness(self, brightness, device, channels):
123
124
125
        tol = 1.0 + 1e-10
        meth_kwargs = {"brightness": brightness}
        _test_class_op(
126
127
128
129
130
131
132
            T.ColorJitter,
            meth_kwargs=meth_kwargs,
            test_exact_match=False,
            device=device,
            tol=tol,
            agg_method="max",
            channels=channels,
133
134
        )

135
    @pytest.mark.parametrize("contrast", [0.2, 0.5, 1.0, 1.5, (0.3, 0.7), [0.4, 0.5]])
136
    def test_color_jitter_contrast(self, contrast, device, channels):
137
138
139
        tol = 1.0 + 1e-10
        meth_kwargs = {"contrast": contrast}
        _test_class_op(
140
141
142
143
144
145
146
            T.ColorJitter,
            meth_kwargs=meth_kwargs,
            test_exact_match=False,
            device=device,
            tol=tol,
            agg_method="max",
            channels=channels,
147
148
        )

149
    @pytest.mark.parametrize("saturation", [0.5, 0.75, 1.0, 1.25, (0.3, 0.7), [0.3, 0.4]])
150
    def test_color_jitter_saturation(self, saturation, device, channels):
151
152
153
        tol = 1.0 + 1e-10
        meth_kwargs = {"saturation": saturation}
        _test_class_op(
154
155
156
157
158
159
160
            T.ColorJitter,
            meth_kwargs=meth_kwargs,
            test_exact_match=False,
            device=device,
            tol=tol,
            agg_method="max",
            channels=channels,
161
162
        )

163
    @pytest.mark.parametrize("hue", [0.2, 0.5, (-0.2, 0.3), [-0.4, 0.5]])
164
    def test_color_jitter_hue(self, hue, device, channels):
165
166
        meth_kwargs = {"hue": hue}
        _test_class_op(
167
168
169
170
171
172
173
            T.ColorJitter,
            meth_kwargs=meth_kwargs,
            test_exact_match=False,
            device=device,
            tol=16.1,
            agg_method="max",
            channels=channels,
174
175
        )

176
    def test_color_jitter_all(self, device, channels):
177
178
179
        # All 4 parameters together
        meth_kwargs = {"brightness": 0.2, "contrast": 0.2, "saturation": 0.2, "hue": 0.2}
        _test_class_op(
180
181
182
183
184
185
186
            T.ColorJitter,
            meth_kwargs=meth_kwargs,
            test_exact_match=False,
            device=device,
            tol=12.1,
            agg_method="max",
            channels=channels,
187
188
189
        )


190
191
192
@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize("m", ["constant", "edge", "reflect", "symmetric"])
@pytest.mark.parametrize("mul", [1, -1])
193
194
195
196
def test_pad(m, mul, device):
    fill = 127 if m == "constant" else 0

    # Test functional.pad (PIL and Tensor) with padding as single int
197
    _test_functional_op(F.pad, fn_kwargs={"padding": mul * 2, "fill": fill, "padding_mode": m}, device=device)
198
    # Test functional.pad and transforms.Pad with padding as [int, ]
199
200
201
202
203
204
205
206
    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)
207
208
    # Test functional.pad and transforms.Pad with padding as list
    fn_kwargs = meth_kwargs = {"padding": [mul * 4, 4], "fill": fill, "padding_mode": m}
209
    _test_op(F.pad, T.Pad, device=device, fn_kwargs=fn_kwargs, meth_kwargs=meth_kwargs)
210
211
    # 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}
212
    _test_op(F.pad, T.Pad, device=device, fn_kwargs=fn_kwargs, meth_kwargs=meth_kwargs)
213
214


215
@pytest.mark.parametrize("device", cpu_and_gpu())
216
217
218
def test_crop(device):
    fn_kwargs = {"top": 2, "left": 3, "height": 4, "width": 5}
    # Test transforms.RandomCrop with size and padding as tuple
219
220
221
222
223
224
    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)
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242

    # 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)


243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
@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],
    ],
)
264
265
266
def test_crop_pad(size, padding_config, device):
    config = dict(padding_config)
    config["size"] = size
267
    _test_class_op(T.RandomCrop, device, meth_kwargs=config)
268
269


270
@pytest.mark.parametrize("device", cpu_and_gpu())
271
def test_center_crop(device, tmpdir):
272
    fn_kwargs = {"output_size": (4, 5)}
273
274
275
276
    meth_kwargs = {
        "size": (4, 5),
    }
    _test_op(F.center_crop, T.CenterCrop, device=device, fn_kwargs=fn_kwargs, meth_kwargs=meth_kwargs)
277
    fn_kwargs = {"output_size": (5,)}
278
    meth_kwargs = {"size": (5,)}
279
    _test_op(F.center_crop, T.CenterCrop, device=device, fn_kwargs=fn_kwargs, meth_kwargs=meth_kwargs)
280
281
282
283
284
285
286
    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, ]
287
288
289
290
291
    f = T.CenterCrop(
        size=[
            5,
        ]
    )
292
293
294
295
296
297
298
299
    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)

300
    scripted_fn.save(os.path.join(tmpdir, "t_center_crop.pt"))
301
302


303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
@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],
    ],
)
324
def test_x_crop(fn, method, out_length, size, device):
325
    meth_kwargs = fn_kwargs = {"size": size}
326
327
328
329
330
331
332
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
    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, ...])


361
@pytest.mark.parametrize("method", ["FiveCrop", "TenCrop"])
362
def test_x_crop_save(method, tmpdir):
363
364
365
366
367
    fn = getattr(T, method)(
        size=[
            5,
        ]
    )
368
    scripted_fn = torch.jit.script(fn)
369
    scripted_fn.save(os.path.join(tmpdir, "t_op_list_{}.pt".format(method)))
370
371
372


class TestResize:
373
    @pytest.mark.parametrize("size", [32, 34, 35, 36, 38])
374
375
376
377
378
379
380
381
382
383
384
    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)

385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
    @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])
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
    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)

415
    def test_resize_save(self, tmpdir):
416
417
418
419
420
        transform = T.Resize(
            size=[
                32,
            ]
        )
421
        s_transform = torch.jit.script(transform)
422
        s_transform.save(os.path.join(tmpdir, "t_resize.pt"))
423

424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
    @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])
443
444
445
446
447
448
449
450
    def test_resized_crop(self, scale, ratio, size, interpolation, device):
        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.RandomResizedCrop(size=size, scale=scale, ratio=ratio, interpolation=interpolation)
        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)

451
    def test_resized_crop_save(self, tmpdir):
452
453
454
455
456
        transform = T.RandomResizedCrop(
            size=[
                32,
            ]
        )
457
        s_transform = torch.jit.script(transform)
458
        s_transform.save(os.path.join(tmpdir, "t_resized_crop.pt"))
459
460


461
462
463
464
465
466
467
468
469
470
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)


471
@pytest.mark.parametrize("device", cpu_and_gpu())
472
def test_random_affine(device, tmpdir):
473
474
    transform = T.RandomAffine(degrees=45.0)
    s_transform = torch.jit.script(transform)
475
    s_transform.save(os.path.join(tmpdir, "t_random_affine.pt"))
476
477


478
479
480
@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]])
481
482
483
484
def test_random_affine_shear(device, interpolation, shear):
    _test_random_affine_helper(device, degrees=0.0, interpolation=interpolation, shear=shear)


485
486
487
@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]])
488
489
490
491
def test_random_affine_scale(device, interpolation, scale):
    _test_random_affine_helper(device, degrees=0.0, interpolation=interpolation, scale=scale)


492
493
494
@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]])
495
496
497
498
def test_random_affine_translate(device, interpolation, translate):
    _test_random_affine_helper(device, degrees=0.0, interpolation=interpolation, translate=translate)


499
500
501
@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]])
502
503
504
505
def test_random_affine_degrees(device, interpolation, degrees):
    _test_random_affine_helper(device, degrees=degrees, interpolation=interpolation)


506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
@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,
    ],
)
521
522
523
524
def test_random_affine_fill(device, interpolation, fill):
    _test_random_affine_helper(device, degrees=0.0, interpolation=interpolation, fill=fill)


525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
@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,
    ],
)
543
544
545
546
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)

547
    transform = T.RandomRotation(degrees=degrees, interpolation=interpolation, expand=expand, center=center, fill=fill)
548
549
550
551
552
553
    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)


554
def test_random_rotate_save(tmpdir):
555
556
    transform = T.RandomRotation(degrees=45.0)
    s_transform = torch.jit.script(transform)
557
    s_transform.save(os.path.join(tmpdir, "t_random_rotate.pt"))
558
559


560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
@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,
    ],
)
576
577
578
579
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)

580
    transform = T.RandomPerspective(distortion_scale=distortion_scale, interpolation=interpolation, fill=fill)
581
582
583
584
585
586
    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)


587
def test_random_perspective_save(tmpdir):
588
589
    transform = T.RandomPerspective()
    s_transform = torch.jit.script(transform)
590
    s_transform.save(os.path.join(tmpdir, "t_perspective.pt"))
591
592


593
594
595
596
597
@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, {})],
)
598
599
def test_to_grayscale(device, Klass, meth_kwargs):
    tol = 1.0 + 1e-10
600
    _test_class_op(Klass, meth_kwargs=meth_kwargs, test_exact_match=False, device=device, tol=tol, agg_method="max")
601
602


603
604
605
@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())
606
607
608
609
610
611
612
613
614
615
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)

616
617
618
    if (in_dtype == torch.float32 and out_dtype in (torch.int32, torch.int64)) or (
        in_dtype == torch.float64 and out_dtype == torch.int64
    ):
619
620
621
622
623
624
625
626
627
628
        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)


629
def test_convert_image_dtype_save(tmpdir):
630
631
    fn = T.ConvertImageDtype(dtype=torch.uint8)
    scripted_fn = torch.jit.script(fn)
632
    scripted_fn.save(os.path.join(tmpdir, "t_convert_dtype.pt"))
633
634


635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
@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,
    ],
)
651
652
653
654
655
656
657
658
659
660
661
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)


662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
@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,
    ],
)
679
680
681
682
683
684
685
686
687
688
689
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)


690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
@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,
    ],
)
705
706
707
708
709
710
711
712
713
714
715
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)


716
@pytest.mark.parametrize("augmentation", [T.AutoAugment, T.RandAugment, T.TrivialAugmentWide])
717
718
def test_autoaugment_save(augmentation, tmpdir):
    transform = augmentation()
719
    s_transform = torch.jit.script(transform)
720
    s_transform.save(os.path.join(tmpdir, "t_autoaugment.pt"))
721
722


723
@pytest.mark.parametrize("device", cpu_and_gpu())
724
@pytest.mark.parametrize(
725
726
    "config",
    [{"value": 0.2}, {"value": "random"}, {"value": (0.2, 0.2, 0.2)}, {"value": "random", "ratio": (0.1, 0.2)}],
727
728
729
730
731
732
733
734
735
736
737
)
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)


738
def test_random_erasing_save(tmpdir):
739
740
    fn = T.RandomErasing(value=0.2)
    scripted_fn = torch.jit.script(fn)
741
    scripted_fn.save(os.path.join(tmpdir, "t_random_erasing.pt"))
742
743
744
745
746
747
748
749
750
751


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)


752
@pytest.mark.parametrize("device", cpu_and_gpu())
753
def test_normalize(device, tmpdir):
754
755
756
757
758
759
760
761
762
763
764
765
766
767
    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)

768
    scripted_fn.save(os.path.join(tmpdir, "t_norm.pt"))
769
770


771
@pytest.mark.parametrize("device", cpu_and_gpu())
772
def test_linear_transformation(device, tmpdir):
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
    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)

794
    scripted_fn.save(os.path.join(tmpdir, "t_norm.pt"))
795
796


797
@pytest.mark.parametrize("device", cpu_and_gpu())
798
799
800
def test_compose(device):
    tensor, _ = _create_data(26, 34, device=device)
    tensor = tensor.to(dtype=torch.float32) / 255.0
801
802
803
804
805
806
    transforms = T.Compose(
        [
            T.CenterCrop(10),
            T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        ]
    )
807
808
809
810
811
812
813
814
815
    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)
    assert_equal(transformed_tensor, transformed_tensor_script, msg="{}".format(transforms))

816
817
818
819
820
    t = T.Compose(
        [
            lambda x: x,
        ]
    )
821
    with pytest.raises(RuntimeError, match="cannot call a value of type 'Tensor'"):
822
823
824
        torch.jit.script(t)


825
@pytest.mark.parametrize("device", cpu_and_gpu())
826
827
828
829
def test_random_apply(device):
    tensor, _ = _create_data(26, 34, device=device)
    tensor = tensor.to(dtype=torch.float32) / 255.0

830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
    transforms = T.RandomApply(
        [
            T.RandomHorizontalFlip(),
            T.ColorJitter(),
        ],
        p=0.4,
    )
    s_transforms = T.RandomApply(
        torch.nn.ModuleList(
            [
                T.RandomHorizontalFlip(),
                T.ColorJitter(),
            ]
        ),
        p=0.4,
    )
846
847
848
849
850
851
852
853
854
855
856

    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)
    assert_equal(transformed_tensor, transformed_tensor_script, msg="{}".format(transforms))

    if device == "cpu":
        # Can't check this twice, otherwise
        # "Can't redefine method: forward on class: __torch__.torchvision.transforms.transforms.RandomApply"
857
858
859
860
861
862
        transforms = T.RandomApply(
            [
                T.ColorJitter(),
            ],
            p=0.3,
        )
863
864
865
866
        with pytest.raises(RuntimeError, match="Module 'RandomApply' has no attribute 'transforms'"):
            torch.jit.script(transforms)


867
868
869
870
871
872
873
874
875
876
877
878
879
@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])
880
def test_gaussian_blur(device, channels, meth_kwargs):
881
    tol = 1.0 + 1e-10
882
    torch.manual_seed(12)
883
    _test_class_op(
884
885
886
887
888
889
890
        T.GaussianBlur,
        meth_kwargs=meth_kwargs,
        channels=channels,
        test_exact_match=False,
        device=device,
        agg_method="max",
        tol=tol,
891
    )