transforms_v2_kernel_infos.py 49.5 KB
Newer Older
1
2
3
4
import functools
import itertools

import numpy as np
5
import PIL.Image
6
7
import pytest
import torch.testing
8
import torchvision.transforms.v2.functional as F
9
from torchvision import tv_tensors
10
11
from torchvision.transforms._functional_tensor import _max_value as get_max_value, _parse_pad_padding
from transforms_v2_legacy_utils import (
12
    ArgsKwargs,
13
    combinations_grid,
14
    DEFAULT_PORTRAIT_SPATIAL_SIZE,
15
16
    get_num_channels,
    ImageLoader,
17
    InfoBase,
18
    make_bounding_box_loader,
19
    make_bounding_box_loaders,
20
    make_detection_mask_loader,
21
22
    make_image_loader,
    make_image_loaders,
23
    make_image_loaders_for_interpolation,
24
    make_mask_loaders,
25
    make_video_loader,
26
    make_video_loaders,
27
28
    mark_framework_limitation,
    TestMark,
29
)
30
31
32
33

__all__ = ["KernelInfo", "KERNEL_INFOS"]


34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
class KernelInfo(InfoBase):
    def __init__(
        self,
        kernel,
        *,
        # Defaults to `kernel.__name__`. Should be set if the function is exposed under a different name
        # TODO: This can probably be removed after roll-out since we shouldn't have any aliasing then
        kernel_name=None,
        # Most common tests use these inputs to check the kernel. As such it should cover all valid code paths, but
        # should not include extensive parameter combinations to keep to overall test count moderate.
        sample_inputs_fn,
        # This function should mirror the kernel. It should have the same signature as the `kernel` and as such also
        # take tensors as inputs. Any conversion into another object type, e.g. PIL images or numpy arrays, should
        # happen inside the function. It should return a tensor or to be more precise an object that can be compared to
        # a tensor by `assert_close`. If omitted, no reference test will be performed.
        reference_fn=None,
        # These inputs are only used for the reference tests and thus can be comprehensive with regard to the parameter
        # values to be tested. If not specified, `sample_inputs_fn` will be used.
        reference_inputs_fn=None,
53
        # If true-ish, triggers a test that checks the kernel for consistency between uint8 and float32 inputs with the
54
        # reference inputs. This is usually used whenever we use a PIL kernel as reference.
55
56
57
58
        # Can be a callable in which case it will be called with `other_args, kwargs`. It should return the same
        # structure, but with adapted parameters. This is useful in case a parameter value is closely tied to the input
        # dtype.
        float32_vs_uint8=False,
59
60
61
        # Some kernels don't have dispatchers that would handle logging the usage. Thus, the kernel has to do it
        # manually. If set, triggers a test that makes sure this happens.
        logs_usage=False,
62
63
64
65
66
67
68
69
70
71
        # See InfoBase
        test_marks=None,
        # See InfoBase
        closeness_kwargs=None,
    ):
        super().__init__(id=kernel_name or kernel.__name__, test_marks=test_marks, closeness_kwargs=closeness_kwargs)
        self.kernel = kernel
        self.sample_inputs_fn = sample_inputs_fn
        self.reference_fn = reference_fn
        self.reference_inputs_fn = reference_inputs_fn
72

73
74
75
        if float32_vs_uint8 and not callable(float32_vs_uint8):
            float32_vs_uint8 = lambda other_args, kwargs: (other_args, kwargs)  # noqa: E731
        self.float32_vs_uint8 = float32_vs_uint8
76
        self.logs_usage = logs_usage
77
78


79
def pixel_difference_closeness_kwargs(uint8_atol, *, dtype=torch.uint8, mae=False):
80
    return dict(atol=uint8_atol / 255 * get_max_value(dtype), rtol=0, mae=mae)
81
82
83
84


def cuda_vs_cpu_pixel_difference(atol=1):
    return {
85
        (("TestKernels", "test_cuda_vs_cpu"), dtype, "cuda"): pixel_difference_closeness_kwargs(atol, dtype=dtype)
86
87
88
89
        for dtype in [torch.uint8, torch.float32]
    }


90
def pil_reference_pixel_difference(atol=1, mae=False):
91
    return {
92
        (("TestKernels", "test_against_reference"), torch.uint8, "cpu"): pixel_difference_closeness_kwargs(
93
            atol, mae=mae
94
95
96
97
        )
    }


98
def float32_vs_uint8_pixel_difference(atol=1, mae=False):
99
100
101
102
103
    return {
        (
            ("TestKernels", "test_float32_vs_uint8"),
            torch.float32,
            "cpu",
104
        ): pixel_difference_closeness_kwargs(atol, dtype=torch.float32, mae=mae)
105
    }
106

107

108
def scripted_vs_eager_float64_tolerances(device, atol=1e-6, rtol=1e-6):
109
110
111
112
113
    return {
        (("TestKernels", "test_scripted_vs_eager"), torch.float64, device): {"atol": atol, "rtol": rtol, "mae": False},
    }


114
115
def pil_reference_wrapper(pil_kernel):
    @functools.wraps(pil_kernel)
116
117
118
119
    def wrapper(input_tensor, *other_args, **kwargs):
        if input_tensor.dtype != torch.uint8:
            raise pytest.UsageError(f"Can only test uint8 tensor images against PIL, but input is {input_tensor.dtype}")
        if input_tensor.ndim > 3:
120
            raise pytest.UsageError(
121
                f"Can only test single tensor images against PIL, but input has shape {input_tensor.shape}"
122
123
            )

124
        input_pil = F.to_pil_image(input_tensor)
125
126
127
128
        output_pil = pil_kernel(input_pil, *other_args, **kwargs)
        if not isinstance(output_pil, PIL.Image.Image):
            return output_pil

129
        output_tensor = F.to_image(output_pil)
130
131
132
133
134
135
136
137

        # 2D mask shenanigans
        if output_tensor.ndim == 2 and input_tensor.ndim == 3:
            output_tensor = output_tensor.unsqueeze(0)
        elif output_tensor.ndim == 3 and input_tensor.ndim == 2:
            output_tensor = output_tensor.squeeze(0)

        return output_tensor
138
139
140
141

    return wrapper


142
143
144
145
def xfail_jit(reason, *, condition=None):
    return TestMark(("TestKernels", "test_scripted_vs_eager"), pytest.mark.xfail(reason=reason), condition=condition)


146
def xfail_jit_python_scalar_arg(name, *, reason=None):
147
148
    return xfail_jit(
        reason or f"Python scalar int or float for `{name}` is not supported when scripting",
149
150
151
152
        condition=lambda args_kwargs: isinstance(args_kwargs.kwargs.get(name), (int, float)),
    )


153
154
155
KERNEL_INFOS = []


156
def get_fills(*, num_channels, dtype):
157
158
    yield None

159
160
161
162
    int_value = get_max_value(dtype)
    float_value = int_value / 2
    yield int_value
    yield float_value
163

164
165
166
    for vector_type in [list, tuple]:
        yield vector_type([int_value])
        yield vector_type([float_value])
167

168
169
170
        if num_channels > 1:
            yield vector_type(float_value * c / 10 for c in range(num_channels))
            yield vector_type(int_value if c % 2 == 0 else 0 for c in range(num_channels))
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185


def float32_vs_uint8_fill_adapter(other_args, kwargs):
    fill = kwargs.get("fill")
    if fill is None:
        return other_args, kwargs

    if isinstance(fill, (int, float)):
        fill /= 255
    else:
        fill = type(fill)(fill_ / 255 for fill_ in fill)

    return other_args, dict(kwargs, fill=fill)


Philip Meier's avatar
Philip Meier committed
186
187
def reference_affine_bounding_boxes_helper(bounding_boxes, *, format, canvas_size, affine_matrix):
    def transform(bbox, affine_matrix_, format_, canvas_size_):
188
189
        # Go to float before converting to prevent precision loss in case of CXCYWH -> XYXY and W or H is 1
        in_dtype = bbox.dtype
190
191
        if not torch.is_floating_point(bbox):
            bbox = bbox.float()
Nicolas Hug's avatar
Nicolas Hug committed
192
        bbox_xyxy = F.convert_bounding_box_format(
193
194
            bbox.as_subclass(torch.Tensor),
            old_format=format_,
195
            new_format=tv_tensors.BoundingBoxFormat.XYXY,
196
            inplace=True,
197
        )
198
199
200
201
202
203
204
205
        points = np.array(
            [
                [bbox_xyxy[0].item(), bbox_xyxy[1].item(), 1.0],
                [bbox_xyxy[2].item(), bbox_xyxy[1].item(), 1.0],
                [bbox_xyxy[0].item(), bbox_xyxy[3].item(), 1.0],
                [bbox_xyxy[2].item(), bbox_xyxy[3].item(), 1.0],
            ]
        )
206
        transformed_points = np.matmul(points, affine_matrix_.T)
207
208
        out_bbox = torch.tensor(
            [
209
210
211
212
                np.min(transformed_points[:, 0]).item(),
                np.min(transformed_points[:, 1]).item(),
                np.max(transformed_points[:, 0]).item(),
                np.max(transformed_points[:, 1]).item(),
213
            ],
214
            dtype=bbox_xyxy.dtype,
215
        )
Nicolas Hug's avatar
Nicolas Hug committed
216
        out_bbox = F.convert_bounding_box_format(
217
            out_bbox, old_format=tv_tensors.BoundingBoxFormat.XYXY, new_format=format_, inplace=True
218
        )
219
        # It is important to clamp before casting, especially for CXCYWH format, dtype=int64
Philip Meier's avatar
Philip Meier committed
220
        out_bbox = F.clamp_bounding_boxes(out_bbox, format=format_, canvas_size=canvas_size_)
221
222
        out_bbox = out_bbox.to(dtype=in_dtype)
        return out_bbox
223

224
225
226
    return torch.stack(
        [transform(b, affine_matrix, format, canvas_size) for b in bounding_boxes.reshape(-1, 4).unbind()]
    ).reshape(bounding_boxes.shape)
227
228


229
230
231
232
233
234
235
_PAD_PARAMS = combinations_grid(
    padding=[[1], [1, 1], [1, 1, 2, 2]],
    padding_mode=["constant", "symmetric", "edge", "reflect"],
)


def sample_inputs_pad_image_tensor():
236
    make_pad_image_loaders = functools.partial(
237
        make_image_loaders, sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], color_spaces=["RGB"], dtypes=[torch.float32]
238
239
240
241
242
243
244
245
246
    )

    for image_loader, padding in itertools.product(
        make_pad_image_loaders(),
        [1, (1,), (1, 2), (1, 2, 3, 4), [1], [1, 2], [1, 2, 3, 4]],
    ):
        yield ArgsKwargs(image_loader, padding=padding)

    for image_loader in make_pad_image_loaders():
247
        for fill in get_fills(num_channels=image_loader.num_channels, dtype=image_loader.dtype):
248
249
250
251
252
253
254
255
256
257
258
259
260
            yield ArgsKwargs(image_loader, padding=[1], fill=fill)

    for image_loader, padding_mode in itertools.product(
        # We branch for non-constant padding and integer inputs
        make_pad_image_loaders(dtypes=[torch.uint8]),
        ["constant", "symmetric", "edge", "reflect"],
    ):
        yield ArgsKwargs(image_loader, padding=[1], padding_mode=padding_mode)

    # `torch.nn.functional.pad` does not support symmetric padding, and thus we have a custom implementation. Besides
    # negative padding, this is already handled by the inputs above.
    for image_loader in make_pad_image_loaders():
        yield ArgsKwargs(image_loader, padding=[-1], padding_mode="symmetric")
261
262
263


def reference_inputs_pad_image_tensor():
264
265
266
267
268
269
270
    for image_loader, params in itertools.product(
        make_image_loaders(extra_dims=[()], dtypes=[torch.uint8]), _PAD_PARAMS
    ):
        for fill in get_fills(
            num_channels=image_loader.num_channels,
            dtype=image_loader.dtype,
        ):
271
272
273
274
            # FIXME: PIL kernel doesn't support sequences of length 1 if the number of channels is larger. Shouldn't it?
            if isinstance(fill, (list, tuple)):
                continue

275
276
277
            yield ArgsKwargs(image_loader, fill=fill, **params)


278
279
def sample_inputs_pad_bounding_boxes():
    for bounding_boxes_loader, padding in itertools.product(
280
281
        make_bounding_box_loaders(), [1, (1,), (1, 2), (1, 2, 3, 4), [1], [1, 2], [1, 2, 3, 4]]
    ):
282
        yield ArgsKwargs(
283
284
            bounding_boxes_loader,
            format=bounding_boxes_loader.format,
Philip Meier's avatar
Philip Meier committed
285
            canvas_size=bounding_boxes_loader.canvas_size,
286
287
            padding=padding,
            padding_mode="constant",
288
        )
289
290
291


def sample_inputs_pad_mask():
292
    for mask_loader in make_mask_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], num_categories=[10], num_objects=[5]):
293
        yield ArgsKwargs(mask_loader, padding=[1])
294
295
296


def reference_inputs_pad_mask():
297
298
299
300
    for mask_loader, fill, params in itertools.product(
        make_mask_loaders(num_objects=[1], extra_dims=[()]), [None, 127], _PAD_PARAMS
    ):
        yield ArgsKwargs(mask_loader, fill=fill, **params)
301
302


303
def sample_inputs_pad_video():
304
    for video_loader in make_video_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], num_frames=[3]):
305
306
307
        yield ArgsKwargs(video_loader, padding=[1])


Philip Meier's avatar
Philip Meier committed
308
def reference_pad_bounding_boxes(bounding_boxes, *, format, canvas_size, padding, padding_mode):
309
310
311
312
313
314
315
316

    left, right, top, bottom = _parse_pad_padding(padding)

    affine_matrix = np.array(
        [
            [1, 0, left],
            [0, 1, top],
        ],
317
        dtype="float64" if bounding_boxes.dtype == torch.float64 else "float32",
318
319
    )

Philip Meier's avatar
Philip Meier committed
320
321
    height = canvas_size[0] + top + bottom
    width = canvas_size[1] + left + right
322

323
    expected_bboxes = reference_affine_bounding_boxes_helper(
Philip Meier's avatar
Philip Meier committed
324
        bounding_boxes, format=format, canvas_size=(height, width), affine_matrix=affine_matrix
325
    )
326
327
328
    return expected_bboxes, (height, width)


329
330
def reference_inputs_pad_bounding_boxes():
    for bounding_boxes_loader, padding in itertools.product(
331
332
333
        make_bounding_box_loaders(extra_dims=((), (4,))), [1, (1,), (1, 2), (1, 2, 3, 4), [1], [1, 2], [1, 2, 3, 4]]
    ):
        yield ArgsKwargs(
334
335
            bounding_boxes_loader,
            format=bounding_boxes_loader.format,
Philip Meier's avatar
Philip Meier committed
336
            canvas_size=bounding_boxes_loader.canvas_size,
337
338
339
340
341
            padding=padding,
            padding_mode="constant",
        )


342
343
344
345
346
347
348
349
350
351
def pad_xfail_jit_fill_condition(args_kwargs):
    fill = args_kwargs.kwargs.get("fill")
    if not isinstance(fill, (list, tuple)):
        return False
    elif isinstance(fill, tuple):
        return True
    else:  # isinstance(fill, list):
        return all(isinstance(f, int) for f in fill)


352
353
354
KERNEL_INFOS.extend(
    [
        KernelInfo(
355
            F.pad_image,
356
            sample_inputs_fn=sample_inputs_pad_image_tensor,
357
            reference_fn=pil_reference_wrapper(F._pad_image_pil),
358
            reference_inputs_fn=reference_inputs_pad_image_tensor,
359
360
            float32_vs_uint8=float32_vs_uint8_fill_adapter,
            closeness_kwargs=float32_vs_uint8_pixel_difference(),
361
            test_marks=[
362
363
364
365
                xfail_jit_python_scalar_arg("padding"),
                xfail_jit(
                    "F.pad only supports vector fills for list of floats", condition=pad_xfail_jit_fill_condition
                ),
366
            ],
367
368
        ),
        KernelInfo(
369
370
371
372
            F.pad_bounding_boxes,
            sample_inputs_fn=sample_inputs_pad_bounding_boxes,
            reference_fn=reference_pad_bounding_boxes,
            reference_inputs_fn=reference_inputs_pad_bounding_boxes,
373
            test_marks=[
374
                xfail_jit_python_scalar_arg("padding"),
375
            ],
376
377
378
379
        ),
        KernelInfo(
            F.pad_mask,
            sample_inputs_fn=sample_inputs_pad_mask,
380
            reference_fn=pil_reference_wrapper(F._pad_image_pil),
381
            reference_inputs_fn=reference_inputs_pad_mask,
382
            float32_vs_uint8=float32_vs_uint8_fill_adapter,
383
        ),
384
385
386
387
        KernelInfo(
            F.pad_video,
            sample_inputs_fn=sample_inputs_pad_video,
        ),
388
389
390
391
392
393
394
    ]
)

_PERSPECTIVE_COEFFS = [
    [1.2405, 0.1772, -6.9113, 0.0463, 1.251, -5.235, 0.00013, 0.0018],
    [0.7366, -0.11724, 1.45775, -0.15012, 0.73406, 2.6019, -0.0072, -0.0063],
]
395
396
_STARTPOINTS = [[0, 1], [2, 3], [4, 5], [6, 7]]
_ENDPOINTS = [[9, 8], [7, 6], [5, 4], [3, 2]]
397
398
399


def sample_inputs_perspective_image_tensor():
400
    for image_loader in make_image_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE]):
401
        for fill in get_fills(num_channels=image_loader.num_channels, dtype=image_loader.dtype):
402
403
404
405
406
            yield ArgsKwargs(
                image_loader, startpoints=None, endpoints=None, fill=fill, coefficients=_PERSPECTIVE_COEFFS[0]
            )

    yield ArgsKwargs(make_image_loader(), startpoints=_STARTPOINTS, endpoints=_ENDPOINTS)
407
408
409


def reference_inputs_perspective_image_tensor():
410
411
412
413
414
415
416
    for image_loader, coefficients, interpolation in itertools.product(
        make_image_loaders_for_interpolation(),
        _PERSPECTIVE_COEFFS,
        [
            F.InterpolationMode.NEAREST,
            F.InterpolationMode.BILINEAR,
        ],
417
418
    ):
        for fill in get_fills(num_channels=image_loader.num_channels, dtype=image_loader.dtype):
419
420
421
422
            # FIXME: PIL kernel doesn't support sequences of length 1 if the number of channels is larger. Shouldn't it?
            if isinstance(fill, (list, tuple)):
                continue

423
424
425
426
427
428
429
430
            yield ArgsKwargs(
                image_loader,
                startpoints=None,
                endpoints=None,
                interpolation=interpolation,
                fill=fill,
                coefficients=coefficients,
            )
431
432


433
434
def sample_inputs_perspective_bounding_boxes():
    for bounding_boxes_loader in make_bounding_box_loaders():
435
        yield ArgsKwargs(
436
437
            bounding_boxes_loader,
            format=bounding_boxes_loader.format,
Philip Meier's avatar
Philip Meier committed
438
            canvas_size=bounding_boxes_loader.canvas_size,
439
440
441
            startpoints=None,
            endpoints=None,
            coefficients=_PERSPECTIVE_COEFFS[0],
442
443
        )

444
    format = tv_tensors.BoundingBoxFormat.XYXY
445
    loader = make_bounding_box_loader(format=format)
446
    yield ArgsKwargs(
Philip Meier's avatar
Philip Meier committed
447
        loader, format=format, canvas_size=loader.canvas_size, startpoints=_STARTPOINTS, endpoints=_ENDPOINTS
448
449
    )

450
451

def sample_inputs_perspective_mask():
452
    for mask_loader in make_mask_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE]):
453
454
455
        yield ArgsKwargs(mask_loader, startpoints=None, endpoints=None, coefficients=_PERSPECTIVE_COEFFS[0])

    yield ArgsKwargs(make_detection_mask_loader(), startpoints=_STARTPOINTS, endpoints=_ENDPOINTS)
456
457
458
459
460
461


def reference_inputs_perspective_mask():
    for mask_loader, perspective_coeffs in itertools.product(
        make_mask_loaders(extra_dims=[()], num_objects=[1]), _PERSPECTIVE_COEFFS
    ):
462
        yield ArgsKwargs(mask_loader, startpoints=None, endpoints=None, coefficients=perspective_coeffs)
463
464


465
def sample_inputs_perspective_video():
466
    for video_loader in make_video_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], num_frames=[3]):
467
468
469
        yield ArgsKwargs(video_loader, startpoints=None, endpoints=None, coefficients=_PERSPECTIVE_COEFFS[0])

    yield ArgsKwargs(make_video_loader(), startpoints=_STARTPOINTS, endpoints=_ENDPOINTS)
470
471


472
473
474
KERNEL_INFOS.extend(
    [
        KernelInfo(
475
            F.perspective_image,
476
            sample_inputs_fn=sample_inputs_perspective_image_tensor,
477
            reference_fn=pil_reference_wrapper(F._perspective_image_pil),
478
            reference_inputs_fn=reference_inputs_perspective_image_tensor,
479
            float32_vs_uint8=float32_vs_uint8_fill_adapter,
480
            closeness_kwargs={
481
                **pil_reference_pixel_difference(2, mae=True),
482
483
                **cuda_vs_cpu_pixel_difference(),
                **float32_vs_uint8_pixel_difference(),
484
485
                **scripted_vs_eager_float64_tolerances("cpu", atol=1e-5, rtol=1e-5),
                **scripted_vs_eager_float64_tolerances("cuda", atol=1e-5, rtol=1e-5),
486
            },
487
            test_marks=[xfail_jit_python_scalar_arg("fill")],
488
489
        ),
        KernelInfo(
490
491
            F.perspective_bounding_boxes,
            sample_inputs_fn=sample_inputs_perspective_bounding_boxes,
492
493
494
495
            closeness_kwargs={
                **scripted_vs_eager_float64_tolerances("cpu", atol=1e-6, rtol=1e-6),
                **scripted_vs_eager_float64_tolerances("cuda", atol=1e-6, rtol=1e-6),
            },
496
497
498
499
        ),
        KernelInfo(
            F.perspective_mask,
            sample_inputs_fn=sample_inputs_perspective_mask,
500
            reference_fn=pil_reference_wrapper(F._perspective_image_pil),
501
            reference_inputs_fn=reference_inputs_perspective_mask,
502
503
504
505
            float32_vs_uint8=True,
            closeness_kwargs={
                (("TestKernels", "test_against_reference"), torch.uint8, "cpu"): dict(atol=10, rtol=0),
            },
506
507
508
509
        ),
        KernelInfo(
            F.perspective_video,
            sample_inputs_fn=sample_inputs_perspective_video,
510
511
            closeness_kwargs={
                **cuda_vs_cpu_pixel_difference(),
512
513
                **scripted_vs_eager_float64_tolerances("cpu", atol=1e-5, rtol=1e-5),
                **scripted_vs_eager_float64_tolerances("cuda", atol=1e-5, rtol=1e-5),
514
            },
515
516
517
518
519
        ),
    ]
)


Philip Meier's avatar
Philip Meier committed
520
521
def _get_elastic_displacement(canvas_size):
    return torch.rand(1, *canvas_size, 2)
522
523
524


def sample_inputs_elastic_image_tensor():
525
    for image_loader in make_image_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE]):
Philip Meier's avatar
Philip Meier committed
526
        displacement = _get_elastic_displacement(image_loader.canvas_size)
527
        for fill in get_fills(num_channels=image_loader.num_channels, dtype=image_loader.dtype):
528
529
530
531
532
            yield ArgsKwargs(image_loader, displacement=displacement, fill=fill)


def reference_inputs_elastic_image_tensor():
    for image_loader, interpolation in itertools.product(
533
        make_image_loaders_for_interpolation(),
534
535
536
537
538
539
        [
            F.InterpolationMode.NEAREST,
            F.InterpolationMode.BILINEAR,
            F.InterpolationMode.BICUBIC,
        ],
    ):
Philip Meier's avatar
Philip Meier committed
540
        displacement = _get_elastic_displacement(image_loader.canvas_size)
541
        for fill in get_fills(num_channels=image_loader.num_channels, dtype=image_loader.dtype):
542
543
544
            yield ArgsKwargs(image_loader, interpolation=interpolation, displacement=displacement, fill=fill)


545
546
def sample_inputs_elastic_bounding_boxes():
    for bounding_boxes_loader in make_bounding_box_loaders():
Philip Meier's avatar
Philip Meier committed
547
        displacement = _get_elastic_displacement(bounding_boxes_loader.canvas_size)
548
        yield ArgsKwargs(
549
550
            bounding_boxes_loader,
            format=bounding_boxes_loader.format,
Philip Meier's avatar
Philip Meier committed
551
            canvas_size=bounding_boxes_loader.canvas_size,
552
553
554
555
556
            displacement=displacement,
        )


def sample_inputs_elastic_mask():
557
    for mask_loader in make_mask_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE]):
558
559
560
561
        displacement = _get_elastic_displacement(mask_loader.shape[-2:])
        yield ArgsKwargs(mask_loader, displacement=displacement)


562
def sample_inputs_elastic_video():
563
    for video_loader in make_video_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], num_frames=[3]):
564
565
566
567
        displacement = _get_elastic_displacement(video_loader.shape[-2:])
        yield ArgsKwargs(video_loader, displacement=displacement)


568
569
570
KERNEL_INFOS.extend(
    [
        KernelInfo(
571
            F.elastic_image,
572
573
            sample_inputs_fn=sample_inputs_elastic_image_tensor,
            reference_inputs_fn=reference_inputs_elastic_image_tensor,
574
            float32_vs_uint8=float32_vs_uint8_fill_adapter,
575
            closeness_kwargs={
576
                **float32_vs_uint8_pixel_difference(6, mae=True),
577
578
                **cuda_vs_cpu_pixel_difference(),
            },
579
            test_marks=[xfail_jit_python_scalar_arg("fill")],
580
581
        ),
        KernelInfo(
582
583
            F.elastic_bounding_boxes,
            sample_inputs_fn=sample_inputs_elastic_bounding_boxes,
584
585
586
587
        ),
        KernelInfo(
            F.elastic_mask,
            sample_inputs_fn=sample_inputs_elastic_mask,
588
589
590
591
        ),
        KernelInfo(
            F.elastic_video,
            sample_inputs_fn=sample_inputs_elastic_video,
592
            closeness_kwargs=cuda_vs_cpu_pixel_difference(),
593
594
595
596
597
        ),
    ]
)


598
_CENTER_CROP_SPATIAL_SIZES = [(16, 16), (7, 33), (31, 9)]
599
_CENTER_CROP_OUTPUT_SIZES = [[4, 3], [42, 70], [4], 3, (5, 2), (6,)]
600
601
602
603


def sample_inputs_center_crop_image_tensor():
    for image_loader, output_size in itertools.product(
604
        make_image_loaders(sizes=[(16, 17)], color_spaces=["RGB"], dtypes=[torch.float32]),
605
606
607
608
609
610
        [
            # valid `output_size` types for which cropping is applied to both dimensions
            *[5, (4,), (2, 3), [6], [3, 2]],
            # `output_size`'s for which at least one dimension needs to be padded
            *[[4, 18], [17, 5], [17, 18]],
        ],
611
612
613
614
615
616
    ):
        yield ArgsKwargs(image_loader, output_size=output_size)


def reference_inputs_center_crop_image_tensor():
    for image_loader, output_size in itertools.product(
617
618
        make_image_loaders(sizes=_CENTER_CROP_SPATIAL_SIZES, extra_dims=[()], dtypes=[torch.uint8]),
        _CENTER_CROP_OUTPUT_SIZES,
619
620
621
622
    ):
        yield ArgsKwargs(image_loader, output_size=output_size)


623
624
def sample_inputs_center_crop_bounding_boxes():
    for bounding_boxes_loader, output_size in itertools.product(make_bounding_box_loaders(), _CENTER_CROP_OUTPUT_SIZES):
625
        yield ArgsKwargs(
626
627
            bounding_boxes_loader,
            format=bounding_boxes_loader.format,
Philip Meier's avatar
Philip Meier committed
628
            canvas_size=bounding_boxes_loader.canvas_size,
629
630
631
632
633
            output_size=output_size,
        )


def sample_inputs_center_crop_mask():
634
    for mask_loader in make_mask_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], num_categories=[10], num_objects=[5]):
635
636
        height, width = mask_loader.shape[-2:]
        yield ArgsKwargs(mask_loader, output_size=(height // 2, width // 2))
637
638
639
640


def reference_inputs_center_crop_mask():
    for mask_loader, output_size in itertools.product(
641
        make_mask_loaders(sizes=_CENTER_CROP_SPATIAL_SIZES, extra_dims=[()], num_objects=[1]), _CENTER_CROP_OUTPUT_SIZES
642
643
644
645
    ):
        yield ArgsKwargs(mask_loader, output_size=output_size)


646
def sample_inputs_center_crop_video():
647
    for video_loader in make_video_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], num_frames=[3]):
648
649
650
651
        height, width = video_loader.shape[-2:]
        yield ArgsKwargs(video_loader, output_size=(height // 2, width // 2))


652
653
654
KERNEL_INFOS.extend(
    [
        KernelInfo(
655
            F.center_crop_image,
656
            sample_inputs_fn=sample_inputs_center_crop_image_tensor,
657
            reference_fn=pil_reference_wrapper(F._center_crop_image_pil),
658
            reference_inputs_fn=reference_inputs_center_crop_image_tensor,
659
            float32_vs_uint8=True,
660
            test_marks=[
661
                xfail_jit_python_scalar_arg("output_size"),
662
            ],
663
664
        ),
        KernelInfo(
665
666
            F.center_crop_bounding_boxes,
            sample_inputs_fn=sample_inputs_center_crop_bounding_boxes,
667
            test_marks=[
668
                xfail_jit_python_scalar_arg("output_size"),
669
            ],
670
671
672
673
        ),
        KernelInfo(
            F.center_crop_mask,
            sample_inputs_fn=sample_inputs_center_crop_mask,
674
            reference_fn=pil_reference_wrapper(F._center_crop_image_pil),
675
            reference_inputs_fn=reference_inputs_center_crop_mask,
676
            float32_vs_uint8=True,
677
            test_marks=[
678
                xfail_jit_python_scalar_arg("output_size"),
679
            ],
680
        ),
681
682
683
684
        KernelInfo(
            F.center_crop_video,
            sample_inputs_fn=sample_inputs_center_crop_video,
        ),
685
686
687
688
689
    ]
)


def sample_inputs_equalize_image_tensor():
690
    for image_loader in make_image_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], color_spaces=("GRAY", "RGB")):
691
692
693
694
        yield ArgsKwargs(image_loader)


def reference_inputs_equalize_image_tensor():
695
696
697
    # We are not using `make_image_loaders` here since that uniformly samples the values over the whole value range.
    # Since the whole point of this kernel is to transform an arbitrary distribution of values into a uniform one,
    # the information gain is low if we already provide something really close to the expected value.
698
    def make_uniform_band_image(shape, dtype, device, *, low_factor, high_factor, memory_format):
699
700
701
702
703
704
705
        if dtype.is_floating_point:
            low = low_factor
            high = high_factor
        else:
            max_value = torch.iinfo(dtype).max
            low = int(low_factor * max_value)
            high = int(high_factor * max_value)
706
707
708
        return torch.testing.make_tensor(shape, dtype=dtype, device=device, low=low, high=high).to(
            memory_format=memory_format, copy=True
        )
709

710
    def make_beta_distributed_image(shape, dtype, device, *, alpha, beta, memory_format):
711
712
713
        image = torch.distributions.Beta(alpha, beta).sample(shape)
        if not dtype.is_floating_point:
            image.mul_(torch.iinfo(dtype).max).round_()
714
        return image.to(dtype=dtype, device=device, memory_format=memory_format, copy=True)
715

Philip Meier's avatar
Philip Meier committed
716
    canvas_size = (256, 256)
717
    for dtype, color_space, fn in itertools.product(
718
        [torch.uint8],
719
        ["GRAY", "RGB"],
720
        [
721
722
            lambda shape, dtype, device, memory_format: torch.zeros(shape, dtype=dtype, device=device).to(
                memory_format=memory_format, copy=True
723
            ),
724
725
726
            lambda shape, dtype, device, memory_format: torch.full(
                shape, 1.0 if dtype.is_floating_point else torch.iinfo(dtype).max, dtype=dtype, device=device
            ).to(memory_format=memory_format, copy=True),
727
            *[
728
729
730
731
732
                functools.partial(make_uniform_band_image, low_factor=low_factor, high_factor=high_factor)
                for low_factor, high_factor in [
                    (0.0, 0.25),
                    (0.25, 0.75),
                    (0.75, 1.0),
733
734
735
                ]
            ],
            *[
736
                functools.partial(make_beta_distributed_image, alpha=alpha, beta=beta)
737
738
739
740
741
742
743
744
                for alpha, beta in [
                    (0.5, 0.5),
                    (2, 2),
                    (2, 5),
                    (5, 2),
                ]
            ],
        ],
745
    ):
Philip Meier's avatar
Philip Meier committed
746
        image_loader = ImageLoader(fn, shape=(get_num_channels(color_space), *canvas_size), dtype=dtype)
747
748
749
        yield ArgsKwargs(image_loader)


750
def sample_inputs_equalize_video():
751
    for video_loader in make_video_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], num_frames=[3]):
752
753
754
755
756
757
        yield ArgsKwargs(video_loader)


KERNEL_INFOS.extend(
    [
        KernelInfo(
758
            F.equalize_image,
759
760
            kernel_name="equalize_image_tensor",
            sample_inputs_fn=sample_inputs_equalize_image_tensor,
761
            reference_fn=pil_reference_wrapper(F._equalize_image_pil),
762
            float32_vs_uint8=True,
763
764
765
766
767
768
769
            reference_inputs_fn=reference_inputs_equalize_image_tensor,
        ),
        KernelInfo(
            F.equalize_video,
            sample_inputs_fn=sample_inputs_equalize_video,
        ),
    ]
770
771
772
773
)


def sample_inputs_invert_image_tensor():
774
    for image_loader in make_image_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], color_spaces=("GRAY", "RGB")):
775
776
777
778
        yield ArgsKwargs(image_loader)


def reference_inputs_invert_image_tensor():
779
    for image_loader in make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]):
780
781
782
        yield ArgsKwargs(image_loader)


783
def sample_inputs_invert_video():
784
    for video_loader in make_video_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], num_frames=[3]):
785
786
787
788
789
790
        yield ArgsKwargs(video_loader)


KERNEL_INFOS.extend(
    [
        KernelInfo(
791
            F.invert_image,
792
793
            kernel_name="invert_image_tensor",
            sample_inputs_fn=sample_inputs_invert_image_tensor,
794
            reference_fn=pil_reference_wrapper(F._invert_image_pil),
795
            reference_inputs_fn=reference_inputs_invert_image_tensor,
796
            float32_vs_uint8=True,
797
798
799
800
801
802
        ),
        KernelInfo(
            F.invert_video,
            sample_inputs_fn=sample_inputs_invert_video,
        ),
    ]
803
804
805
806
807
808
809
)


_POSTERIZE_BITS = [1, 4, 8]


def sample_inputs_posterize_image_tensor():
810
    for image_loader in make_image_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], color_spaces=("GRAY", "RGB")):
811
812
813
814
815
        yield ArgsKwargs(image_loader, bits=_POSTERIZE_BITS[0])


def reference_inputs_posterize_image_tensor():
    for image_loader, bits in itertools.product(
816
        make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]),
817
818
819
820
821
        _POSTERIZE_BITS,
    ):
        yield ArgsKwargs(image_loader, bits=bits)


822
def sample_inputs_posterize_video():
823
    for video_loader in make_video_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], num_frames=[3]):
824
825
826
827
828
829
        yield ArgsKwargs(video_loader, bits=_POSTERIZE_BITS[0])


KERNEL_INFOS.extend(
    [
        KernelInfo(
830
            F.posterize_image,
831
832
            kernel_name="posterize_image_tensor",
            sample_inputs_fn=sample_inputs_posterize_image_tensor,
833
            reference_fn=pil_reference_wrapper(F._posterize_image_pil),
834
            reference_inputs_fn=reference_inputs_posterize_image_tensor,
835
836
            float32_vs_uint8=True,
            closeness_kwargs=float32_vs_uint8_pixel_difference(),
837
838
839
840
841
842
        ),
        KernelInfo(
            F.posterize_video,
            sample_inputs_fn=sample_inputs_posterize_video,
        ),
    ]
843
844
845
846
847
848
849
850
851
852
)


def _get_solarize_thresholds(dtype):
    for factor in [0.1, 0.5]:
        max_value = get_max_value(dtype)
        yield (float if dtype.is_floating_point else int)(max_value * factor)


def sample_inputs_solarize_image_tensor():
853
    for image_loader in make_image_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], color_spaces=("GRAY", "RGB")):
854
855
856
857
        yield ArgsKwargs(image_loader, threshold=next(_get_solarize_thresholds(image_loader.dtype)))


def reference_inputs_solarize_image_tensor():
858
    for image_loader in make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]):
859
860
861
862
        for threshold in _get_solarize_thresholds(image_loader.dtype):
            yield ArgsKwargs(image_loader, threshold=threshold)


863
864
865
866
def uint8_to_float32_threshold_adapter(other_args, kwargs):
    return other_args, dict(threshold=kwargs["threshold"] / 255)


867
def sample_inputs_solarize_video():
868
    for video_loader in make_video_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], num_frames=[3]):
869
870
871
872
873
874
        yield ArgsKwargs(video_loader, threshold=next(_get_solarize_thresholds(video_loader.dtype)))


KERNEL_INFOS.extend(
    [
        KernelInfo(
875
            F.solarize_image,
876
877
            kernel_name="solarize_image_tensor",
            sample_inputs_fn=sample_inputs_solarize_image_tensor,
878
            reference_fn=pil_reference_wrapper(F._solarize_image_pil),
879
            reference_inputs_fn=reference_inputs_solarize_image_tensor,
880
881
            float32_vs_uint8=uint8_to_float32_threshold_adapter,
            closeness_kwargs=float32_vs_uint8_pixel_difference(),
882
883
884
885
886
887
        ),
        KernelInfo(
            F.solarize_video,
            sample_inputs_fn=sample_inputs_solarize_video,
        ),
    ]
888
889
890
891
)


def sample_inputs_autocontrast_image_tensor():
892
    for image_loader in make_image_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], color_spaces=("GRAY", "RGB")):
893
894
895
896
        yield ArgsKwargs(image_loader)


def reference_inputs_autocontrast_image_tensor():
897
    for image_loader in make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]):
898
899
900
        yield ArgsKwargs(image_loader)


901
def sample_inputs_autocontrast_video():
902
    for video_loader in make_video_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], num_frames=[3]):
903
904
905
906
907
908
        yield ArgsKwargs(video_loader)


KERNEL_INFOS.extend(
    [
        KernelInfo(
909
            F.autocontrast_image,
910
911
            kernel_name="autocontrast_image_tensor",
            sample_inputs_fn=sample_inputs_autocontrast_image_tensor,
912
            reference_fn=pil_reference_wrapper(F._autocontrast_image_pil),
913
            reference_inputs_fn=reference_inputs_autocontrast_image_tensor,
914
915
916
917
918
            float32_vs_uint8=True,
            closeness_kwargs={
                **pil_reference_pixel_difference(),
                **float32_vs_uint8_pixel_difference(),
            },
919
920
921
922
923
924
        ),
        KernelInfo(
            F.autocontrast_video,
            sample_inputs_fn=sample_inputs_autocontrast_video,
        ),
    ]
925
926
927
928
929
930
931
)

_ADJUST_SHARPNESS_FACTORS = [0.1, 0.5]


def sample_inputs_adjust_sharpness_image_tensor():
    for image_loader in make_image_loaders(
932
        sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE, (2, 2)],
933
        color_spaces=("GRAY", "RGB"),
934
935
936
937
938
939
    ):
        yield ArgsKwargs(image_loader, sharpness_factor=_ADJUST_SHARPNESS_FACTORS[0])


def reference_inputs_adjust_sharpness_image_tensor():
    for image_loader, sharpness_factor in itertools.product(
940
        make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]),
941
942
943
944
945
        _ADJUST_SHARPNESS_FACTORS,
    ):
        yield ArgsKwargs(image_loader, sharpness_factor=sharpness_factor)


946
def sample_inputs_adjust_sharpness_video():
947
    for video_loader in make_video_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], num_frames=[3]):
948
949
950
951
952
953
        yield ArgsKwargs(video_loader, sharpness_factor=_ADJUST_SHARPNESS_FACTORS[0])


KERNEL_INFOS.extend(
    [
        KernelInfo(
954
            F.adjust_sharpness_image,
955
956
            kernel_name="adjust_sharpness_image_tensor",
            sample_inputs_fn=sample_inputs_adjust_sharpness_image_tensor,
957
            reference_fn=pil_reference_wrapper(F._adjust_sharpness_image_pil),
958
            reference_inputs_fn=reference_inputs_adjust_sharpness_image_tensor,
959
960
            float32_vs_uint8=True,
            closeness_kwargs=float32_vs_uint8_pixel_difference(2),
961
962
963
964
965
966
        ),
        KernelInfo(
            F.adjust_sharpness_video,
            sample_inputs_fn=sample_inputs_adjust_sharpness_video,
        ),
    ]
967
968
969
)


970
971
972
973
_ADJUST_CONTRAST_FACTORS = [0.1, 0.5]


def sample_inputs_adjust_contrast_image_tensor():
974
    for image_loader in make_image_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], color_spaces=("GRAY", "RGB")):
975
976
977
978
979
        yield ArgsKwargs(image_loader, contrast_factor=_ADJUST_CONTRAST_FACTORS[0])


def reference_inputs_adjust_contrast_image_tensor():
    for image_loader, contrast_factor in itertools.product(
980
        make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]),
981
982
983
984
985
        _ADJUST_CONTRAST_FACTORS,
    ):
        yield ArgsKwargs(image_loader, contrast_factor=contrast_factor)


986
def sample_inputs_adjust_contrast_video():
987
    for video_loader in make_video_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], num_frames=[3]):
988
989
990
991
992
993
        yield ArgsKwargs(video_loader, contrast_factor=_ADJUST_CONTRAST_FACTORS[0])


KERNEL_INFOS.extend(
    [
        KernelInfo(
994
            F.adjust_contrast_image,
995
996
            kernel_name="adjust_contrast_image_tensor",
            sample_inputs_fn=sample_inputs_adjust_contrast_image_tensor,
997
            reference_fn=pil_reference_wrapper(F._adjust_contrast_image_pil),
998
            reference_inputs_fn=reference_inputs_adjust_contrast_image_tensor,
999
1000
1001
1002
            float32_vs_uint8=True,
            closeness_kwargs={
                **pil_reference_pixel_difference(),
                **float32_vs_uint8_pixel_difference(2),
1003
                **cuda_vs_cpu_pixel_difference(),
1004
                (("TestKernels", "test_against_reference"), torch.uint8, "cpu"): pixel_difference_closeness_kwargs(1),
1005
            },
1006
1007
1008
1009
        ),
        KernelInfo(
            F.adjust_contrast_video,
            sample_inputs_fn=sample_inputs_adjust_contrast_video,
1010
1011
1012
1013
            closeness_kwargs={
                **cuda_vs_cpu_pixel_difference(),
                (("TestKernels", "test_against_reference"), torch.uint8, "cpu"): pixel_difference_closeness_kwargs(1),
            },
1014
1015
        ),
    ]
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
)

_ADJUST_GAMMA_GAMMAS_GAINS = [
    (0.5, 2.0),
    (0.0, 1.0),
]


def sample_inputs_adjust_gamma_image_tensor():
    gamma, gain = _ADJUST_GAMMA_GAMMAS_GAINS[0]
1026
    for image_loader in make_image_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], color_spaces=("GRAY", "RGB")):
1027
1028
1029
1030
1031
        yield ArgsKwargs(image_loader, gamma=gamma, gain=gain)


def reference_inputs_adjust_gamma_image_tensor():
    for image_loader, (gamma, gain) in itertools.product(
1032
        make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]),
1033
1034
1035
1036
1037
        _ADJUST_GAMMA_GAMMAS_GAINS,
    ):
        yield ArgsKwargs(image_loader, gamma=gamma, gain=gain)


1038
1039
def sample_inputs_adjust_gamma_video():
    gamma, gain = _ADJUST_GAMMA_GAMMAS_GAINS[0]
1040
    for video_loader in make_video_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], num_frames=[3]):
1041
1042
1043
1044
1045
1046
        yield ArgsKwargs(video_loader, gamma=gamma, gain=gain)


KERNEL_INFOS.extend(
    [
        KernelInfo(
1047
            F.adjust_gamma_image,
1048
1049
            kernel_name="adjust_gamma_image_tensor",
            sample_inputs_fn=sample_inputs_adjust_gamma_image_tensor,
1050
            reference_fn=pil_reference_wrapper(F._adjust_gamma_image_pil),
1051
            reference_inputs_fn=reference_inputs_adjust_gamma_image_tensor,
1052
1053
1054
1055
1056
            float32_vs_uint8=True,
            closeness_kwargs={
                **pil_reference_pixel_difference(),
                **float32_vs_uint8_pixel_difference(),
            },
1057
1058
1059
1060
1061
1062
        ),
        KernelInfo(
            F.adjust_gamma_video,
            sample_inputs_fn=sample_inputs_adjust_gamma_video,
        ),
    ]
1063
1064
1065
1066
1067
1068
1069
)


_ADJUST_HUE_FACTORS = [-0.1, 0.5]


def sample_inputs_adjust_hue_image_tensor():
1070
    for image_loader in make_image_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], color_spaces=("GRAY", "RGB")):
1071
1072
1073
1074
1075
        yield ArgsKwargs(image_loader, hue_factor=_ADJUST_HUE_FACTORS[0])


def reference_inputs_adjust_hue_image_tensor():
    for image_loader, hue_factor in itertools.product(
1076
        make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]),
1077
1078
1079
1080
1081
        _ADJUST_HUE_FACTORS,
    ):
        yield ArgsKwargs(image_loader, hue_factor=hue_factor)


1082
def sample_inputs_adjust_hue_video():
1083
    for video_loader in make_video_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], num_frames=[3]):
1084
1085
1086
1087
1088
1089
        yield ArgsKwargs(video_loader, hue_factor=_ADJUST_HUE_FACTORS[0])


KERNEL_INFOS.extend(
    [
        KernelInfo(
1090
            F.adjust_hue_image,
1091
1092
            kernel_name="adjust_hue_image_tensor",
            sample_inputs_fn=sample_inputs_adjust_hue_image_tensor,
1093
            reference_fn=pil_reference_wrapper(F._adjust_hue_image_pil),
1094
            reference_inputs_fn=reference_inputs_adjust_hue_image_tensor,
1095
1096
            float32_vs_uint8=True,
            closeness_kwargs={
1097
                **pil_reference_pixel_difference(2, mae=True),
1098
1099
                **float32_vs_uint8_pixel_difference(),
            },
1100
1101
1102
1103
1104
1105
        ),
        KernelInfo(
            F.adjust_hue_video,
            sample_inputs_fn=sample_inputs_adjust_hue_video,
        ),
    ]
1106
1107
1108
1109
1110
1111
)

_ADJUST_SATURATION_FACTORS = [0.1, 0.5]


def sample_inputs_adjust_saturation_image_tensor():
1112
    for image_loader in make_image_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], color_spaces=("GRAY", "RGB")):
1113
1114
1115
1116
1117
        yield ArgsKwargs(image_loader, saturation_factor=_ADJUST_SATURATION_FACTORS[0])


def reference_inputs_adjust_saturation_image_tensor():
    for image_loader, saturation_factor in itertools.product(
1118
        make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]),
1119
1120
1121
1122
1123
        _ADJUST_SATURATION_FACTORS,
    ):
        yield ArgsKwargs(image_loader, saturation_factor=saturation_factor)


1124
def sample_inputs_adjust_saturation_video():
1125
    for video_loader in make_video_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], num_frames=[3]):
1126
1127
1128
1129
1130
1131
        yield ArgsKwargs(video_loader, saturation_factor=_ADJUST_SATURATION_FACTORS[0])


KERNEL_INFOS.extend(
    [
        KernelInfo(
1132
            F.adjust_saturation_image,
1133
1134
            kernel_name="adjust_saturation_image_tensor",
            sample_inputs_fn=sample_inputs_adjust_saturation_image_tensor,
1135
            reference_fn=pil_reference_wrapper(F._adjust_saturation_image_pil),
1136
            reference_inputs_fn=reference_inputs_adjust_saturation_image_tensor,
1137
1138
1139
1140
            float32_vs_uint8=True,
            closeness_kwargs={
                **pil_reference_pixel_difference(),
                **float32_vs_uint8_pixel_difference(2),
1141
                **cuda_vs_cpu_pixel_difference(),
1142
            },
1143
1144
1145
1146
        ),
        KernelInfo(
            F.adjust_saturation_video,
            sample_inputs_fn=sample_inputs_adjust_saturation_video,
1147
            closeness_kwargs=cuda_vs_cpu_pixel_difference(),
1148
1149
        ),
    ]
1150
1151
1152
)


1153
1154
def sample_inputs_clamp_bounding_boxes():
    for bounding_boxes_loader in make_bounding_box_loaders():
1155
        yield ArgsKwargs(
1156
1157
            bounding_boxes_loader,
            format=bounding_boxes_loader.format,
Philip Meier's avatar
Philip Meier committed
1158
            canvas_size=bounding_boxes_loader.canvas_size,
1159
1160
1161
1162
1163
        )


KERNEL_INFOS.append(
    KernelInfo(
1164
1165
        F.clamp_bounding_boxes,
        sample_inputs_fn=sample_inputs_clamp_bounding_boxes,
1166
        logs_usage=True,
1167
1168
1169
1170
1171
1172
    )
)

_FIVE_TEN_CROP_SIZES = [7, (6,), [5], (6, 5), [7, 6]]


Philip Meier's avatar
Philip Meier committed
1173
def _get_five_ten_crop_canvas_size(size):
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
    if isinstance(size, int):
        crop_height = crop_width = size
    elif len(size) == 1:
        crop_height = crop_width = size[0]
    else:
        crop_height, crop_width = size
    return 2 * crop_height, 2 * crop_width


def sample_inputs_five_crop_image_tensor():
    for size in _FIVE_TEN_CROP_SIZES:
1185
        for image_loader in make_image_loaders(
Philip Meier's avatar
Philip Meier committed
1186
            sizes=[_get_five_ten_crop_canvas_size(size)],
1187
            color_spaces=["RGB"],
1188
            dtypes=[torch.float32],
1189
        ):
1190
1191
1192
1193
1194
            yield ArgsKwargs(image_loader, size=size)


def reference_inputs_five_crop_image_tensor():
    for size in _FIVE_TEN_CROP_SIZES:
1195
        for image_loader in make_image_loaders(
Philip Meier's avatar
Philip Meier committed
1196
            sizes=[_get_five_ten_crop_canvas_size(size)], extra_dims=[()], dtypes=[torch.uint8]
1197
        ):
1198
1199
1200
            yield ArgsKwargs(image_loader, size=size)


1201
1202
def sample_inputs_five_crop_video():
    size = _FIVE_TEN_CROP_SIZES[0]
Philip Meier's avatar
Philip Meier committed
1203
    for video_loader in make_video_loaders(sizes=[_get_five_ten_crop_canvas_size(size)]):
1204
1205
1206
        yield ArgsKwargs(video_loader, size=size)


1207
1208
def sample_inputs_ten_crop_image_tensor():
    for size, vertical_flip in itertools.product(_FIVE_TEN_CROP_SIZES, [False, True]):
1209
        for image_loader in make_image_loaders(
Philip Meier's avatar
Philip Meier committed
1210
            sizes=[_get_five_ten_crop_canvas_size(size)],
1211
            color_spaces=["RGB"],
1212
            dtypes=[torch.float32],
1213
        ):
1214
1215
1216
1217
1218
            yield ArgsKwargs(image_loader, size=size, vertical_flip=vertical_flip)


def reference_inputs_ten_crop_image_tensor():
    for size, vertical_flip in itertools.product(_FIVE_TEN_CROP_SIZES, [False, True]):
1219
        for image_loader in make_image_loaders(
Philip Meier's avatar
Philip Meier committed
1220
            sizes=[_get_five_ten_crop_canvas_size(size)], extra_dims=[()], dtypes=[torch.uint8]
1221
        ):
1222
1223
1224
            yield ArgsKwargs(image_loader, size=size, vertical_flip=vertical_flip)


1225
1226
def sample_inputs_ten_crop_video():
    size = _FIVE_TEN_CROP_SIZES[0]
Philip Meier's avatar
Philip Meier committed
1227
    for video_loader in make_video_loaders(sizes=[_get_five_ten_crop_canvas_size(size)]):
1228
1229
1230
        yield ArgsKwargs(video_loader, size=size)


1231
1232
1233
1234
def multi_crop_pil_reference_wrapper(pil_kernel):
    def wrapper(input_tensor, *other_args, **kwargs):
        output = pil_reference_wrapper(pil_kernel)(input_tensor, *other_args, **kwargs)
        return type(output)(
1235
            F.to_dtype_image(F.to_image(output_pil), dtype=input_tensor.dtype, scale=True) for output_pil in output
1236
1237
1238
1239
1240
        )

    return wrapper


1241
1242
1243
1244
1245
_common_five_ten_crop_marks = [
    xfail_jit_python_scalar_arg("size"),
    mark_framework_limitation(("TestKernels", "test_batched_vs_single"), "Custom batching needed."),
]

1246
1247
1248
KERNEL_INFOS.extend(
    [
        KernelInfo(
1249
            F.five_crop_image,
1250
            sample_inputs_fn=sample_inputs_five_crop_image_tensor,
1251
            reference_fn=multi_crop_pil_reference_wrapper(F._five_crop_image_pil),
1252
            reference_inputs_fn=reference_inputs_five_crop_image_tensor,
1253
            test_marks=_common_five_ten_crop_marks,
1254
        ),
1255
1256
1257
1258
1259
        KernelInfo(
            F.five_crop_video,
            sample_inputs_fn=sample_inputs_five_crop_video,
            test_marks=_common_five_ten_crop_marks,
        ),
1260
        KernelInfo(
1261
            F.ten_crop_image,
1262
            sample_inputs_fn=sample_inputs_ten_crop_image_tensor,
1263
            reference_fn=multi_crop_pil_reference_wrapper(F._ten_crop_image_pil),
1264
            reference_inputs_fn=reference_inputs_ten_crop_image_tensor,
1265
            test_marks=_common_five_ten_crop_marks,
1266
        ),
1267
1268
1269
1270
1271
        KernelInfo(
            F.ten_crop_video,
            sample_inputs_fn=sample_inputs_ten_crop_video,
            test_marks=_common_five_ten_crop_marks,
        ),
1272
1273
1274
1275
1276
1277
    ]
)

_NORMALIZE_MEANS_STDS = [
    ((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ([0.0, 0.0, 0.0], [1.0, 1.0, 1.0]),
1278
    (0.5, 2.0),
1279
1280
1281
1282
1283
]


def sample_inputs_normalize_image_tensor():
    for image_loader, (mean, std) in itertools.product(
1284
        make_image_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], color_spaces=["RGB"], dtypes=[torch.float32]),
1285
1286
1287
1288
1289
        _NORMALIZE_MEANS_STDS,
    ):
        yield ArgsKwargs(image_loader, mean=mean, std=std)


1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
def reference_normalize_image_tensor(image, mean, std, inplace=False):
    mean = torch.tensor(mean).view(-1, 1, 1)
    std = torch.tensor(std).view(-1, 1, 1)

    sub = torch.Tensor.sub_ if inplace else torch.Tensor.sub
    return sub(image, mean).div_(std)


def reference_inputs_normalize_image_tensor():
    yield ArgsKwargs(
1300
        make_image_loader(size=(32, 32), color_space="RGB", extra_dims=[1]),
1301
1302
1303
1304
1305
        mean=[0.5, 0.5, 0.5],
        std=[1.0, 1.0, 1.0],
    )


1306
1307
1308
def sample_inputs_normalize_video():
    mean, std = _NORMALIZE_MEANS_STDS[0]
    for video_loader in make_video_loaders(
1309
        sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], color_spaces=["RGB"], num_frames=[3], dtypes=[torch.float32]
1310
1311
1312
1313
1314
1315
1316
    ):
        yield ArgsKwargs(video_loader, mean=mean, std=std)


KERNEL_INFOS.extend(
    [
        KernelInfo(
1317
            F.normalize_image,
1318
1319
            kernel_name="normalize_image_tensor",
            sample_inputs_fn=sample_inputs_normalize_image_tensor,
1320
1321
            reference_fn=reference_normalize_image_tensor,
            reference_inputs_fn=reference_inputs_normalize_image_tensor,
1322
1323
1324
1325
            test_marks=[
                xfail_jit_python_scalar_arg("mean"),
                xfail_jit_python_scalar_arg("std"),
            ],
1326
1327
1328
1329
1330
1331
        ),
        KernelInfo(
            F.normalize_video,
            sample_inputs_fn=sample_inputs_normalize_video,
        ),
    ]
1332
)
1333
1334


1335
def sample_inputs_uniform_temporal_subsample_video():
1336
    for video_loader in make_video_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], num_frames=[4]):
1337
        yield ArgsKwargs(video_loader, num_samples=2)
1338
1339


1340
def reference_uniform_temporal_subsample_video(x, num_samples):
1341
1342
    # Copy-pasted from
    # https://github.com/facebookresearch/pytorchvideo/blob/c8d23d8b7e597586a9e2d18f6ed31ad8aa379a7a/pytorchvideo/transforms/functional.py#L19
1343
    t = x.shape[-4]
1344
1345
1346
1347
    assert num_samples > 0 and t > 0
    # Sample by nearest neighbor interpolation if num_samples > t.
    indices = torch.linspace(0, t - 1, num_samples)
    indices = torch.clamp(indices, 0, t - 1).long()
1348
    return torch.index_select(x, -4, indices)
1349
1350
1351


def reference_inputs_uniform_temporal_subsample_video():
1352
1353
1354
    for video_loader in make_video_loaders(
        sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], color_spaces=["RGB"], num_frames=[10]
    ):
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
        for num_samples in range(1, video_loader.shape[-4] + 1):
            yield ArgsKwargs(video_loader, num_samples)


KERNEL_INFOS.append(
    KernelInfo(
        F.uniform_temporal_subsample_video,
        sample_inputs_fn=sample_inputs_uniform_temporal_subsample_video,
        reference_fn=reference_uniform_temporal_subsample_video,
        reference_inputs_fn=reference_inputs_uniform_temporal_subsample_video,
    )
)