_geometry.py 88.3 KB
Newer Older
1
import math
2
import numbers
3
import warnings
4
from typing import Any, List, Optional, Sequence, Tuple, Union
5
6
7

import PIL.Image
import torch
8
from torch.nn.functional import grid_sample, interpolate, pad as torch_pad
9

10
from torchvision import datapoints
11
12
from torchvision.transforms import _functional_pil as _FP
from torchvision.transforms._functional_tensor import _pad_symmetric
13
from torchvision.transforms.functional import (
14
    _check_antialias,
15
    _compute_resized_output_size as __compute_resized_output_size,
16
    _get_perspective_coeffs,
17
    _interpolation_modes_from_int,
18
    InterpolationMode,
19
    pil_modes_mapping,
20
21
    pil_to_tensor,
    to_pil_image,
22
)
23

24
25
from torchvision.utils import _log_api_usage_once

Philip Meier's avatar
Philip Meier committed
26
from ._meta import clamp_bounding_boxes, convert_format_bounding_boxes, get_size_image_pil
27

28
29
30
31
32
33
34
from ._utils import (
    _get_kernel,
    _register_explicit_noop,
    _register_five_ten_crop_kernel,
    _register_kernel_internal,
    is_simple_tensor,
)
35

36

37
38
39
40
41
42
43
44
45
46
47
def _check_interpolation(interpolation: Union[InterpolationMode, int]) -> InterpolationMode:
    if isinstance(interpolation, int):
        interpolation = _interpolation_modes_from_int(interpolation)
    elif not isinstance(interpolation, InterpolationMode):
        raise ValueError(
            f"Argument interpolation should be an `InterpolationMode` or a corresponding Pillow integer constant, "
            f"but got {interpolation}."
        )
    return interpolation


48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
def horizontal_flip(inpt: datapoints._InputTypeJIT) -> datapoints._InputTypeJIT:
    if not torch.jit.is_scripting():
        _log_api_usage_once(horizontal_flip)

    if torch.jit.is_scripting() or is_simple_tensor(inpt):
        return horizontal_flip_image_tensor(inpt)
    elif isinstance(inpt, datapoints.Datapoint):
        kernel = _get_kernel(horizontal_flip, type(inpt))
        return kernel(
            inpt,
        )
    elif isinstance(inpt, PIL.Image.Image):
        return horizontal_flip_image_pil(inpt)
    else:
        raise TypeError(
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
            f"but got {type(inpt)} instead."
        )


@_register_kernel_internal(horizontal_flip, datapoints.Image)
69
70
71
72
def horizontal_flip_image_tensor(image: torch.Tensor) -> torch.Tensor:
    return image.flip(-1)


73
74
def horizontal_flip_image_pil(image: PIL.Image.Image) -> PIL.Image.Image:
    return _FP.hflip(image)
75
76


77
@_register_kernel_internal(horizontal_flip, datapoints.Mask)
78
79
def horizontal_flip_mask(mask: torch.Tensor) -> torch.Tensor:
    return horizontal_flip_image_tensor(mask)
80
81


82
def horizontal_flip_bounding_boxes(
Philip Meier's avatar
Philip Meier committed
83
    bounding_boxes: torch.Tensor, format: datapoints.BoundingBoxFormat, canvas_size: Tuple[int, int]
84
) -> torch.Tensor:
85
    shape = bounding_boxes.shape
86

87
    bounding_boxes = bounding_boxes.clone().reshape(-1, 4)
88

89
    if format == datapoints.BoundingBoxFormat.XYXY:
Philip Meier's avatar
Philip Meier committed
90
        bounding_boxes[:, [2, 0]] = bounding_boxes[:, [0, 2]].sub_(canvas_size[1]).neg_()
91
    elif format == datapoints.BoundingBoxFormat.XYWH:
Philip Meier's avatar
Philip Meier committed
92
        bounding_boxes[:, 0].add_(bounding_boxes[:, 2]).sub_(canvas_size[1]).neg_()
93
    else:  # format == datapoints.BoundingBoxFormat.CXCYWH:
Philip Meier's avatar
Philip Meier committed
94
        bounding_boxes[:, 0].sub_(canvas_size[1]).neg_()
95

96
    return bounding_boxes.reshape(shape)
97
98


99
100
101
102
103
104
105
106
107
@_register_kernel_internal(horizontal_flip, datapoints.BoundingBoxes, datapoint_wrapper=False)
def _horizontal_flip_bounding_boxes_dispatch(inpt: datapoints.BoundingBoxes) -> datapoints.BoundingBoxes:
    output = horizontal_flip_bounding_boxes(
        inpt.as_subclass(torch.Tensor), format=inpt.format, canvas_size=inpt.canvas_size
    )
    return datapoints.BoundingBoxes.wrap_like(inpt, output)


@_register_kernel_internal(horizontal_flip, datapoints.Video)
108
109
110
111
def horizontal_flip_video(video: torch.Tensor) -> torch.Tensor:
    return horizontal_flip_image_tensor(video)


112
def vertical_flip(inpt: datapoints._InputTypeJIT) -> datapoints._InputTypeJIT:
113
    if not torch.jit.is_scripting():
114
        _log_api_usage_once(vertical_flip)
115

116
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
117
118
119
120
121
122
        return vertical_flip_image_tensor(inpt)
    elif isinstance(inpt, datapoints.Datapoint):
        kernel = _get_kernel(vertical_flip, type(inpt))
        return kernel(
            inpt,
        )
123
    elif isinstance(inpt, PIL.Image.Image):
124
        return vertical_flip_image_pil(inpt)
125
126
    else:
        raise TypeError(
127
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
128
129
            f"but got {type(inpt)} instead."
        )
130
131


132
@_register_kernel_internal(vertical_flip, datapoints.Image)
133
134
135
136
def vertical_flip_image_tensor(image: torch.Tensor) -> torch.Tensor:
    return image.flip(-2)


Philip Meier's avatar
Philip Meier committed
137
138
def vertical_flip_image_pil(image: PIL.Image) -> PIL.Image:
    return _FP.vflip(image)
139
140


141
@_register_kernel_internal(vertical_flip, datapoints.Mask)
142
143
def vertical_flip_mask(mask: torch.Tensor) -> torch.Tensor:
    return vertical_flip_image_tensor(mask)
144
145


146
def vertical_flip_bounding_boxes(
Philip Meier's avatar
Philip Meier committed
147
    bounding_boxes: torch.Tensor, format: datapoints.BoundingBoxFormat, canvas_size: Tuple[int, int]
148
) -> torch.Tensor:
149
    shape = bounding_boxes.shape
150

151
    bounding_boxes = bounding_boxes.clone().reshape(-1, 4)
152

153
    if format == datapoints.BoundingBoxFormat.XYXY:
Philip Meier's avatar
Philip Meier committed
154
        bounding_boxes[:, [1, 3]] = bounding_boxes[:, [3, 1]].sub_(canvas_size[0]).neg_()
155
    elif format == datapoints.BoundingBoxFormat.XYWH:
Philip Meier's avatar
Philip Meier committed
156
        bounding_boxes[:, 1].add_(bounding_boxes[:, 3]).sub_(canvas_size[0]).neg_()
157
    else:  # format == datapoints.BoundingBoxFormat.CXCYWH:
Philip Meier's avatar
Philip Meier committed
158
        bounding_boxes[:, 1].sub_(canvas_size[0]).neg_()
159

160
    return bounding_boxes.reshape(shape)
161
162


163
164
165
166
167
168
@_register_kernel_internal(vertical_flip, datapoints.BoundingBoxes, datapoint_wrapper=False)
def _vertical_flip_bounding_boxes_dispatch(inpt: datapoints.BoundingBoxes) -> datapoints.BoundingBoxes:
    output = vertical_flip_bounding_boxes(
        inpt.as_subclass(torch.Tensor), format=inpt.format, canvas_size=inpt.canvas_size
    )
    return datapoints.BoundingBoxes.wrap_like(inpt, output)
169

170

171
172
173
@_register_kernel_internal(vertical_flip, datapoints.Video)
def vertical_flip_video(video: torch.Tensor) -> torch.Tensor:
    return vertical_flip_image_tensor(video)
174
175


176
177
178
179
180
181
# We changed the names to align them with the transforms, i.e. `RandomHorizontalFlip`. Still, `hflip` and `vflip` are
# prevalent and well understood. Thus, we just alias them without deprecating the old names.
hflip = horizontal_flip
vflip = vertical_flip


182
def _compute_resized_output_size(
Philip Meier's avatar
Philip Meier committed
183
    canvas_size: Tuple[int, int], size: List[int], max_size: Optional[int] = None
184
185
186
) -> List[int]:
    if isinstance(size, int):
        size = [size]
187
188
189
190
191
    elif max_size is not None and len(size) != 1:
        raise ValueError(
            "max_size should only be passed if size specifies the length of the smaller edge, "
            "i.e. size should be an int or a sequence of length 1 in torchscript mode."
        )
Philip Meier's avatar
Philip Meier committed
192
    return __compute_resized_output_size(canvas_size, size=size, max_size=max_size)
193
194


195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
def resize(
    inpt: datapoints._InputTypeJIT,
    size: List[int],
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
    max_size: Optional[int] = None,
    antialias: Optional[Union[str, bool]] = "warn",
) -> datapoints._InputTypeJIT:
    if not torch.jit.is_scripting():
        _log_api_usage_once(resize)
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
        return resize_image_tensor(inpt, size, interpolation=interpolation, max_size=max_size, antialias=antialias)
    elif isinstance(inpt, datapoints.Datapoint):
        kernel = _get_kernel(resize, type(inpt))
        return kernel(inpt, size, interpolation=interpolation, max_size=max_size, antialias=antialias)
    elif isinstance(inpt, PIL.Image.Image):
        if antialias is False:
            warnings.warn("Anti-alias option is always applied for PIL Image input. Argument antialias is ignored.")
        return resize_image_pil(inpt, size, interpolation=interpolation, max_size=max_size)
    else:
        raise TypeError(
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
            f"but got {type(inpt)} instead."
        )


@_register_kernel_internal(resize, datapoints.Image)
221
222
223
def resize_image_tensor(
    image: torch.Tensor,
    size: List[int],
224
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
225
    max_size: Optional[int] = None,
226
    antialias: Optional[Union[str, bool]] = "warn",
227
) -> torch.Tensor:
228
    interpolation = _check_interpolation(interpolation)
229
230
    antialias = _check_antialias(img=image, antialias=antialias, interpolation=interpolation)
    assert not isinstance(antialias, str)
231
    antialias = False if antialias is None else antialias
232
233
234
    align_corners: Optional[bool] = None
    if interpolation == InterpolationMode.BILINEAR or interpolation == InterpolationMode.BICUBIC:
        align_corners = False
235
236
237
238
    else:
        # The default of antialias should be True from 0.17, so we don't warn or
        # error if other interpolation modes are used. This is documented.
        antialias = False
239

240
    shape = image.shape
241
    numel = image.numel()
242
    num_channels, old_height, old_width = shape[-3:]
vfdev's avatar
vfdev committed
243
    new_height, new_width = _compute_resized_output_size((old_height, old_width), size=size, max_size=max_size)
244

245
246
    if (new_height, new_width) == (old_height, old_width):
        return image
247
    elif numel > 0:
248
        image = image.reshape(-1, num_channels, old_height, old_width)
249

250
        dtype = image.dtype
251
252
253
254
        acceptable_dtypes = [torch.float32, torch.float64]
        if interpolation == InterpolationMode.NEAREST or interpolation == InterpolationMode.NEAREST_EXACT:
            # uint8 dtype can be included for cpu and cuda input if nearest mode
            acceptable_dtypes.append(torch.uint8)
255
256
257
258
259
260
261
        elif image.device.type == "cpu":
            # uint8 dtype support for bilinear and bicubic is limited to cpu and
            # according to our benchmarks, non-AVX CPUs should still prefer u8->f32->interpolate->u8 path for bilinear
            if (interpolation == InterpolationMode.BILINEAR and "AVX2" in torch.backends.cpu.get_cpu_capability()) or (
                interpolation == InterpolationMode.BICUBIC
            ):
                acceptable_dtypes.append(torch.uint8)
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277

        strides = image.stride()
        if image.is_contiguous(memory_format=torch.channels_last) and image.shape[0] == 1 and numel != strides[0]:
            # There is a weird behaviour in torch core where the output tensor of `interpolate()` can be allocated as
            # contiguous even though the input is un-ambiguously channels_last (https://github.com/pytorch/pytorch/issues/68430).
            # In particular this happens for the typical torchvision use-case of single CHW images where we fake the batch dim
            # to become 1CHW. Below, we restride those tensors to trick torch core into properly allocating the output as
            # channels_last, thus preserving the memory format of the input. This is not just for format consistency:
            # for uint8 bilinear images, this also avoids an extra copy (re-packing) of the output and saves time.
            # TODO: when https://github.com/pytorch/pytorch/issues/68430 is fixed (possibly by https://github.com/pytorch/pytorch/pull/100373),
            # we should be able to remove this hack.
            new_strides = list(strides)
            new_strides[0] = numel
            image = image.as_strided((1, num_channels, old_height, old_width), new_strides)

        need_cast = dtype not in acceptable_dtypes
278
279
280
281
        if need_cast:
            image = image.to(dtype=torch.float32)

        image = interpolate(
282
283
            image,
            size=[new_height, new_width],
284
285
            mode=interpolation.value,
            align_corners=align_corners,
286
287
            antialias=antialias,
        )
288

289
290
        if need_cast:
            if interpolation == InterpolationMode.BICUBIC and dtype == torch.uint8:
291
                # This path is hit on non-AVX archs, or on GPU.
292
                image = image.clamp_(min=0, max=255)
293
294
295
            if dtype in (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64):
                image = image.round_()
            image = image.to(dtype=dtype)
296

297
    return image.reshape(shape[:-3] + (num_channels, new_height, new_width))
298
299


300
@torch.jit.unused
301
def resize_image_pil(
302
    image: PIL.Image.Image,
303
    size: Union[Sequence[int], int],
304
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
305
306
    max_size: Optional[int] = None,
) -> PIL.Image.Image:
307
308
309
310
311
312
313
    old_height, old_width = image.height, image.width
    new_height, new_width = _compute_resized_output_size(
        (old_height, old_width),
        size=size,  # type: ignore[arg-type]
        max_size=max_size,
    )

314
    interpolation = _check_interpolation(interpolation)
315
316
317
318
319

    if (new_height, new_width) == (old_height, old_width):
        return image

    return image.resize((new_width, new_height), resample=pil_modes_mapping[interpolation])
320
321


322
323
324
def resize_mask(mask: torch.Tensor, size: List[int], max_size: Optional[int] = None) -> torch.Tensor:
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
325
326
327
328
        needs_squeeze = True
    else:
        needs_squeeze = False

329
    output = resize_image_tensor(mask, size=size, interpolation=InterpolationMode.NEAREST, max_size=max_size)
330
331
332
333
334

    if needs_squeeze:
        output = output.squeeze(0)

    return output
335
336


337
338
339
340
341
342
343
344
@_register_kernel_internal(resize, datapoints.Mask, datapoint_wrapper=False)
def _resize_mask_dispatch(
    inpt: datapoints.Mask, size: List[int], max_size: Optional[int] = None, **kwargs: Any
) -> datapoints.Mask:
    output = resize_mask(inpt.as_subclass(torch.Tensor), size, max_size=max_size)
    return datapoints.Mask.wrap_like(inpt, output)


345
def resize_bounding_boxes(
Philip Meier's avatar
Philip Meier committed
346
    bounding_boxes: torch.Tensor, canvas_size: Tuple[int, int], size: List[int], max_size: Optional[int] = None
347
) -> Tuple[torch.Tensor, Tuple[int, int]]:
Philip Meier's avatar
Philip Meier committed
348
349
    old_height, old_width = canvas_size
    new_height, new_width = _compute_resized_output_size(canvas_size, size=size, max_size=max_size)
350
351

    if (new_height, new_width) == (old_height, old_width):
Philip Meier's avatar
Philip Meier committed
352
        return bounding_boxes, canvas_size
353

354
355
    w_ratio = new_width / old_width
    h_ratio = new_height / old_height
356
    ratios = torch.tensor([w_ratio, h_ratio, w_ratio, h_ratio], device=bounding_boxes.device)
357
    return (
358
        bounding_boxes.mul(ratios).to(bounding_boxes.dtype),
359
360
        (new_height, new_width),
    )
361
362


363
364
365
366
367
368
369
370
371
372
373
@_register_kernel_internal(resize, datapoints.BoundingBoxes, datapoint_wrapper=False)
def _resize_bounding_boxes_dispatch(
    inpt: datapoints.BoundingBoxes, size: List[int], max_size: Optional[int] = None, **kwargs: Any
) -> datapoints.BoundingBoxes:
    output, canvas_size = resize_bounding_boxes(
        inpt.as_subclass(torch.Tensor), inpt.canvas_size, size, max_size=max_size
    )
    return datapoints.BoundingBoxes.wrap_like(inpt, output, canvas_size=canvas_size)


@_register_kernel_internal(resize, datapoints.Video)
374
375
376
def resize_video(
    video: torch.Tensor,
    size: List[int],
377
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
378
    max_size: Optional[int] = None,
379
    antialias: Optional[Union[str, bool]] = "warn",
380
381
382
383
) -> torch.Tensor:
    return resize_image_tensor(video, size=size, interpolation=interpolation, max_size=max_size, antialias=antialias)


384
def affine(
Philip Meier's avatar
Philip Meier committed
385
    inpt: datapoints._InputTypeJIT,
386
387
388
389
390
391
392
    angle: Union[int, float],
    translate: List[float],
    scale: float,
    shear: List[float],
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
    fill: datapoints._FillTypeJIT = None,
    center: Optional[List[float]] = None,
Philip Meier's avatar
Philip Meier committed
393
) -> datapoints._InputTypeJIT:
394
    if not torch.jit.is_scripting():
395
396
397
        _log_api_usage_once(affine)

    # TODO: consider deprecating integers from angle and shear on the future
398
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
        return affine_image_tensor(
            inpt,
            angle,
            translate=translate,
            scale=scale,
            shear=shear,
            interpolation=interpolation,
            fill=fill,
            center=center,
        )
    elif isinstance(inpt, datapoints.Datapoint):
        kernel = _get_kernel(affine, type(inpt))
        return kernel(
            inpt,
            angle,
            translate=translate,
            scale=scale,
            shear=shear,
            interpolation=interpolation,
            fill=fill,
            center=center,
        )
421
    elif isinstance(inpt, PIL.Image.Image):
422
423
424
425
426
427
428
429
430
431
        return affine_image_pil(
            inpt,
            angle,
            translate=translate,
            scale=scale,
            shear=shear,
            interpolation=interpolation,
            fill=fill,
            center=center,
        )
432
433
    else:
        raise TypeError(
434
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
435
436
            f"but got {type(inpt)} instead."
        )
437
438


439
def _affine_parse_args(
440
    angle: Union[int, float],
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
    translate: List[float],
    scale: float,
    shear: List[float],
    interpolation: InterpolationMode = InterpolationMode.NEAREST,
    center: Optional[List[float]] = None,
) -> Tuple[float, List[float], List[float], Optional[List[float]]]:
    if not isinstance(angle, (int, float)):
        raise TypeError("Argument angle should be int or float")

    if not isinstance(translate, (list, tuple)):
        raise TypeError("Argument translate should be a sequence")

    if len(translate) != 2:
        raise ValueError("Argument translate should be a sequence of length 2")

    if scale <= 0.0:
        raise ValueError("Argument scale should be positive")

    if not isinstance(shear, (numbers.Number, (list, tuple))):
        raise TypeError("Shear should be either a single value or a sequence of two values")

    if not isinstance(interpolation, InterpolationMode):
        raise TypeError("Argument interpolation should be a InterpolationMode")

    if isinstance(angle, int):
        angle = float(angle)

    if isinstance(translate, tuple):
        translate = list(translate)

    if isinstance(shear, numbers.Number):
        shear = [shear, 0.0]

    if isinstance(shear, tuple):
        shear = list(shear)

    if len(shear) == 1:
        shear = [shear[0], shear[0]]

    if len(shear) != 2:
        raise ValueError(f"Shear should be a sequence containing two values. Got {shear}")

483
484
485
486
487
    if center is not None:
        if not isinstance(center, (list, tuple)):
            raise TypeError("Argument center should be a sequence")
        else:
            center = [float(c) for c in center]
488
489
490
491

    return angle, translate, shear, center


492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
def _get_inverse_affine_matrix(
    center: List[float], angle: float, translate: List[float], scale: float, shear: List[float], inverted: bool = True
) -> List[float]:
    # Helper method to compute inverse matrix for affine transformation

    # Pillow requires inverse affine transformation matrix:
    # Affine matrix is : M = T * C * RotateScaleShear * C^-1
    #
    # where T is translation matrix: [1, 0, tx | 0, 1, ty | 0, 0, 1]
    #       C is translation matrix to keep center: [1, 0, cx | 0, 1, cy | 0, 0, 1]
    #       RotateScaleShear is rotation with scale and shear matrix
    #
    #       RotateScaleShear(a, s, (sx, sy)) =
    #       = R(a) * S(s) * SHy(sy) * SHx(sx)
    #       = [ s*cos(a - sy)/cos(sy), s*(-cos(a - sy)*tan(sx)/cos(sy) - sin(a)), 0 ]
    #         [ s*sin(a - sy)/cos(sy), s*(-sin(a - sy)*tan(sx)/cos(sy) + cos(a)), 0 ]
    #         [ 0                    , 0                                      , 1 ]
    # where R is a rotation matrix, S is a scaling matrix, and SHx and SHy are the shears:
    # SHx(s) = [1, -tan(s)] and SHy(s) = [1      , 0]
    #          [0, 1      ]              [-tan(s), 1]
    #
    # Thus, the inverse is M^-1 = C * RotateScaleShear^-1 * C^-1 * T^-1

    rot = math.radians(angle)
    sx = math.radians(shear[0])
    sy = math.radians(shear[1])

    cx, cy = center
    tx, ty = translate

    # Cached results
    cos_sy = math.cos(sy)
    tan_sx = math.tan(sx)
    rot_minus_sy = rot - sy
    cx_plus_tx = cx + tx
    cy_plus_ty = cy + ty

    # Rotate Scale Shear (RSS) without scaling
    a = math.cos(rot_minus_sy) / cos_sy
    b = -(a * tan_sx + math.sin(rot))
    c = math.sin(rot_minus_sy) / cos_sy
    d = math.cos(rot) - c * tan_sx

    if inverted:
        # Inverted rotation matrix with scale and shear
        # det([[a, b], [c, d]]) == 1, since det(rotation) = 1 and det(shear) = 1
        matrix = [d / scale, -b / scale, 0.0, -c / scale, a / scale, 0.0]
        # Apply inverse of translation and of center translation: RSS^-1 * C^-1 * T^-1
        # and then apply center translation: C * RSS^-1 * C^-1 * T^-1
        matrix[2] += cx - matrix[0] * cx_plus_tx - matrix[1] * cy_plus_ty
        matrix[5] += cy - matrix[3] * cx_plus_tx - matrix[4] * cy_plus_ty
    else:
        matrix = [a * scale, b * scale, 0.0, c * scale, d * scale, 0.0]
        # Apply inverse of center translation: RSS * C^-1
        # and then apply translation and center : T * C * RSS * C^-1
        matrix[2] += cx_plus_tx - matrix[0] * cx - matrix[1] * cy
        matrix[5] += cy_plus_ty - matrix[3] * cx - matrix[4] * cy

    return matrix


def _compute_affine_output_size(matrix: List[float], w: int, h: int) -> Tuple[int, int]:
    # Inspired of PIL implementation:
    # https://github.com/python-pillow/Pillow/blob/11de3318867e4398057373ee9f12dcb33db7335c/src/PIL/Image.py#L2054

    # pts are Top-Left, Top-Right, Bottom-Left, Bottom-Right points.
    # Points are shifted due to affine matrix torch convention about
    # the center point. Center is (0, 0) for image center pivot point (w * 0.5, h * 0.5)
    half_w = 0.5 * w
    half_h = 0.5 * h
    pts = torch.tensor(
        [
            [-half_w, -half_h, 1.0],
            [-half_w, half_h, 1.0],
            [half_w, half_h, 1.0],
            [half_w, -half_h, 1.0],
        ]
    )
    theta = torch.tensor(matrix, dtype=torch.float).view(2, 3)
    new_pts = torch.matmul(pts, theta.T)
    min_vals, max_vals = new_pts.aminmax(dim=0)

    # shift points to [0, w] and [0, h] interval to match PIL results
    halfs = torch.tensor((half_w, half_h))
    min_vals.add_(halfs)
    max_vals.add_(halfs)

    # Truncate precision to 1e-4 to avoid ceil of Xe-15 to 1.0
    tol = 1e-4
    inv_tol = 1.0 / tol
    cmax = max_vals.mul_(inv_tol).trunc_().mul_(tol).ceil_()
    cmin = min_vals.mul_(inv_tol).trunc_().mul_(tol).floor_()
    size = cmax.sub_(cmin)
    return int(size[0]), int(size[1])  # w, h


def _apply_grid_transform(
Philip Meier's avatar
Philip Meier committed
589
    img: torch.Tensor, grid: torch.Tensor, mode: str, fill: datapoints._FillTypeJIT
590
591
) -> torch.Tensor:

592
593
594
595
    # We are using context knowledge that grid should have float dtype
    fp = img.dtype == grid.dtype
    float_img = img if fp else img.to(grid.dtype)

596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
    shape = float_img.shape
    if shape[0] > 1:
        # Apply same grid to a batch of images
        grid = grid.expand(shape[0], -1, -1, -1)

    # Append a dummy mask for customized fill colors, should be faster than grid_sample() twice
    if fill is not None:
        mask = torch.ones((shape[0], 1, shape[2], shape[3]), dtype=float_img.dtype, device=float_img.device)
        float_img = torch.cat((float_img, mask), dim=1)

    float_img = grid_sample(float_img, grid, mode=mode, padding_mode="zeros", align_corners=False)

    # Fill with required color
    if fill is not None:
        float_img, mask = torch.tensor_split(float_img, indices=(-1,), dim=-3)
        mask = mask.expand_as(float_img)
612
        fill_list = fill if isinstance(fill, (tuple, list)) else [float(fill)]  # type: ignore[arg-type]
613
614
615
616
617
618
619
620
621
        fill_img = torch.tensor(fill_list, dtype=float_img.dtype, device=float_img.device).view(1, -1, 1, 1)
        if mode == "nearest":
            bool_mask = mask < 0.5
            float_img[bool_mask] = fill_img.expand_as(float_img)[bool_mask]
        else:  # 'bilinear'
            # The following is mathematically equivalent to:
            # img * mask + (1.0 - mask) * fill = img * mask - fill * mask + fill = mask * (img - fill) + fill
            float_img = float_img.sub_(fill_img).mul_(mask).add_(fill_img)

622
623
624
    img = float_img.round_().to(img.dtype) if not fp else float_img

    return img
625
626
627
628
629
630


def _assert_grid_transform_inputs(
    image: torch.Tensor,
    matrix: Optional[List[float]],
    interpolation: str,
Philip Meier's avatar
Philip Meier committed
631
    fill: datapoints._FillTypeJIT,
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
    supported_interpolation_modes: List[str],
    coeffs: Optional[List[float]] = None,
) -> None:
    if matrix is not None:
        if not isinstance(matrix, list):
            raise TypeError("Argument matrix should be a list")
        elif len(matrix) != 6:
            raise ValueError("Argument matrix should have 6 float values")

    if coeffs is not None and len(coeffs) != 8:
        raise ValueError("Argument coeffs should have 8 float values")

    if fill is not None:
        if isinstance(fill, (tuple, list)):
            length = len(fill)
            num_channels = image.shape[-3]
            if length > 1 and length != num_channels:
                raise ValueError(
                    "The number of elements in 'fill' cannot broadcast to match the number of "
                    f"channels of the image ({length} != {num_channels})"
                )
        elif not isinstance(fill, (int, float)):
            raise ValueError("Argument fill should be either int, float, tuple or list")

    if interpolation not in supported_interpolation_modes:
        raise ValueError(f"Interpolation mode '{interpolation}' is unsupported with Tensor input")


def _affine_grid(
    theta: torch.Tensor,
    w: int,
    h: int,
    ow: int,
    oh: int,
) -> torch.Tensor:
    # https://github.com/pytorch/pytorch/blob/74b65c32be68b15dc7c9e8bb62459efbfbde33d8/aten/src/ATen/native/
    # AffineGridGenerator.cpp#L18
    # Difference with AffineGridGenerator is that:
    # 1) we normalize grid values after applying theta
    # 2) we can normalize by other image size, such that it covers "extend" option like in PIL.Image.rotate
    dtype = theta.dtype
    device = theta.device

    base_grid = torch.empty(1, oh, ow, 3, dtype=dtype, device=device)
    x_grid = torch.linspace((1.0 - ow) * 0.5, (ow - 1.0) * 0.5, steps=ow, device=device)
    base_grid[..., 0].copy_(x_grid)
    y_grid = torch.linspace((1.0 - oh) * 0.5, (oh - 1.0) * 0.5, steps=oh, device=device).unsqueeze_(-1)
    base_grid[..., 1].copy_(y_grid)
    base_grid[..., 2].fill_(1)

    rescaled_theta = theta.transpose(1, 2).div_(torch.tensor([0.5 * w, 0.5 * h], dtype=dtype, device=device))
    output_grid = base_grid.view(1, oh * ow, 3).bmm(rescaled_theta)
    return output_grid.view(1, oh, ow, 2)


687
@_register_kernel_internal(affine, datapoints.Image)
688
def affine_image_tensor(
689
    image: torch.Tensor,
690
    angle: Union[int, float],
691
692
693
    translate: List[float],
    scale: float,
    shear: List[float],
694
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
Philip Meier's avatar
Philip Meier committed
695
    fill: datapoints._FillTypeJIT = None,
696
697
    center: Optional[List[float]] = None,
) -> torch.Tensor:
698
699
    interpolation = _check_interpolation(interpolation)

700
701
    if image.numel() == 0:
        return image
702

703
    shape = image.shape
704
    ndim = image.ndim
705

706
707
708
709
710
711
712
713
714
715
    if ndim > 4:
        image = image.reshape((-1,) + shape[-3:])
        needs_unsquash = True
    elif ndim == 3:
        image = image.unsqueeze(0)
        needs_unsquash = True
    else:
        needs_unsquash = False

    height, width = shape[-2:]
716
717
718
719
720
    angle, translate, shear, center = _affine_parse_args(angle, translate, scale, shear, interpolation, center)

    center_f = [0.0, 0.0]
    if center is not None:
        # Center values should be in pixel coordinates but translated such that (0, 0) corresponds to image center.
721
        center_f = [(c - s * 0.5) for c, s in zip(center, [width, height])]
722

723
    translate_f = [float(t) for t in translate]
724
725
    matrix = _get_inverse_affine_matrix(center_f, angle, translate_f, scale, shear)

726
727
    _assert_grid_transform_inputs(image, matrix, interpolation.value, fill, ["nearest", "bilinear"])

728
    dtype = image.dtype if torch.is_floating_point(image) else torch.float32
729
730
    theta = torch.tensor(matrix, dtype=dtype, device=image.device).reshape(1, 2, 3)
    grid = _affine_grid(theta, w=width, h=height, ow=width, oh=height)
731
    output = _apply_grid_transform(image, grid, interpolation.value, fill=fill)
732
733
734
735
736

    if needs_unsquash:
        output = output.reshape(shape)

    return output
737
738


739
@torch.jit.unused
740
def affine_image_pil(
741
    image: PIL.Image.Image,
742
    angle: Union[int, float],
743
744
745
    translate: List[float],
    scale: float,
    shear: List[float],
746
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
Philip Meier's avatar
Philip Meier committed
747
    fill: datapoints._FillTypeJIT = None,
748
749
    center: Optional[List[float]] = None,
) -> PIL.Image.Image:
750
    interpolation = _check_interpolation(interpolation)
751
752
753
754
755
756
    angle, translate, shear, center = _affine_parse_args(angle, translate, scale, shear, interpolation, center)

    # center = (img_size[0] * 0.5 + 0.5, img_size[1] * 0.5 + 0.5)
    # it is visually better to estimate the center without 0.5 offset
    # otherwise image rotated by 90 degrees is shifted vs output image of torch.rot90 or F_t.affine
    if center is None:
Philip Meier's avatar
Philip Meier committed
757
        height, width = get_size_image_pil(image)
758
759
760
        center = [width * 0.5, height * 0.5]
    matrix = _get_inverse_affine_matrix(center, angle, translate, scale, shear)

761
    return _FP.affine(image, matrix, interpolation=pil_modes_mapping[interpolation], fill=fill)
762
763


764
765
def _affine_bounding_boxes_with_expand(
    bounding_boxes: torch.Tensor,
766
    format: datapoints.BoundingBoxFormat,
Philip Meier's avatar
Philip Meier committed
767
    canvas_size: Tuple[int, int],
768
769
770
771
    angle: Union[int, float],
    translate: List[float],
    scale: float,
    shear: List[float],
772
    center: Optional[List[float]] = None,
773
    expand: bool = False,
774
) -> Tuple[torch.Tensor, Tuple[int, int]]:
775
    if bounding_boxes.numel() == 0:
Philip Meier's avatar
Philip Meier committed
776
        return bounding_boxes, canvas_size
777
778
779
780
781
782
783
784
785

    original_shape = bounding_boxes.shape
    original_dtype = bounding_boxes.dtype
    bounding_boxes = bounding_boxes.clone() if bounding_boxes.is_floating_point() else bounding_boxes.float()
    dtype = bounding_boxes.dtype
    device = bounding_boxes.device
    bounding_boxes = (
        convert_format_bounding_boxes(
            bounding_boxes, old_format=format, new_format=datapoints.BoundingBoxFormat.XYXY, inplace=True
786
787
788
        )
    ).reshape(-1, 4)

789
790
791
    angle, translate, shear, center = _affine_parse_args(
        angle, translate, scale, shear, InterpolationMode.NEAREST, center
    )
792

793
    if center is None:
Philip Meier's avatar
Philip Meier committed
794
        height, width = canvas_size
795
796
        center = [width * 0.5, height * 0.5]

797
798
799
800
801
802
803
    affine_vector = _get_inverse_affine_matrix(center, angle, translate, scale, shear, inverted=False)
    transposed_affine_matrix = (
        torch.tensor(
            affine_vector,
            dtype=dtype,
            device=device,
        )
804
        .reshape(2, 3)
805
806
        .T
    )
807
808
809
810
    # 1) Let's transform bboxes into a tensor of 4 points (top-left, top-right, bottom-left, bottom-right corners).
    # Tensor of points has shape (N * 4, 3), where N is the number of bboxes
    # Single point structure is similar to
    # [(xmin, ymin, 1), (xmax, ymin, 1), (xmax, ymax, 1), (xmin, ymax, 1)]
811
    points = bounding_boxes[:, [[0, 1], [2, 1], [2, 3], [0, 3]]].reshape(-1, 2)
812
    points = torch.cat([points, torch.ones(points.shape[0], 1, device=device, dtype=dtype)], dim=-1)
813
    # 2) Now let's transform the points using affine matrix
814
    transformed_points = torch.matmul(points, transposed_affine_matrix)
815
816
    # 3) Reshape transformed points to [N boxes, 4 points, x/y coords]
    # and compute bounding box from 4 transformed points:
817
    transformed_points = transformed_points.reshape(-1, 4, 2)
818
    out_bbox_mins, out_bbox_maxs = torch.aminmax(transformed_points, dim=1)
819
    out_bboxes = torch.cat([out_bbox_mins, out_bbox_maxs], dim=1)
820
821
822
823

    if expand:
        # Compute minimum point for transformed image frame:
        # Points are Top-Left, Top-Right, Bottom-Left, Bottom-Right points.
Philip Meier's avatar
Philip Meier committed
824
        height, width = canvas_size
825
826
827
        points = torch.tensor(
            [
                [0.0, 0.0, 1.0],
828
829
830
                [0.0, float(height), 1.0],
                [float(width), float(height), 1.0],
                [float(width), 0.0, 1.0],
831
832
833
834
            ],
            dtype=dtype,
            device=device,
        )
835
        new_points = torch.matmul(points, transposed_affine_matrix)
836
        tr = torch.amin(new_points, dim=0, keepdim=True)
837
        # Translate bounding boxes
838
        out_bboxes.sub_(tr.repeat((1, 2)))
839
840
        # Estimate meta-data for image with inverted=True and with center=[0,0]
        affine_vector = _get_inverse_affine_matrix([0.0, 0.0], angle, translate, scale, shear)
841
        new_width, new_height = _compute_affine_output_size(affine_vector, width, height)
Philip Meier's avatar
Philip Meier committed
842
        canvas_size = (new_height, new_width)
843

Philip Meier's avatar
Philip Meier committed
844
    out_bboxes = clamp_bounding_boxes(out_bboxes, format=datapoints.BoundingBoxFormat.XYXY, canvas_size=canvas_size)
845
    out_bboxes = convert_format_bounding_boxes(
846
847
848
849
        out_bboxes, old_format=datapoints.BoundingBoxFormat.XYXY, new_format=format, inplace=True
    ).reshape(original_shape)

    out_bboxes = out_bboxes.to(original_dtype)
Philip Meier's avatar
Philip Meier committed
850
    return out_bboxes, canvas_size
851
852


853
854
def affine_bounding_boxes(
    bounding_boxes: torch.Tensor,
855
    format: datapoints.BoundingBoxFormat,
Philip Meier's avatar
Philip Meier committed
856
    canvas_size: Tuple[int, int],
857
    angle: Union[int, float],
858
859
860
861
862
    translate: List[float],
    scale: float,
    shear: List[float],
    center: Optional[List[float]] = None,
) -> torch.Tensor:
863
864
    out_box, _ = _affine_bounding_boxes_with_expand(
        bounding_boxes,
865
        format=format,
Philip Meier's avatar
Philip Meier committed
866
        canvas_size=canvas_size,
867
868
869
870
871
872
873
874
        angle=angle,
        translate=translate,
        scale=scale,
        shear=shear,
        center=center,
        expand=False,
    )
    return out_box
875
876


877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
@_register_kernel_internal(affine, datapoints.BoundingBoxes, datapoint_wrapper=False)
def _affine_bounding_boxes_dispatch(
    inpt: datapoints.BoundingBoxes,
    angle: Union[int, float],
    translate: List[float],
    scale: float,
    shear: List[float],
    center: Optional[List[float]] = None,
    **kwargs,
) -> datapoints.BoundingBoxes:
    output = affine_bounding_boxes(
        inpt.as_subclass(torch.Tensor),
        format=inpt.format,
        canvas_size=inpt.canvas_size,
        angle=angle,
        translate=translate,
        scale=scale,
        shear=shear,
        center=center,
    )
    return datapoints.BoundingBoxes.wrap_like(inpt, output)


900
901
def affine_mask(
    mask: torch.Tensor,
902
    angle: Union[int, float],
903
904
905
    translate: List[float],
    scale: float,
    shear: List[float],
Philip Meier's avatar
Philip Meier committed
906
    fill: datapoints._FillTypeJIT = None,
907
908
    center: Optional[List[float]] = None,
) -> torch.Tensor:
909
910
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
911
912
913
914
915
        needs_squeeze = True
    else:
        needs_squeeze = False

    output = affine_image_tensor(
916
        mask,
917
918
919
920
921
        angle=angle,
        translate=translate,
        scale=scale,
        shear=shear,
        interpolation=InterpolationMode.NEAREST,
922
        fill=fill,
923
924
925
        center=center,
    )

926
927
928
929
930
    if needs_squeeze:
        output = output.squeeze(0)

    return output

931

932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
@_register_kernel_internal(affine, datapoints.Mask, datapoint_wrapper=False)
def _affine_mask_dispatch(
    inpt: datapoints.Mask,
    angle: Union[int, float],
    translate: List[float],
    scale: float,
    shear: List[float],
    fill: datapoints._FillTypeJIT = None,
    center: Optional[List[float]] = None,
    **kwargs,
) -> datapoints.Mask:
    output = affine_mask(
        inpt.as_subclass(torch.Tensor),
        angle=angle,
        translate=translate,
        scale=scale,
        shear=shear,
        fill=fill,
        center=center,
    )
    return datapoints.Mask.wrap_like(inpt, output)


@_register_kernel_internal(affine, datapoints.Video)
956
957
958
959
960
961
def affine_video(
    video: torch.Tensor,
    angle: Union[int, float],
    translate: List[float],
    scale: float,
    shear: List[float],
962
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
Philip Meier's avatar
Philip Meier committed
963
    fill: datapoints._FillTypeJIT = None,
964
965
966
967
968
969
970
971
972
973
974
975
976
977
    center: Optional[List[float]] = None,
) -> torch.Tensor:
    return affine_image_tensor(
        video,
        angle=angle,
        translate=translate,
        scale=scale,
        shear=shear,
        interpolation=interpolation,
        fill=fill,
        center=center,
    )


978
def rotate(
Philip Meier's avatar
Philip Meier committed
979
    inpt: datapoints._InputTypeJIT,
980
    angle: float,
981
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
982
    expand: bool = False,
983
    center: Optional[List[float]] = None,
984
    fill: datapoints._FillTypeJIT = None,
Philip Meier's avatar
Philip Meier committed
985
) -> datapoints._InputTypeJIT:
986
    if not torch.jit.is_scripting():
987
        _log_api_usage_once(rotate)
988

989
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
990
991
992
993
        return rotate_image_tensor(inpt, angle, interpolation=interpolation, expand=expand, fill=fill, center=center)
    elif isinstance(inpt, datapoints.Datapoint):
        kernel = _get_kernel(rotate, type(inpt))
        return kernel(inpt, angle, interpolation=interpolation, expand=expand, fill=fill, center=center)
994
    elif isinstance(inpt, PIL.Image.Image):
995
        return rotate_image_pil(inpt, angle, interpolation=interpolation, expand=expand, fill=fill, center=center)
996
997
    else:
        raise TypeError(
998
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
999
1000
            f"but got {type(inpt)} instead."
        )
1001
1002


1003
@_register_kernel_internal(rotate, datapoints.Image)
1004
def rotate_image_tensor(
1005
    image: torch.Tensor,
1006
    angle: float,
1007
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
1008
1009
    expand: bool = False,
    center: Optional[List[float]] = None,
Philip Meier's avatar
Philip Meier committed
1010
    fill: datapoints._FillTypeJIT = None,
1011
) -> torch.Tensor:
1012
1013
    interpolation = _check_interpolation(interpolation)

1014
1015
    shape = image.shape
    num_channels, height, width = shape[-3:]
1016

1017
1018
    center_f = [0.0, 0.0]
    if center is not None:
1019
        if expand:
1020
            # TODO: Do we actually want to warn, or just document this?
1021
            warnings.warn("The provided center argument has no effect on the result if expand is True")
1022
1023
        # Center values should be in pixel coordinates but translated such that (0, 0) corresponds to image center.
        center_f = [(c - s * 0.5) for c, s in zip(center, [width, height])]
1024
1025
1026
1027

    # due to current incoherence of rotation angle direction between affine and rotate implementations
    # we need to set -angle.
    matrix = _get_inverse_affine_matrix(center_f, -angle, [0.0, 0.0], 1.0, [0.0, 0.0])
1028

1029
    if image.numel() > 0:
1030
1031
1032
1033
1034
        image = image.reshape(-1, num_channels, height, width)

        _assert_grid_transform_inputs(image, matrix, interpolation.value, fill, ["nearest", "bilinear"])

        ow, oh = _compute_affine_output_size(matrix, width, height) if expand else (width, height)
1035
        dtype = image.dtype if torch.is_floating_point(image) else torch.float32
1036
1037
        theta = torch.tensor(matrix, dtype=dtype, device=image.device).reshape(1, 2, 3)
        grid = _affine_grid(theta, w=width, h=height, ow=ow, oh=oh)
1038
        output = _apply_grid_transform(image, grid, interpolation.value, fill=fill)
1039
1040

        new_height, new_width = output.shape[-2:]
1041
    else:
1042
1043
        output = image
        new_width, new_height = _compute_affine_output_size(matrix, width, height) if expand else (width, height)
1044

1045
    return output.reshape(shape[:-3] + (num_channels, new_height, new_width))
1046
1047


1048
@torch.jit.unused
1049
def rotate_image_pil(
1050
    image: PIL.Image.Image,
1051
    angle: float,
1052
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
1053
1054
    expand: bool = False,
    center: Optional[List[float]] = None,
Philip Meier's avatar
Philip Meier committed
1055
    fill: datapoints._FillTypeJIT = None,
1056
) -> PIL.Image.Image:
1057
1058
    interpolation = _check_interpolation(interpolation)

1059
    if center is not None and expand:
1060
        warnings.warn("The provided center argument has no effect on the result if expand is True")
1061

1062
    return _FP.rotate(
1063
        image, angle, interpolation=pil_modes_mapping[interpolation], expand=expand, fill=fill, center=center
1064
1065
1066
    )


1067
1068
def rotate_bounding_boxes(
    bounding_boxes: torch.Tensor,
1069
    format: datapoints.BoundingBoxFormat,
Philip Meier's avatar
Philip Meier committed
1070
    canvas_size: Tuple[int, int],
1071
1072
1073
    angle: float,
    expand: bool = False,
    center: Optional[List[float]] = None,
1074
) -> Tuple[torch.Tensor, Tuple[int, int]]:
1075
1076
1077
    if center is not None and expand:
        warnings.warn("The provided center argument has no effect on the result if expand is True")

1078
1079
    return _affine_bounding_boxes_with_expand(
        bounding_boxes,
1080
        format=format,
Philip Meier's avatar
Philip Meier committed
1081
        canvas_size=canvas_size,
1082
1083
1084
1085
1086
1087
1088
        angle=-angle,
        translate=[0.0, 0.0],
        scale=1.0,
        shear=[0.0, 0.0],
        center=center,
        expand=expand,
    )
1089
1090


1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
@_register_kernel_internal(rotate, datapoints.BoundingBoxes, datapoint_wrapper=False)
def _rotate_bounding_boxes_dispatch(
    inpt: datapoints.BoundingBoxes, angle: float, expand: bool = False, center: Optional[List[float]] = None, **kwargs
) -> datapoints.BoundingBoxes:
    output, canvas_size = rotate_bounding_boxes(
        inpt.as_subclass(torch.Tensor),
        format=inpt.format,
        canvas_size=inpt.canvas_size,
        angle=angle,
        expand=expand,
        center=center,
    )
    return datapoints.BoundingBoxes.wrap_like(inpt, output, canvas_size=canvas_size)


1106
1107
def rotate_mask(
    mask: torch.Tensor,
1108
1109
1110
    angle: float,
    expand: bool = False,
    center: Optional[List[float]] = None,
Philip Meier's avatar
Philip Meier committed
1111
    fill: datapoints._FillTypeJIT = None,
1112
) -> torch.Tensor:
1113
1114
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
1115
1116
1117
1118
1119
        needs_squeeze = True
    else:
        needs_squeeze = False

    output = rotate_image_tensor(
1120
        mask,
1121
1122
1123
        angle=angle,
        expand=expand,
        interpolation=InterpolationMode.NEAREST,
1124
        fill=fill,
1125
1126
1127
        center=center,
    )

1128
1129
1130
1131
1132
    if needs_squeeze:
        output = output.squeeze(0)

    return output

1133

1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
@_register_kernel_internal(rotate, datapoints.Mask, datapoint_wrapper=False)
def _rotate_mask_dispatch(
    inpt: datapoints.Mask,
    angle: float,
    expand: bool = False,
    center: Optional[List[float]] = None,
    fill: datapoints._FillTypeJIT = None,
    **kwargs,
) -> datapoints.Mask:
    output = rotate_mask(inpt.as_subclass(torch.Tensor), angle=angle, expand=expand, fill=fill, center=center)
    return datapoints.Mask.wrap_like(inpt, output)


@_register_kernel_internal(rotate, datapoints.Video)
1148
1149
1150
def rotate_video(
    video: torch.Tensor,
    angle: float,
1151
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
1152
1153
    expand: bool = False,
    center: Optional[List[float]] = None,
Philip Meier's avatar
Philip Meier committed
1154
    fill: datapoints._FillTypeJIT = None,
1155
1156
1157
1158
) -> torch.Tensor:
    return rotate_image_tensor(video, angle, interpolation=interpolation, expand=expand, fill=fill, center=center)


1159
def pad(
Philip Meier's avatar
Philip Meier committed
1160
    inpt: datapoints._InputTypeJIT,
1161
1162
1163
    padding: List[int],
    fill: Optional[Union[int, float, List[float]]] = None,
    padding_mode: str = "constant",
Philip Meier's avatar
Philip Meier committed
1164
) -> datapoints._InputTypeJIT:
1165
    if not torch.jit.is_scripting():
1166
        _log_api_usage_once(pad)
1167

1168
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
1169
1170
1171
1172
1173
        return pad_image_tensor(inpt, padding, fill=fill, padding_mode=padding_mode)

    elif isinstance(inpt, datapoints.Datapoint):
        kernel = _get_kernel(pad, type(inpt))
        return kernel(inpt, padding, fill=fill, padding_mode=padding_mode)
1174
    elif isinstance(inpt, PIL.Image.Image):
1175
        return pad_image_pil(inpt, padding, fill=fill, padding_mode=padding_mode)
1176
1177
    else:
        raise TypeError(
1178
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
1179
1180
            f"but got {type(inpt)} instead."
        )
1181
1182


1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
def _parse_pad_padding(padding: Union[int, List[int]]) -> List[int]:
    if isinstance(padding, int):
        pad_left = pad_right = pad_top = pad_bottom = padding
    elif isinstance(padding, (tuple, list)):
        if len(padding) == 1:
            pad_left = pad_right = pad_top = pad_bottom = padding[0]
        elif len(padding) == 2:
            pad_left = pad_right = padding[0]
            pad_top = pad_bottom = padding[1]
        elif len(padding) == 4:
            pad_left = padding[0]
            pad_top = padding[1]
            pad_right = padding[2]
            pad_bottom = padding[3]
        else:
            raise ValueError(
                f"Padding must be an int or a 1, 2, or 4 element tuple, not a {len(padding)} element tuple"
            )
    else:
        raise TypeError(f"`padding` should be an integer or tuple or list of integers, but got {padding}")

    return [pad_left, pad_right, pad_top, pad_bottom]
1205

1206

1207
@_register_kernel_internal(pad, datapoints.Image)
1208
def pad_image_tensor(
1209
    image: torch.Tensor,
1210
1211
    padding: List[int],
    fill: Optional[Union[int, float, List[float]]] = None,
1212
1213
    padding_mode: str = "constant",
) -> torch.Tensor:
1214
1215
1216
1217
1218
    # Be aware that while `padding` has order `[left, top, right, bottom]` has order, `torch_padding` uses
    # `[left, right, top, bottom]`. This stems from the fact that we align our API with PIL, but need to use `torch_pad`
    # internally.
    torch_padding = _parse_pad_padding(padding)

1219
    if padding_mode not in ("constant", "edge", "reflect", "symmetric"):
1220
1221
1222
1223
1224
        raise ValueError(
            f"`padding_mode` should be either `'constant'`, `'edge'`, `'reflect'` or `'symmetric'`, "
            f"but got `'{padding_mode}'`."
        )

1225
    if fill is None:
1226
1227
1228
1229
1230
1231
        fill = 0

    if isinstance(fill, (int, float)):
        return _pad_with_scalar_fill(image, torch_padding, fill=fill, padding_mode=padding_mode)
    elif len(fill) == 1:
        return _pad_with_scalar_fill(image, torch_padding, fill=fill[0], padding_mode=padding_mode)
1232
    else:
1233
        return _pad_with_vector_fill(image, torch_padding, fill=fill, padding_mode=padding_mode)
1234
1235
1236


def _pad_with_scalar_fill(
1237
    image: torch.Tensor,
1238
1239
1240
    torch_padding: List[int],
    fill: Union[int, float],
    padding_mode: str,
1241
) -> torch.Tensor:
1242
1243
    shape = image.shape
    num_channels, height, width = shape[-3:]
1244

1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
    batch_size = 1
    for s in shape[:-3]:
        batch_size *= s

    image = image.reshape(batch_size, num_channels, height, width)

    if padding_mode == "edge":
        # Similar to the padding order, `torch_pad`'s PIL's padding modes don't have the same names. Thus, we map
        # the PIL name for the padding mode, which we are also using for our API, to the corresponding `torch_pad`
        # name.
        padding_mode = "replicate"

    if padding_mode == "constant":
        image = torch_pad(image, torch_padding, mode=padding_mode, value=float(fill))
    elif padding_mode in ("reflect", "replicate"):
        # `torch_pad` only supports `"reflect"` or `"replicate"` padding for floating point inputs.
        # TODO: See https://github.com/pytorch/pytorch/issues/40763
        dtype = image.dtype
        if not image.is_floating_point():
            needs_cast = True
            image = image.to(torch.float32)
        else:
            needs_cast = False
1268

1269
1270
1271
1272
1273
        image = torch_pad(image, torch_padding, mode=padding_mode)

        if needs_cast:
            image = image.to(dtype)
    else:  # padding_mode == "symmetric"
1274
        image = _pad_symmetric(image, torch_padding)
1275
1276

    new_height, new_width = image.shape[-2:]
1277

1278
    return image.reshape(shape[:-3] + (num_channels, new_height, new_width))
1279
1280


1281
# TODO: This should be removed once torch_pad supports non-scalar padding values
1282
def _pad_with_vector_fill(
1283
    image: torch.Tensor,
1284
    torch_padding: List[int],
1285
    fill: List[float],
1286
    padding_mode: str,
1287
1288
1289
1290
) -> torch.Tensor:
    if padding_mode != "constant":
        raise ValueError(f"Padding mode '{padding_mode}' is not supported if fill is not scalar")

1291
1292
    output = _pad_with_scalar_fill(image, torch_padding, fill=0, padding_mode="constant")
    left, right, top, bottom = torch_padding
1293
    fill = torch.tensor(fill, dtype=image.dtype, device=image.device).reshape(-1, 1, 1)
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305

    if top > 0:
        output[..., :top, :] = fill
    if left > 0:
        output[..., :, :left] = fill
    if bottom > 0:
        output[..., -bottom:, :] = fill
    if right > 0:
        output[..., :, -right:] = fill
    return output


1306
1307
1308
pad_image_pil = _FP.pad


1309
@_register_kernel_internal(pad, datapoints.Mask)
1310
1311
def pad_mask(
    mask: torch.Tensor,
1312
1313
    padding: List[int],
    fill: Optional[Union[int, float, List[float]]] = None,
1314
1315
    padding_mode: str = "constant",
) -> torch.Tensor:
1316
1317
1318
    if fill is None:
        fill = 0

1319
    if isinstance(fill, (tuple, list)):
1320
1321
        raise ValueError("Non-scalar fill value is not supported")

1322
1323
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
1324
1325
1326
1327
        needs_squeeze = True
    else:
        needs_squeeze = False

1328
    output = pad_image_tensor(mask, padding=padding, fill=fill, padding_mode=padding_mode)
1329
1330
1331
1332
1333

    if needs_squeeze:
        output = output.squeeze(0)

    return output
1334
1335


1336
1337
def pad_bounding_boxes(
    bounding_boxes: torch.Tensor,
1338
    format: datapoints.BoundingBoxFormat,
Philip Meier's avatar
Philip Meier committed
1339
    canvas_size: Tuple[int, int],
1340
    padding: List[int],
vfdev's avatar
vfdev committed
1341
    padding_mode: str = "constant",
1342
) -> Tuple[torch.Tensor, Tuple[int, int]]:
vfdev's avatar
vfdev committed
1343
1344
1345
1346
    if padding_mode not in ["constant"]:
        # TODO: add support of other padding modes
        raise ValueError(f"Padding mode '{padding_mode}' is not supported with bounding boxes")

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

1349
    if format == datapoints.BoundingBoxFormat.XYXY:
1350
1351
1352
        pad = [left, top, left, top]
    else:
        pad = [left, top, 0, 0]
1353
    bounding_boxes = bounding_boxes + torch.tensor(pad, dtype=bounding_boxes.dtype, device=bounding_boxes.device)
1354

Philip Meier's avatar
Philip Meier committed
1355
    height, width = canvas_size
1356
1357
    height += top + bottom
    width += left + right
Philip Meier's avatar
Philip Meier committed
1358
    canvas_size = (height, width)
1359

Philip Meier's avatar
Philip Meier committed
1360
    return clamp_bounding_boxes(bounding_boxes, format=format, canvas_size=canvas_size), canvas_size
1361
1362


1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
@_register_kernel_internal(pad, datapoints.BoundingBoxes, datapoint_wrapper=False)
def _pad_bounding_boxes_dispatch(
    inpt: datapoints.BoundingBoxes, padding: List[int], padding_mode: str = "constant", **kwargs
) -> datapoints.BoundingBoxes:
    output, canvas_size = pad_bounding_boxes(
        inpt.as_subclass(torch.Tensor),
        format=inpt.format,
        canvas_size=inpt.canvas_size,
        padding=padding,
        padding_mode=padding_mode,
    )
    return datapoints.BoundingBoxes.wrap_like(inpt, output, canvas_size=canvas_size)


@_register_kernel_internal(pad, datapoints.Video)
1378
1379
def pad_video(
    video: torch.Tensor,
1380
1381
    padding: List[int],
    fill: Optional[Union[int, float, List[float]]] = None,
1382
1383
1384
1385
1386
    padding_mode: str = "constant",
) -> torch.Tensor:
    return pad_image_tensor(video, padding, fill=fill, padding_mode=padding_mode)


1387
def crop(inpt: datapoints._InputTypeJIT, top: int, left: int, height: int, width: int) -> datapoints._InputTypeJIT:
1388
    if not torch.jit.is_scripting():
1389
        _log_api_usage_once(crop)
1390

1391
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
1392
1393
1394
1395
        return crop_image_tensor(inpt, top, left, height, width)
    elif isinstance(inpt, datapoints.Datapoint):
        kernel = _get_kernel(crop, type(inpt))
        return kernel(inpt, top, left, height, width)
1396
    elif isinstance(inpt, PIL.Image.Image):
1397
        return crop_image_pil(inpt, top, left, height, width)
1398
1399
    else:
        raise TypeError(
1400
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
1401
1402
            f"but got {type(inpt)} instead."
        )
1403
1404


1405
@_register_kernel_internal(crop, datapoints.Image)
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
def crop_image_tensor(image: torch.Tensor, top: int, left: int, height: int, width: int) -> torch.Tensor:
    h, w = image.shape[-2:]

    right = left + width
    bottom = top + height

    if left < 0 or top < 0 or right > w or bottom > h:
        image = image[..., max(top, 0) : bottom, max(left, 0) : right]
        torch_padding = [
            max(min(right, 0) - left, 0),
            max(right - max(w, left), 0),
            max(min(bottom, 0) - top, 0),
            max(bottom - max(h, top), 0),
        ]
        return _pad_with_scalar_fill(image, torch_padding, fill=0, padding_mode="constant")
    return image[..., top:bottom, left:right]


1424
1425
1426
crop_image_pil = _FP.crop


1427
1428
def crop_bounding_boxes(
    bounding_boxes: torch.Tensor,
1429
    format: datapoints.BoundingBoxFormat,
1430
1431
    top: int,
    left: int,
1432
1433
1434
    height: int,
    width: int,
) -> Tuple[torch.Tensor, Tuple[int, int]]:
1435

1436
    # Crop or implicit pad if left and/or top have negative values:
1437
    if format == datapoints.BoundingBoxFormat.XYXY:
1438
        sub = [left, top, left, top]
1439
    else:
1440
1441
        sub = [left, top, 0, 0]

1442
    bounding_boxes = bounding_boxes - torch.tensor(sub, dtype=bounding_boxes.dtype, device=bounding_boxes.device)
Philip Meier's avatar
Philip Meier committed
1443
    canvas_size = (height, width)
1444

Philip Meier's avatar
Philip Meier committed
1445
    return clamp_bounding_boxes(bounding_boxes, format=format, canvas_size=canvas_size), canvas_size
1446
1447


1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
@_register_kernel_internal(crop, datapoints.BoundingBoxes, datapoint_wrapper=False)
def _crop_bounding_boxes_dispatch(
    inpt: datapoints.BoundingBoxes, top: int, left: int, height: int, width: int
) -> datapoints.BoundingBoxes:
    output, canvas_size = crop_bounding_boxes(
        inpt.as_subclass(torch.Tensor), format=inpt.format, top=top, left=left, height=height, width=width
    )
    return datapoints.BoundingBoxes.wrap_like(inpt, output, canvas_size=canvas_size)


@_register_kernel_internal(crop, datapoints.Mask)
1459
def crop_mask(mask: torch.Tensor, top: int, left: int, height: int, width: int) -> torch.Tensor:
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
        needs_squeeze = True
    else:
        needs_squeeze = False

    output = crop_image_tensor(mask, top, left, height, width)

    if needs_squeeze:
        output = output.squeeze(0)

    return output
1472
1473


1474
@_register_kernel_internal(crop, datapoints.Video)
1475
1476
1477
1478
def crop_video(video: torch.Tensor, top: int, left: int, height: int, width: int) -> torch.Tensor:
    return crop_image_tensor(video, top, left, height, width)


1479
1480
1481
1482
1483
1484
1485
1486
def perspective(
    inpt: datapoints._InputTypeJIT,
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
    fill: datapoints._FillTypeJIT = None,
    coefficients: Optional[List[float]] = None,
) -> datapoints._InputTypeJIT:
1487
    if not torch.jit.is_scripting():
1488
        _log_api_usage_once(perspective)
1489
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
1490
1491
1492
1493
1494
1495
        return perspective_image_tensor(
            inpt, startpoints, endpoints, interpolation=interpolation, fill=fill, coefficients=coefficients
        )
    elif isinstance(inpt, datapoints.Datapoint):
        kernel = _get_kernel(perspective, type(inpt))
        return kernel(inpt, startpoints, endpoints, interpolation=interpolation, fill=fill, coefficients=coefficients)
1496
    elif isinstance(inpt, PIL.Image.Image):
1497
1498
1499
        return perspective_image_pil(
            inpt, startpoints, endpoints, interpolation=interpolation, fill=fill, coefficients=coefficients
        )
1500
1501
    else:
        raise TypeError(
1502
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
1503
1504
            f"but got {type(inpt)} instead."
        )
1505
1506


1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
def _perspective_grid(coeffs: List[float], ow: int, oh: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
    # https://github.com/python-pillow/Pillow/blob/4634eafe3c695a014267eefdce830b4a825beed7/
    # src/libImaging/Geometry.c#L394

    #
    # x_out = (coeffs[0] * x + coeffs[1] * y + coeffs[2]) / (coeffs[6] * x + coeffs[7] * y + 1)
    # y_out = (coeffs[3] * x + coeffs[4] * y + coeffs[5]) / (coeffs[6] * x + coeffs[7] * y + 1)
    #
    theta1 = torch.tensor(
        [[[coeffs[0], coeffs[1], coeffs[2]], [coeffs[3], coeffs[4], coeffs[5]]]], dtype=dtype, device=device
    )
    theta2 = torch.tensor([[[coeffs[6], coeffs[7], 1.0], [coeffs[6], coeffs[7], 1.0]]], dtype=dtype, device=device)

    d = 0.5
    base_grid = torch.empty(1, oh, ow, 3, dtype=dtype, device=device)
1522
    x_grid = torch.linspace(d, ow + d - 1.0, steps=ow, device=device, dtype=dtype)
1523
    base_grid[..., 0].copy_(x_grid)
1524
    y_grid = torch.linspace(d, oh + d - 1.0, steps=oh, device=device, dtype=dtype).unsqueeze_(-1)
1525
1526
1527
1528
    base_grid[..., 1].copy_(y_grid)
    base_grid[..., 2].fill_(1)

    rescaled_theta1 = theta1.transpose(1, 2).div_(torch.tensor([0.5 * ow, 0.5 * oh], dtype=dtype, device=device))
1529
1530
1531
    shape = (1, oh * ow, 3)
    output_grid1 = base_grid.view(shape).bmm(rescaled_theta1)
    output_grid2 = base_grid.view(shape).bmm(theta2.transpose(1, 2))
1532
1533
1534
1535
1536

    output_grid = output_grid1.div_(output_grid2).sub_(1.0)
    return output_grid.view(1, oh, ow, 2)


1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
def _perspective_coefficients(
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
    coefficients: Optional[List[float]],
) -> List[float]:
    if coefficients is not None:
        if startpoints is not None and endpoints is not None:
            raise ValueError("The startpoints/endpoints and the coefficients shouldn't be defined concurrently.")
        elif len(coefficients) != 8:
            raise ValueError("Argument coefficients should have 8 float values")
        return coefficients
    elif startpoints is not None and endpoints is not None:
        return _get_perspective_coeffs(startpoints, endpoints)
    else:
        raise ValueError("Either the startpoints/endpoints or the coefficients must have non `None` values.")


1554
@_register_kernel_internal(perspective, datapoints.Image)
1555
def perspective_image_tensor(
1556
    image: torch.Tensor,
1557
1558
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
1559
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
Philip Meier's avatar
Philip Meier committed
1560
    fill: datapoints._FillTypeJIT = None,
1561
    coefficients: Optional[List[float]] = None,
1562
) -> torch.Tensor:
1563
    perspective_coeffs = _perspective_coefficients(startpoints, endpoints, coefficients)
1564
1565
    interpolation = _check_interpolation(interpolation)

1566
1567
1568
1569
    if image.numel() == 0:
        return image

    shape = image.shape
1570
    ndim = image.ndim
1571

1572
    if ndim > 4:
1573
        image = image.reshape((-1,) + shape[-3:])
1574
        needs_unsquash = True
1575
1576
1577
    elif ndim == 3:
        image = image.unsqueeze(0)
        needs_unsquash = True
1578
1579
1580
    else:
        needs_unsquash = False

1581
    _assert_grid_transform_inputs(
1582
1583
1584
1585
1586
1587
1588
1589
        image,
        matrix=None,
        interpolation=interpolation.value,
        fill=fill,
        supported_interpolation_modes=["nearest", "bilinear"],
        coeffs=perspective_coeffs,
    )

1590
    oh, ow = shape[-2:]
1591
    dtype = image.dtype if torch.is_floating_point(image) else torch.float32
1592
    grid = _perspective_grid(perspective_coeffs, ow=ow, oh=oh, dtype=dtype, device=image.device)
1593
    output = _apply_grid_transform(image, grid, interpolation.value, fill=fill)
1594
1595

    if needs_unsquash:
1596
        output = output.reshape(shape)
1597
1598

    return output
1599
1600


1601
@torch.jit.unused
1602
def perspective_image_pil(
1603
    image: PIL.Image.Image,
1604
1605
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
1606
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BICUBIC,
Philip Meier's avatar
Philip Meier committed
1607
    fill: datapoints._FillTypeJIT = None,
1608
    coefficients: Optional[List[float]] = None,
1609
) -> PIL.Image.Image:
1610
    perspective_coeffs = _perspective_coefficients(startpoints, endpoints, coefficients)
1611
    interpolation = _check_interpolation(interpolation)
1612
    return _FP.perspective(image, perspective_coeffs, interpolation=pil_modes_mapping[interpolation], fill=fill)
1613
1614


1615
1616
def perspective_bounding_boxes(
    bounding_boxes: torch.Tensor,
1617
    format: datapoints.BoundingBoxFormat,
Philip Meier's avatar
Philip Meier committed
1618
    canvas_size: Tuple[int, int],
1619
1620
1621
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
    coefficients: Optional[List[float]] = None,
1622
) -> torch.Tensor:
1623
1624
    if bounding_boxes.numel() == 0:
        return bounding_boxes
1625

1626
    perspective_coeffs = _perspective_coefficients(startpoints, endpoints, coefficients)
1627

1628
1629
1630
1631
    original_shape = bounding_boxes.shape
    # TODO: first cast to float if bbox is int64 before convert_format_bounding_boxes
    bounding_boxes = (
        convert_format_bounding_boxes(bounding_boxes, old_format=format, new_format=datapoints.BoundingBoxFormat.XYXY)
1632
    ).reshape(-1, 4)
1633

1634
1635
    dtype = bounding_boxes.dtype if torch.is_floating_point(bounding_boxes) else torch.float32
    device = bounding_boxes.device
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666

    # perspective_coeffs are computed as endpoint -> start point
    # We have to invert perspective_coeffs for bboxes:
    # (x, y) - end point and (x_out, y_out) - start point
    #   x_out = (coeffs[0] * x + coeffs[1] * y + coeffs[2]) / (coeffs[6] * x + coeffs[7] * y + 1)
    #   y_out = (coeffs[3] * x + coeffs[4] * y + coeffs[5]) / (coeffs[6] * x + coeffs[7] * y + 1)
    # and we would like to get:
    # x = (inv_coeffs[0] * x_out + inv_coeffs[1] * y_out + inv_coeffs[2])
    #       / (inv_coeffs[6] * x_out + inv_coeffs[7] * y_out + 1)
    # y = (inv_coeffs[3] * x_out + inv_coeffs[4] * y_out + inv_coeffs[5])
    #       / (inv_coeffs[6] * x_out + inv_coeffs[7] * y_out + 1)
    # and compute inv_coeffs in terms of coeffs

    denom = perspective_coeffs[0] * perspective_coeffs[4] - perspective_coeffs[1] * perspective_coeffs[3]
    if denom == 0:
        raise RuntimeError(
            f"Provided perspective_coeffs {perspective_coeffs} can not be inverted to transform bounding boxes. "
            f"Denominator is zero, denom={denom}"
        )

    inv_coeffs = [
        (perspective_coeffs[4] - perspective_coeffs[5] * perspective_coeffs[7]) / denom,
        (-perspective_coeffs[1] + perspective_coeffs[2] * perspective_coeffs[7]) / denom,
        (perspective_coeffs[1] * perspective_coeffs[5] - perspective_coeffs[2] * perspective_coeffs[4]) / denom,
        (-perspective_coeffs[3] + perspective_coeffs[5] * perspective_coeffs[6]) / denom,
        (perspective_coeffs[0] - perspective_coeffs[2] * perspective_coeffs[6]) / denom,
        (-perspective_coeffs[0] * perspective_coeffs[5] + perspective_coeffs[2] * perspective_coeffs[3]) / denom,
        (-perspective_coeffs[4] * perspective_coeffs[6] + perspective_coeffs[3] * perspective_coeffs[7]) / denom,
        (-perspective_coeffs[0] * perspective_coeffs[7] + perspective_coeffs[1] * perspective_coeffs[6]) / denom,
    ]

1667
1668
    theta1 = torch.tensor(
        [[inv_coeffs[0], inv_coeffs[1], inv_coeffs[2]], [inv_coeffs[3], inv_coeffs[4], inv_coeffs[5]]],
1669
1670
1671
1672
        dtype=dtype,
        device=device,
    )

1673
1674
1675
1676
    theta2 = torch.tensor(
        [[inv_coeffs[6], inv_coeffs[7], 1.0], [inv_coeffs[6], inv_coeffs[7], 1.0]], dtype=dtype, device=device
    )

1677
1678
1679
1680
    # 1) Let's transform bboxes into a tensor of 4 points (top-left, top-right, bottom-left, bottom-right corners).
    # Tensor of points has shape (N * 4, 3), where N is the number of bboxes
    # Single point structure is similar to
    # [(xmin, ymin, 1), (xmax, ymin, 1), (xmax, ymax, 1), (xmin, ymax, 1)]
1681
    points = bounding_boxes[:, [[0, 1], [2, 1], [2, 3], [0, 3]]].reshape(-1, 2)
1682
1683
1684
1685
1686
    points = torch.cat([points, torch.ones(points.shape[0], 1, device=points.device)], dim=-1)
    # 2) Now let's transform the points using perspective matrices
    #   x_out = (coeffs[0] * x + coeffs[1] * y + coeffs[2]) / (coeffs[6] * x + coeffs[7] * y + 1)
    #   y_out = (coeffs[3] * x + coeffs[4] * y + coeffs[5]) / (coeffs[6] * x + coeffs[7] * y + 1)

1687
1688
    numer_points = torch.matmul(points, theta1.T)
    denom_points = torch.matmul(points, theta2.T)
1689
    transformed_points = numer_points.div_(denom_points)
1690
1691
1692

    # 3) Reshape transformed points to [N boxes, 4 points, x/y coords]
    # and compute bounding box from 4 transformed points:
1693
    transformed_points = transformed_points.reshape(-1, 4, 2)
1694
1695
    out_bbox_mins, out_bbox_maxs = torch.aminmax(transformed_points, dim=1)

1696
1697
    out_bboxes = clamp_bounding_boxes(
        torch.cat([out_bbox_mins, out_bbox_maxs], dim=1).to(bounding_boxes.dtype),
1698
        format=datapoints.BoundingBoxFormat.XYXY,
Philip Meier's avatar
Philip Meier committed
1699
        canvas_size=canvas_size,
1700
    )
1701
1702
1703

    # out_bboxes should be of shape [N boxes, 4]

1704
    return convert_format_bounding_boxes(
1705
        out_bboxes, old_format=datapoints.BoundingBoxFormat.XYXY, new_format=format, inplace=True
1706
    ).reshape(original_shape)
1707
1708


1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
@_register_kernel_internal(perspective, datapoints.BoundingBoxes, datapoint_wrapper=False)
def _perspective_bounding_boxes_dispatch(
    inpt: datapoints.BoundingBoxes,
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
    coefficients: Optional[List[float]] = None,
    **kwargs,
) -> datapoints.BoundingBoxes:
    output = perspective_bounding_boxes(
        inpt.as_subclass(torch.Tensor),
        format=inpt.format,
        canvas_size=inpt.canvas_size,
        startpoints=startpoints,
        endpoints=endpoints,
        coefficients=coefficients,
    )
    return datapoints.BoundingBoxes.wrap_like(inpt, output)


1728
1729
def perspective_mask(
    mask: torch.Tensor,
1730
1731
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
Philip Meier's avatar
Philip Meier committed
1732
    fill: datapoints._FillTypeJIT = None,
1733
    coefficients: Optional[List[float]] = None,
1734
) -> torch.Tensor:
1735
1736
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
1737
1738
1739
1740
1741
        needs_squeeze = True
    else:
        needs_squeeze = False

    output = perspective_image_tensor(
1742
        mask, startpoints, endpoints, interpolation=InterpolationMode.NEAREST, fill=fill, coefficients=coefficients
1743
    )
1744

1745
1746
1747
1748
1749
    if needs_squeeze:
        output = output.squeeze(0)

    return output

1750

1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
@_register_kernel_internal(perspective, datapoints.Mask, datapoint_wrapper=False)
def _perspective_mask_dispatch(
    inpt: datapoints.Mask,
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
    fill: datapoints._FillTypeJIT = None,
    coefficients: Optional[List[float]] = None,
    **kwargs,
) -> datapoints.Mask:
    output = perspective_mask(
        inpt.as_subclass(torch.Tensor),
        startpoints=startpoints,
        endpoints=endpoints,
        fill=fill,
        coefficients=coefficients,
    )
    return datapoints.Mask.wrap_like(inpt, output)


@_register_kernel_internal(perspective, datapoints.Video)
1771
1772
def perspective_video(
    video: torch.Tensor,
1773
1774
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
1775
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
Philip Meier's avatar
Philip Meier committed
1776
    fill: datapoints._FillTypeJIT = None,
1777
    coefficients: Optional[List[float]] = None,
1778
) -> torch.Tensor:
1779
1780
1781
    return perspective_image_tensor(
        video, startpoints, endpoints, interpolation=interpolation, fill=fill, coefficients=coefficients
    )
1782
1783


1784
def elastic(
Philip Meier's avatar
Philip Meier committed
1785
    inpt: datapoints._InputTypeJIT,
1786
    displacement: torch.Tensor,
1787
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
Philip Meier's avatar
Philip Meier committed
1788
1789
    fill: datapoints._FillTypeJIT = None,
) -> datapoints._InputTypeJIT:
1790
    if not torch.jit.is_scripting():
1791
1792
1793
1794
1795
        _log_api_usage_once(elastic)

    if not isinstance(displacement, torch.Tensor):
        raise TypeError("Argument displacement should be a Tensor")

1796
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
1797
1798
1799
1800
        return elastic_image_tensor(inpt, displacement, interpolation=interpolation, fill=fill)
    elif isinstance(inpt, datapoints.Datapoint):
        kernel = _get_kernel(elastic, type(inpt))
        return kernel(inpt, displacement, interpolation=interpolation, fill=fill)
1801
    elif isinstance(inpt, PIL.Image.Image):
1802
        return elastic_image_pil(inpt, displacement, interpolation=interpolation, fill=fill)
1803
1804
    else:
        raise TypeError(
1805
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
1806
1807
            f"but got {type(inpt)} instead."
        )
1808
1809


1810
1811
1812
1813
elastic_transform = elastic


@_register_kernel_internal(elastic, datapoints.Image)
1814
def elastic_image_tensor(
1815
    image: torch.Tensor,
1816
    displacement: torch.Tensor,
1817
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
Philip Meier's avatar
Philip Meier committed
1818
    fill: datapoints._FillTypeJIT = None,
1819
) -> torch.Tensor:
1820
1821
    interpolation = _check_interpolation(interpolation)

1822
1823
1824
1825
    if image.numel() == 0:
        return image

    shape = image.shape
1826
    ndim = image.ndim
1827

1828
    device = image.device
1829
    dtype = image.dtype if torch.is_floating_point(image) else torch.float32
1830
1831
1832
1833
1834
1835
1836

    # Patch: elastic transform should support (cpu,f16) input
    is_cpu_half = device.type == "cpu" and dtype == torch.float16
    if is_cpu_half:
        image = image.to(torch.float32)
        dtype = torch.float32

1837
1838
1839
    # We are aware that if input image dtype is uint8 and displacement is float64 then
    # displacement will be casted to float32 and all computations will be done with float32
    # We can fix this later if needed
1840

1841
1842
1843
1844
    expected_shape = (1,) + shape[-2:] + (2,)
    if expected_shape != displacement.shape:
        raise ValueError(f"Argument displacement shape should be {expected_shape}, but given {displacement.shape}")

1845
    if ndim > 4:
1846
        image = image.reshape((-1,) + shape[-3:])
1847
        needs_unsquash = True
1848
1849
1850
    elif ndim == 3:
        image = image.unsqueeze(0)
        needs_unsquash = True
1851
1852
1853
    else:
        needs_unsquash = False

1854
1855
    if displacement.dtype != dtype or displacement.device != device:
        displacement = displacement.to(dtype=dtype, device=device)
1856

1857
1858
1859
    image_height, image_width = shape[-2:]
    grid = _create_identity_grid((image_height, image_width), device=device, dtype=dtype).add_(displacement)
    output = _apply_grid_transform(image, grid, interpolation.value, fill=fill)
1860
1861

    if needs_unsquash:
1862
        output = output.reshape(shape)
1863

1864
1865
1866
    if is_cpu_half:
        output = output.to(torch.float16)

1867
    return output
1868
1869


1870
@torch.jit.unused
1871
def elastic_image_pil(
1872
    image: PIL.Image.Image,
1873
    displacement: torch.Tensor,
1874
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
Philip Meier's avatar
Philip Meier committed
1875
    fill: datapoints._FillTypeJIT = None,
1876
) -> PIL.Image.Image:
1877
    t_img = pil_to_tensor(image)
1878
    output = elastic_image_tensor(t_img, displacement, interpolation=interpolation, fill=fill)
1879
    return to_pil_image(output, mode=image.mode)
1880
1881


1882
def _create_identity_grid(size: Tuple[int, int], device: torch.device, dtype: torch.dtype) -> torch.Tensor:
1883
    sy, sx = size
1884
1885
    base_grid = torch.empty(1, sy, sx, 2, device=device, dtype=dtype)
    x_grid = torch.linspace((-sx + 1) / sx, (sx - 1) / sx, sx, device=device, dtype=dtype)
1886
1887
    base_grid[..., 0].copy_(x_grid)

1888
    y_grid = torch.linspace((-sy + 1) / sy, (sy - 1) / sy, sy, device=device, dtype=dtype).unsqueeze_(-1)
1889
1890
1891
1892
1893
    base_grid[..., 1].copy_(y_grid)

    return base_grid


1894
1895
def elastic_bounding_boxes(
    bounding_boxes: torch.Tensor,
1896
    format: datapoints.BoundingBoxFormat,
Philip Meier's avatar
Philip Meier committed
1897
    canvas_size: Tuple[int, int],
1898
1899
    displacement: torch.Tensor,
) -> torch.Tensor:
1900
1901
    if bounding_boxes.numel() == 0:
        return bounding_boxes
1902

1903
    # TODO: add in docstring about approximation we are doing for grid inversion
1904
1905
    device = bounding_boxes.device
    dtype = bounding_boxes.dtype if torch.is_floating_point(bounding_boxes) else torch.float32
1906
1907
1908

    if displacement.dtype != dtype or displacement.device != device:
        displacement = displacement.to(dtype=dtype, device=device)
1909

1910
1911
1912
1913
    original_shape = bounding_boxes.shape
    # TODO: first cast to float if bbox is int64 before convert_format_bounding_boxes
    bounding_boxes = (
        convert_format_bounding_boxes(bounding_boxes, old_format=format, new_format=datapoints.BoundingBoxFormat.XYXY)
1914
    ).reshape(-1, 4)
1915

Philip Meier's avatar
Philip Meier committed
1916
    id_grid = _create_identity_grid(canvas_size, device=device, dtype=dtype)
1917
1918
    # We construct an approximation of inverse grid as inv_grid = id_grid - displacement
    # This is not an exact inverse of the grid
1919
    inv_grid = id_grid.sub_(displacement)
1920
1921

    # Get points from bboxes
1922
    points = bounding_boxes[:, [[0, 1], [2, 1], [2, 3], [0, 3]]].reshape(-1, 2)
1923
1924
1925
1926
1927
    if points.is_floating_point():
        points = points.ceil_()
    index_xy = points.to(dtype=torch.long)
    index_x, index_y = index_xy[:, 0], index_xy[:, 1]

1928
    # Transform points:
Philip Meier's avatar
Philip Meier committed
1929
    t_size = torch.tensor(canvas_size[::-1], device=displacement.device, dtype=displacement.dtype)
1930
    transformed_points = inv_grid[0, index_y, index_x, :].add_(1).mul_(0.5 * t_size).sub_(0.5)
1931

1932
    transformed_points = transformed_points.reshape(-1, 4, 2)
1933
    out_bbox_mins, out_bbox_maxs = torch.aminmax(transformed_points, dim=1)
1934
1935
    out_bboxes = clamp_bounding_boxes(
        torch.cat([out_bbox_mins, out_bbox_maxs], dim=1).to(bounding_boxes.dtype),
1936
        format=datapoints.BoundingBoxFormat.XYXY,
Philip Meier's avatar
Philip Meier committed
1937
        canvas_size=canvas_size,
1938
    )
1939

1940
    return convert_format_bounding_boxes(
1941
        out_bboxes, old_format=datapoints.BoundingBoxFormat.XYXY, new_format=format, inplace=True
1942
    ).reshape(original_shape)
1943
1944


1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
@_register_kernel_internal(elastic, datapoints.BoundingBoxes, datapoint_wrapper=False)
def _elastic_bounding_boxes_dispatch(
    inpt: datapoints.BoundingBoxes, displacement: torch.Tensor, **kwargs
) -> datapoints.BoundingBoxes:
    output = elastic_bounding_boxes(
        inpt.as_subclass(torch.Tensor), format=inpt.format, canvas_size=inpt.canvas_size, displacement=displacement
    )
    return datapoints.BoundingBoxes.wrap_like(inpt, output)


1955
1956
1957
def elastic_mask(
    mask: torch.Tensor,
    displacement: torch.Tensor,
Philip Meier's avatar
Philip Meier committed
1958
    fill: datapoints._FillTypeJIT = None,
1959
) -> torch.Tensor:
1960
1961
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
1962
1963
1964
1965
        needs_squeeze = True
    else:
        needs_squeeze = False

1966
    output = elastic_image_tensor(mask, displacement=displacement, interpolation=InterpolationMode.NEAREST, fill=fill)
1967
1968
1969
1970
1971

    if needs_squeeze:
        output = output.squeeze(0)

    return output
1972
1973


1974
1975
1976
1977
1978
1979
1980
1981
1982
@_register_kernel_internal(elastic, datapoints.Mask, datapoint_wrapper=False)
def _elastic_mask_dispatch(
    inpt: datapoints.Mask, displacement: torch.Tensor, fill: datapoints._FillTypeJIT = None, **kwargs
) -> datapoints.Mask:
    output = elastic_mask(inpt.as_subclass(torch.Tensor), displacement=displacement, fill=fill)
    return datapoints.Mask.wrap_like(inpt, output)


@_register_kernel_internal(elastic, datapoints.Video)
1983
1984
1985
def elastic_video(
    video: torch.Tensor,
    displacement: torch.Tensor,
1986
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
Philip Meier's avatar
Philip Meier committed
1987
    fill: datapoints._FillTypeJIT = None,
1988
) -> torch.Tensor:
1989
    return elastic_image_tensor(video, displacement, interpolation=interpolation, fill=fill)
1990
1991


1992
def center_crop(inpt: datapoints._InputTypeJIT, output_size: List[int]) -> datapoints._InputTypeJIT:
1993
    if not torch.jit.is_scripting():
1994
        _log_api_usage_once(center_crop)
1995

1996
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
1997
1998
1999
2000
        return center_crop_image_tensor(inpt, output_size)
    elif isinstance(inpt, datapoints.Datapoint):
        kernel = _get_kernel(center_crop, type(inpt))
        return kernel(inpt, output_size)
2001
    elif isinstance(inpt, PIL.Image.Image):
2002
        return center_crop_image_pil(inpt, output_size)
2003
2004
    else:
        raise TypeError(
2005
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
2006
2007
            f"but got {type(inpt)} instead."
        )
2008
2009


2010
2011
def _center_crop_parse_output_size(output_size: List[int]) -> List[int]:
    if isinstance(output_size, numbers.Number):
2012
2013
        s = int(output_size)
        return [s, s]
2014
    elif isinstance(output_size, (tuple, list)) and len(output_size) == 1:
2015
        return [output_size[0], output_size[0]]
2016
2017
    else:
        return list(output_size)
2018
2019
2020
2021
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036


def _center_crop_compute_padding(crop_height: int, crop_width: int, image_height: int, image_width: int) -> List[int]:
    return [
        (crop_width - image_width) // 2 if crop_width > image_width else 0,
        (crop_height - image_height) // 2 if crop_height > image_height else 0,
        (crop_width - image_width + 1) // 2 if crop_width > image_width else 0,
        (crop_height - image_height + 1) // 2 if crop_height > image_height else 0,
    ]


def _center_crop_compute_crop_anchor(
    crop_height: int, crop_width: int, image_height: int, image_width: int
) -> Tuple[int, int]:
    crop_top = int(round((image_height - crop_height) / 2.0))
    crop_left = int(round((image_width - crop_width) / 2.0))
    return crop_top, crop_left


2037
@_register_kernel_internal(center_crop, datapoints.Image)
2038
def center_crop_image_tensor(image: torch.Tensor, output_size: List[int]) -> torch.Tensor:
2039
    crop_height, crop_width = _center_crop_parse_output_size(output_size)
2040
2041
2042
2043
    shape = image.shape
    if image.numel() == 0:
        return image.reshape(shape[:-2] + (crop_height, crop_width))
    image_height, image_width = shape[-2:]
2044
2045
2046

    if crop_height > image_height or crop_width > image_width:
        padding_ltrb = _center_crop_compute_padding(crop_height, crop_width, image_height, image_width)
2047
        image = torch_pad(image, _parse_pad_padding(padding_ltrb), value=0.0)
2048

2049
        image_height, image_width = image.shape[-2:]
2050
        if crop_width == image_width and crop_height == image_height:
2051
            return image
2052
2053

    crop_top, crop_left = _center_crop_compute_crop_anchor(crop_height, crop_width, image_height, image_width)
2054
    return image[..., crop_top : (crop_top + crop_height), crop_left : (crop_left + crop_width)]
2055
2056


2057
@torch.jit.unused
2058
def center_crop_image_pil(image: PIL.Image.Image, output_size: List[int]) -> PIL.Image.Image:
2059
    crop_height, crop_width = _center_crop_parse_output_size(output_size)
Philip Meier's avatar
Philip Meier committed
2060
    image_height, image_width = get_size_image_pil(image)
2061
2062
2063

    if crop_height > image_height or crop_width > image_width:
        padding_ltrb = _center_crop_compute_padding(crop_height, crop_width, image_height, image_width)
2064
        image = pad_image_pil(image, padding_ltrb, fill=0)
2065

Philip Meier's avatar
Philip Meier committed
2066
        image_height, image_width = get_size_image_pil(image)
2067
        if crop_width == image_width and crop_height == image_height:
2068
            return image
2069
2070

    crop_top, crop_left = _center_crop_compute_crop_anchor(crop_height, crop_width, image_height, image_width)
2071
    return crop_image_pil(image, crop_top, crop_left, crop_height, crop_width)
2072
2073


2074
2075
def center_crop_bounding_boxes(
    bounding_boxes: torch.Tensor,
2076
    format: datapoints.BoundingBoxFormat,
Philip Meier's avatar
Philip Meier committed
2077
    canvas_size: Tuple[int, int],
2078
    output_size: List[int],
2079
) -> Tuple[torch.Tensor, Tuple[int, int]]:
2080
    crop_height, crop_width = _center_crop_parse_output_size(output_size)
Philip Meier's avatar
Philip Meier committed
2081
    crop_top, crop_left = _center_crop_compute_crop_anchor(crop_height, crop_width, *canvas_size)
2082
2083
2084
    return crop_bounding_boxes(
        bounding_boxes, format, top=crop_top, left=crop_left, height=crop_height, width=crop_width
    )
2085
2086


2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
@_register_kernel_internal(center_crop, datapoints.BoundingBoxes, datapoint_wrapper=False)
def _center_crop_bounding_boxes_dispatch(
    inpt: datapoints.BoundingBoxes, output_size: List[int]
) -> datapoints.BoundingBoxes:
    output, canvas_size = center_crop_bounding_boxes(
        inpt.as_subclass(torch.Tensor), format=inpt.format, canvas_size=inpt.canvas_size, output_size=output_size
    )
    return datapoints.BoundingBoxes.wrap_like(inpt, output, canvas_size=canvas_size)


@_register_kernel_internal(center_crop, datapoints.Mask)
2098
2099
2100
def center_crop_mask(mask: torch.Tensor, output_size: List[int]) -> torch.Tensor:
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
2101
2102
2103
2104
        needs_squeeze = True
    else:
        needs_squeeze = False

2105
    output = center_crop_image_tensor(image=mask, output_size=output_size)
2106
2107
2108
2109
2110

    if needs_squeeze:
        output = output.squeeze(0)

    return output
2111
2112


2113
@_register_kernel_internal(center_crop, datapoints.Video)
2114
2115
2116
2117
def center_crop_video(video: torch.Tensor, output_size: List[int]) -> torch.Tensor:
    return center_crop_image_tensor(video, output_size)


2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
def resized_crop(
    inpt: datapoints._InputTypeJIT,
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
    antialias: Optional[Union[str, bool]] = "warn",
) -> datapoints._InputTypeJIT:
2128
    if not torch.jit.is_scripting():
2129
        _log_api_usage_once(resized_crop)
2130

2131
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
2132
2133
2134
2135
2136
2137
        return resized_crop_image_tensor(
            inpt, top, left, height, width, antialias=antialias, size=size, interpolation=interpolation
        )
    elif isinstance(inpt, datapoints.Datapoint):
        kernel = _get_kernel(resized_crop, type(inpt))
        return kernel(inpt, top, left, height, width, antialias=antialias, size=size, interpolation=interpolation)
2138
    elif isinstance(inpt, PIL.Image.Image):
2139
        return resized_crop_image_pil(inpt, top, left, height, width, size=size, interpolation=interpolation)
2140
2141
    else:
        raise TypeError(
2142
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
2143
2144
            f"but got {type(inpt)} instead."
        )
2145
2146


2147
@_register_kernel_internal(resized_crop, datapoints.Image)
2148
def resized_crop_image_tensor(
2149
    image: torch.Tensor,
2150
2151
2152
2153
2154
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
2155
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
2156
    antialias: Optional[Union[str, bool]] = "warn",
2157
) -> torch.Tensor:
2158
2159
    image = crop_image_tensor(image, top, left, height, width)
    return resize_image_tensor(image, size, interpolation=interpolation, antialias=antialias)
2160
2161


2162
@torch.jit.unused
2163
def resized_crop_image_pil(
2164
    image: PIL.Image.Image,
2165
2166
2167
2168
2169
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
2170
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
2171
) -> PIL.Image.Image:
2172
2173
    image = crop_image_pil(image, top, left, height, width)
    return resize_image_pil(image, size, interpolation=interpolation)
2174
2175


2176
2177
def resized_crop_bounding_boxes(
    bounding_boxes: torch.Tensor,
2178
    format: datapoints.BoundingBoxFormat,
2179
2180
2181
2182
2183
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
2184
) -> Tuple[torch.Tensor, Tuple[int, int]]:
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
    bounding_boxes, canvas_size = crop_bounding_boxes(bounding_boxes, format, top, left, height, width)
    return resize_bounding_boxes(bounding_boxes, canvas_size=canvas_size, size=size)


@_register_kernel_internal(resized_crop, datapoints.BoundingBoxes, datapoint_wrapper=False)
def _resized_crop_bounding_boxes_dispatch(
    inpt: datapoints.BoundingBoxes, top: int, left: int, height: int, width: int, size: List[int], **kwargs
) -> datapoints.BoundingBoxes:
    output, canvas_size = resized_crop_bounding_boxes(
        inpt.as_subclass(torch.Tensor), format=inpt.format, top=top, left=left, height=height, width=width, size=size
    )
    return datapoints.BoundingBoxes.wrap_like(inpt, output, canvas_size=canvas_size)
2197
2198


2199
def resized_crop_mask(
2200
2201
2202
2203
2204
2205
2206
    mask: torch.Tensor,
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
) -> torch.Tensor:
2207
2208
    mask = crop_mask(mask, top, left, height, width)
    return resize_mask(mask, size)
2209
2210


2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
@_register_kernel_internal(resized_crop, datapoints.Mask, datapoint_wrapper=False)
def _resized_crop_mask_dispatch(
    inpt: datapoints.Mask, top: int, left: int, height: int, width: int, size: List[int], **kwargs
) -> datapoints.Mask:
    output = resized_crop_mask(
        inpt.as_subclass(torch.Tensor), top=top, left=left, height=height, width=width, size=size
    )
    return datapoints.Mask.wrap_like(inpt, output)


@_register_kernel_internal(resized_crop, datapoints.Video)
2222
2223
2224
2225
2226
2227
2228
def resized_crop_video(
    video: torch.Tensor,
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
2229
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
2230
    antialias: Optional[Union[str, bool]] = "warn",
2231
2232
2233
2234
2235
2236
) -> torch.Tensor:
    return resized_crop_image_tensor(
        video, top, left, height, width, antialias=antialias, size=size, interpolation=interpolation
    )


2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
@_register_explicit_noop(datapoints.BoundingBoxes, datapoints.Mask, warn_passthrough=True)
def five_crop(
    inpt: datapoints._InputTypeJIT, size: List[int]
) -> Tuple[
    datapoints._InputTypeJIT,
    datapoints._InputTypeJIT,
    datapoints._InputTypeJIT,
    datapoints._InputTypeJIT,
    datapoints._InputTypeJIT,
]:
2247
    if not torch.jit.is_scripting():
2248
        _log_api_usage_once(five_crop)
2249

2250
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
2251
2252
2253
2254
        return five_crop_image_tensor(inpt, size)
    elif isinstance(inpt, datapoints.Datapoint):
        kernel = _get_kernel(five_crop, type(inpt))
        return kernel(inpt, size)
2255
    elif isinstance(inpt, PIL.Image.Image):
2256
        return five_crop_image_pil(inpt, size)
2257
2258
    else:
        raise TypeError(
2259
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
2260
2261
            f"but got {type(inpt)} instead."
        )
2262
2263


2264
2265
def _parse_five_crop_size(size: List[int]) -> List[int]:
    if isinstance(size, numbers.Number):
2266
2267
        s = int(size)
        size = [s, s]
2268
    elif isinstance(size, (tuple, list)) and len(size) == 1:
2269
2270
        s = size[0]
        size = [s, s]
2271
2272
2273
2274
2275
2276
2277

    if len(size) != 2:
        raise ValueError("Please provide only two dimensions (h, w) for size.")

    return size


2278
@_register_five_ten_crop_kernel(five_crop, datapoints.Image)
2279
def five_crop_image_tensor(
2280
    image: torch.Tensor, size: List[int]
2281
2282
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    crop_height, crop_width = _parse_five_crop_size(size)
2283
    image_height, image_width = image.shape[-2:]
2284
2285

    if crop_width > image_width or crop_height > image_height:
2286
        raise ValueError(f"Requested crop size {size} is bigger than input size {(image_height, image_width)}")
2287

2288
2289
2290
2291
2292
    tl = crop_image_tensor(image, 0, 0, crop_height, crop_width)
    tr = crop_image_tensor(image, 0, image_width - crop_width, crop_height, crop_width)
    bl = crop_image_tensor(image, image_height - crop_height, 0, crop_height, crop_width)
    br = crop_image_tensor(image, image_height - crop_height, image_width - crop_width, crop_height, crop_width)
    center = center_crop_image_tensor(image, [crop_height, crop_width])
2293
2294
2295
2296

    return tl, tr, bl, br, center


2297
@torch.jit.unused
2298
def five_crop_image_pil(
2299
    image: PIL.Image.Image, size: List[int]
2300
2301
) -> Tuple[PIL.Image.Image, PIL.Image.Image, PIL.Image.Image, PIL.Image.Image, PIL.Image.Image]:
    crop_height, crop_width = _parse_five_crop_size(size)
Philip Meier's avatar
Philip Meier committed
2302
    image_height, image_width = get_size_image_pil(image)
2303
2304

    if crop_width > image_width or crop_height > image_height:
2305
        raise ValueError(f"Requested crop size {size} is bigger than input size {(image_height, image_width)}")
2306

2307
2308
2309
2310
2311
    tl = crop_image_pil(image, 0, 0, crop_height, crop_width)
    tr = crop_image_pil(image, 0, image_width - crop_width, crop_height, crop_width)
    bl = crop_image_pil(image, image_height - crop_height, 0, crop_height, crop_width)
    br = crop_image_pil(image, image_height - crop_height, image_width - crop_width, crop_height, crop_width)
    center = center_crop_image_pil(image, [crop_height, crop_width])
2312
2313
2314
2315

    return tl, tr, bl, br, center


2316
@_register_five_ten_crop_kernel(five_crop, datapoints.Video)
2317
2318
2319
2320
2321
2322
def five_crop_video(
    video: torch.Tensor, size: List[int]
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    return five_crop_image_tensor(video, size)


2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
@_register_explicit_noop(datapoints.BoundingBoxes, datapoints.Mask, warn_passthrough=True)
def ten_crop(
    inpt: Union[datapoints._ImageTypeJIT, datapoints._VideoTypeJIT], size: List[int], vertical_flip: bool = False
) -> Tuple[
    datapoints._InputTypeJIT,
    datapoints._InputTypeJIT,
    datapoints._InputTypeJIT,
    datapoints._InputTypeJIT,
    datapoints._InputTypeJIT,
    datapoints._InputTypeJIT,
    datapoints._InputTypeJIT,
    datapoints._InputTypeJIT,
    datapoints._InputTypeJIT,
    datapoints._InputTypeJIT,
]:
2338
    if not torch.jit.is_scripting():
2339
        _log_api_usage_once(ten_crop)
2340

2341
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
2342
2343
2344
2345
        return ten_crop_image_tensor(inpt, size, vertical_flip=vertical_flip)
    elif isinstance(inpt, datapoints.Datapoint):
        kernel = _get_kernel(ten_crop, type(inpt))
        return kernel(inpt, size, vertical_flip=vertical_flip)
2346
    elif isinstance(inpt, PIL.Image.Image):
2347
        return ten_crop_image_pil(inpt, size, vertical_flip=vertical_flip)
2348
2349
    else:
        raise TypeError(
2350
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
2351
2352
            f"but got {type(inpt)} instead."
        )
2353
2354


2355
@_register_five_ten_crop_kernel(ten_crop, datapoints.Image)
Philip Meier's avatar
Philip Meier committed
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
def ten_crop_image_tensor(
    image: torch.Tensor, size: List[int], vertical_flip: bool = False
) -> Tuple[
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
]:
    non_flipped = five_crop_image_tensor(image, size)
2371
2372

    if vertical_flip:
2373
        image = vertical_flip_image_tensor(image)
2374
    else:
2375
        image = horizontal_flip_image_tensor(image)
2376

Philip Meier's avatar
Philip Meier committed
2377
    flipped = five_crop_image_tensor(image, size)
2378

Philip Meier's avatar
Philip Meier committed
2379
    return non_flipped + flipped
2380
2381


2382
@torch.jit.unused
Philip Meier's avatar
Philip Meier committed
2383
2384
2385
2386
2387
2388
2389
2390
2391
2392
2393
2394
2395
2396
2397
def ten_crop_image_pil(
    image: PIL.Image.Image, size: List[int], vertical_flip: bool = False
) -> Tuple[
    PIL.Image.Image,
    PIL.Image.Image,
    PIL.Image.Image,
    PIL.Image.Image,
    PIL.Image.Image,
    PIL.Image.Image,
    PIL.Image.Image,
    PIL.Image.Image,
    PIL.Image.Image,
    PIL.Image.Image,
]:
    non_flipped = five_crop_image_pil(image, size)
2398
2399

    if vertical_flip:
2400
        image = vertical_flip_image_pil(image)
2401
    else:
2402
        image = horizontal_flip_image_pil(image)
2403

Philip Meier's avatar
Philip Meier committed
2404
2405
2406
2407
2408
    flipped = five_crop_image_pil(image, size)

    return non_flipped + flipped


2409
@_register_five_ten_crop_kernel(ten_crop, datapoints.Video)
Philip Meier's avatar
Philip Meier committed
2410
2411
2412
2413
2414
2415
2416
2417
2418
2419
2420
2421
2422
2423
def ten_crop_video(
    video: torch.Tensor, size: List[int], vertical_flip: bool = False
) -> Tuple[
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
]:
2424
    return ten_crop_image_tensor(video, size, vertical_flip=vertical_flip)