transforms_v2_dispatcher_infos.py 13.2 KB
Newer Older
1
import collections.abc
2
3

import pytest
4
import torchvision.transforms.v2.functional as F
5
from common_utils import InfoBase, TestMark
6
from torchvision import datapoints
7
from transforms_v2_kernel_infos import KERNEL_INFOS, pad_xfail_jit_fill_condition
8
9
10
11

__all__ = ["DispatcherInfo", "DISPATCHER_INFOS"]


12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
class PILKernelInfo(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,
    ):
        super().__init__(id=kernel_name or kernel.__name__)
        self.kernel = kernel


class DispatcherInfo(InfoBase):
    _KERNEL_INFO_MAP = {info.kernel: info for info in KERNEL_INFOS}

    def __init__(
        self,
        dispatcher,
        *,
        # Dictionary of types that map to the kernel the dispatcher dispatches to.
        kernels,
        # If omitted, no PIL dispatch test will be performed.
        pil_kernel_info=None,
        # See InfoBase
        test_marks=None,
        # See InfoBase
        closeness_kwargs=None,
    ):
        super().__init__(id=dispatcher.__name__, test_marks=test_marks, closeness_kwargs=closeness_kwargs)
        self.dispatcher = dispatcher
        self.kernels = kernels
        self.pil_kernel_info = pil_kernel_info

        kernel_infos = {}
47
        for datapoint_type, kernel in self.kernels.items():
48
49
50
            kernel_info = self._KERNEL_INFO_MAP.get(kernel)
            if not kernel_info:
                raise pytest.UsageError(
51
                    f"Can't register {kernel.__name__} for type {datapoint_type} since there is no `KernelInfo` for it. "
52
                    f"Please add a `KernelInfo` for it in `transforms_v2_kernel_infos.py`."
53
                )
54
            kernel_infos[datapoint_type] = kernel_info
55
        self.kernel_infos = kernel_infos
56

57
58
59
    def sample_inputs(self, *datapoint_types, filter_metadata=True):
        for datapoint_type in datapoint_types or self.kernel_infos.keys():
            kernel_info = self.kernel_infos.get(datapoint_type)
60
61
62
63
            if not kernel_info:
                raise pytest.UsageError(f"There is no kernel registered for type {type.__name__}")

            sample_inputs = kernel_info.sample_inputs_fn()
64
65
66

            if not filter_metadata:
                yield from sample_inputs
67
                return
68

69
70
71
72
73
74
75
76
77
78
79
80
81
            import itertools

            for args_kwargs in sample_inputs:
                for name in itertools.chain(
                    datapoint_type.__annotations__.keys(),
                    # FIXME: this seems ok for conversion dispatchers, but we should probably handle this on a
                    #  per-dispatcher level. However, so far there is no option for that.
                    (f"old_{name}" for name in datapoint_type.__annotations__.keys()),
                ):
                    if name in args_kwargs.kwargs:
                        del args_kwargs.kwargs[name]

                yield args_kwargs
82

83

84
def xfail_jit(reason, *, condition=None):
85
86
87
    return TestMark(
        ("TestDispatchers", "test_scripted_smoke"),
        pytest.mark.xfail(reason=reason),
88
        condition=condition,
89
90
    )

91

92
93
94
95
96
def xfail_jit_python_scalar_arg(name, *, reason=None):
    return xfail_jit(
        reason or f"Python scalar int or float for `{name}` is not supported when scripting",
        condition=lambda args_kwargs: isinstance(args_kwargs.kwargs.get(name), (int, float)),
    )
97
98


99
100
101
skip_dispatch_datapoint = TestMark(
    ("TestDispatchers", "test_dispatch_datapoint"),
    pytest.mark.skip(reason="Dispatcher doesn't support arbitrary datapoint dispatch."),
102
103
)

104
105
106
107
108
109
110
111
112
multi_crop_skips = [
    TestMark(
        ("TestDispatchers", test_name),
        pytest.mark.skip(reason="Multi-crop dispatchers return a sequence of items rather than a single one."),
    )
    for test_name in ["test_simple_tensor_output_type", "test_pil_output_type", "test_datapoint_output_type"]
]
multi_crop_skips.append(skip_dispatch_datapoint)

113

114
115
116
117
118
119
120
def xfails_pil(reason, *, condition=None):
    return [
        TestMark(("TestDispatchers", test_name), pytest.mark.xfail(reason=reason), condition=condition)
        for test_name in ["test_dispatch_pil", "test_pil_output_type"]
    ]


121
def fill_sequence_needs_broadcast(args_kwargs):
122
123
124
125
126
127
128
129
130
131
132
133
    (image_loader, *_), kwargs = args_kwargs
    try:
        fill = kwargs["fill"]
    except KeyError:
        return False

    if not isinstance(fill, collections.abc.Sequence) or len(fill) > 1:
        return False

    return image_loader.num_channels > 1


134
135
xfails_pil_if_fill_sequence_needs_broadcast = xfails_pil(
    "PIL kernel doesn't support sequences of length 1 for `fill` if the number of color channels is larger.",
136
137
138
139
    condition=fill_sequence_needs_broadcast,
)


140
141
142
143
DISPATCHER_INFOS = [
    DispatcherInfo(
        F.crop,
        kernels={
144
145
            datapoints.Image: F.crop_image_tensor,
            datapoints.Video: F.crop_video,
146
            datapoints.BoundingBoxes: F.crop_bounding_boxes,
147
            datapoints.Mask: F.crop_mask,
148
        },
149
        pil_kernel_info=PILKernelInfo(F.crop_image_pil, kernel_name="crop_image_pil"),
150
151
152
153
    ),
    DispatcherInfo(
        F.resized_crop,
        kernels={
154
155
            datapoints.Image: F.resized_crop_image_tensor,
            datapoints.Video: F.resized_crop_video,
156
            datapoints.BoundingBoxes: F.resized_crop_bounding_boxes,
157
            datapoints.Mask: F.resized_crop_mask,
158
        },
159
        pil_kernel_info=PILKernelInfo(F.resized_crop_image_pil),
160
161
162
163
    ),
    DispatcherInfo(
        F.pad,
        kernels={
164
165
            datapoints.Image: F.pad_image_tensor,
            datapoints.Video: F.pad_video,
166
            datapoints.BoundingBoxes: F.pad_bounding_boxes,
167
            datapoints.Mask: F.pad_mask,
168
        },
169
        pil_kernel_info=PILKernelInfo(F.pad_image_pil, kernel_name="pad_image_pil"),
170
        test_marks=[
171
172
173
174
            *xfails_pil(
                reason=(
                    "PIL kernel doesn't support sequences of length 1 for argument `fill` and "
                    "`padding_mode='constant'`, if the number of color channels is larger."
175
176
177
                ),
                condition=lambda args_kwargs: fill_sequence_needs_broadcast(args_kwargs)
                and args_kwargs.kwargs.get("padding_mode", "constant") == "constant",
178
            ),
179
180
            xfail_jit("F.pad only supports vector fills for list of floats", condition=pad_xfail_jit_fill_condition),
            xfail_jit_python_scalar_arg("padding"),
181
        ],
182
    ),
183
184
185
    DispatcherInfo(
        F.perspective,
        kernels={
186
187
            datapoints.Image: F.perspective_image_tensor,
            datapoints.Video: F.perspective_video,
188
            datapoints.BoundingBoxes: F.perspective_bounding_boxes,
189
            datapoints.Mask: F.perspective_mask,
190
        },
191
        pil_kernel_info=PILKernelInfo(F.perspective_image_pil),
192
        test_marks=[
193
194
            *xfails_pil_if_fill_sequence_needs_broadcast,
            xfail_jit_python_scalar_arg("fill"),
195
        ],
196
    ),
197
198
199
    DispatcherInfo(
        F.elastic,
        kernels={
200
201
            datapoints.Image: F.elastic_image_tensor,
            datapoints.Video: F.elastic_video,
202
            datapoints.BoundingBoxes: F.elastic_bounding_boxes,
203
            datapoints.Mask: F.elastic_mask,
204
        },
205
        pil_kernel_info=PILKernelInfo(F.elastic_image_pil),
206
        test_marks=[xfail_jit_python_scalar_arg("fill")],
207
    ),
208
209
210
    DispatcherInfo(
        F.center_crop,
        kernels={
211
212
            datapoints.Image: F.center_crop_image_tensor,
            datapoints.Video: F.center_crop_video,
213
            datapoints.BoundingBoxes: F.center_crop_bounding_boxes,
214
            datapoints.Mask: F.center_crop_mask,
215
        },
216
        pil_kernel_info=PILKernelInfo(F.center_crop_image_pil),
217
        test_marks=[
218
            xfail_jit_python_scalar_arg("output_size"),
219
        ],
220
221
222
223
    ),
    DispatcherInfo(
        F.gaussian_blur,
        kernels={
224
225
            datapoints.Image: F.gaussian_blur_image_tensor,
            datapoints.Video: F.gaussian_blur_video,
226
        },
227
        pil_kernel_info=PILKernelInfo(F.gaussian_blur_image_pil),
228
        test_marks=[
229
230
            xfail_jit_python_scalar_arg("kernel_size"),
            xfail_jit_python_scalar_arg("sigma"),
231
        ],
232
233
234
235
    ),
    DispatcherInfo(
        F.equalize,
        kernels={
236
237
            datapoints.Image: F.equalize_image_tensor,
            datapoints.Video: F.equalize_video,
238
        },
239
        pil_kernel_info=PILKernelInfo(F.equalize_image_pil, kernel_name="equalize_image_pil"),
240
241
242
243
    ),
    DispatcherInfo(
        F.invert,
        kernels={
244
245
            datapoints.Image: F.invert_image_tensor,
            datapoints.Video: F.invert_video,
246
        },
247
        pil_kernel_info=PILKernelInfo(F.invert_image_pil, kernel_name="invert_image_pil"),
248
249
250
251
    ),
    DispatcherInfo(
        F.posterize,
        kernels={
252
253
            datapoints.Image: F.posterize_image_tensor,
            datapoints.Video: F.posterize_video,
254
        },
255
        pil_kernel_info=PILKernelInfo(F.posterize_image_pil, kernel_name="posterize_image_pil"),
256
257
258
259
    ),
    DispatcherInfo(
        F.solarize,
        kernels={
260
261
            datapoints.Image: F.solarize_image_tensor,
            datapoints.Video: F.solarize_video,
262
        },
263
        pil_kernel_info=PILKernelInfo(F.solarize_image_pil, kernel_name="solarize_image_pil"),
264
265
266
267
    ),
    DispatcherInfo(
        F.autocontrast,
        kernels={
268
269
            datapoints.Image: F.autocontrast_image_tensor,
            datapoints.Video: F.autocontrast_video,
270
        },
271
        pil_kernel_info=PILKernelInfo(F.autocontrast_image_pil, kernel_name="autocontrast_image_pil"),
272
273
274
275
    ),
    DispatcherInfo(
        F.adjust_sharpness,
        kernels={
276
277
            datapoints.Image: F.adjust_sharpness_image_tensor,
            datapoints.Video: F.adjust_sharpness_video,
278
        },
279
        pil_kernel_info=PILKernelInfo(F.adjust_sharpness_image_pil, kernel_name="adjust_sharpness_image_pil"),
280
281
282
283
    ),
    DispatcherInfo(
        F.erase,
        kernels={
284
285
            datapoints.Image: F.erase_image_tensor,
            datapoints.Video: F.erase_video,
286
        },
287
        pil_kernel_info=PILKernelInfo(F.erase_image_pil),
288
        test_marks=[
289
            skip_dispatch_datapoint,
290
        ],
291
    ),
292
293
294
    DispatcherInfo(
        F.adjust_brightness,
        kernels={
295
296
            datapoints.Image: F.adjust_brightness_image_tensor,
            datapoints.Video: F.adjust_brightness_video,
297
        },
298
        pil_kernel_info=PILKernelInfo(F.adjust_brightness_image_pil, kernel_name="adjust_brightness_image_pil"),
299
300
301
302
    ),
    DispatcherInfo(
        F.adjust_contrast,
        kernels={
303
304
            datapoints.Image: F.adjust_contrast_image_tensor,
            datapoints.Video: F.adjust_contrast_video,
305
        },
306
        pil_kernel_info=PILKernelInfo(F.adjust_contrast_image_pil, kernel_name="adjust_contrast_image_pil"),
307
308
309
310
    ),
    DispatcherInfo(
        F.adjust_gamma,
        kernels={
311
312
            datapoints.Image: F.adjust_gamma_image_tensor,
            datapoints.Video: F.adjust_gamma_video,
313
        },
314
        pil_kernel_info=PILKernelInfo(F.adjust_gamma_image_pil, kernel_name="adjust_gamma_image_pil"),
315
316
317
318
    ),
    DispatcherInfo(
        F.adjust_hue,
        kernels={
319
320
            datapoints.Image: F.adjust_hue_image_tensor,
            datapoints.Video: F.adjust_hue_video,
321
        },
322
        pil_kernel_info=PILKernelInfo(F.adjust_hue_image_pil, kernel_name="adjust_hue_image_pil"),
323
324
325
326
    ),
    DispatcherInfo(
        F.adjust_saturation,
        kernels={
327
328
            datapoints.Image: F.adjust_saturation_image_tensor,
            datapoints.Video: F.adjust_saturation_video,
329
        },
330
        pil_kernel_info=PILKernelInfo(F.adjust_saturation_image_pil, kernel_name="adjust_saturation_image_pil"),
331
332
333
334
    ),
    DispatcherInfo(
        F.five_crop,
        kernels={
335
336
            datapoints.Image: F.five_crop_image_tensor,
            datapoints.Video: F.five_crop_video,
337
        },
338
        pil_kernel_info=PILKernelInfo(F.five_crop_image_pil),
339
        test_marks=[
340
            xfail_jit_python_scalar_arg("size"),
341
            *multi_crop_skips,
342
343
344
345
346
        ],
    ),
    DispatcherInfo(
        F.ten_crop,
        kernels={
347
348
            datapoints.Image: F.ten_crop_image_tensor,
            datapoints.Video: F.ten_crop_video,
349
        },
350
        test_marks=[
351
            xfail_jit_python_scalar_arg("size"),
352
            *multi_crop_skips,
353
        ],
354
        pil_kernel_info=PILKernelInfo(F.ten_crop_image_pil),
355
356
357
358
    ),
    DispatcherInfo(
        F.normalize,
        kernels={
359
360
            datapoints.Image: F.normalize_image_tensor,
            datapoints.Video: F.normalize_video,
361
        },
362
        test_marks=[
363
364
            xfail_jit_python_scalar_arg("mean"),
            xfail_jit_python_scalar_arg("std"),
365
        ],
366
    ),
367
368
369
    DispatcherInfo(
        F.uniform_temporal_subsample,
        kernels={
370
            datapoints.Video: F.uniform_temporal_subsample_video,
371
372
        },
        test_marks=[
373
            skip_dispatch_datapoint,
374
375
        ],
    ),
376
    DispatcherInfo(
377
378
        F.clamp_bounding_boxes,
        kernels={datapoints.BoundingBoxes: F.clamp_bounding_boxes},
379
380
381
382
383
        test_marks=[
            skip_dispatch_datapoint,
        ],
    ),
    DispatcherInfo(
384
385
        F.convert_format_bounding_boxes,
        kernels={datapoints.BoundingBoxes: F.convert_format_bounding_boxes},
386
387
388
389
        test_marks=[
            skip_dispatch_datapoint,
        ],
    ),
390
]