_geometry.py 85.5 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 tv_tensors
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
    _compute_resized_output_size as __compute_resized_output_size,
15
    _get_perspective_coeffs,
16
    _interpolation_modes_from_int,
17
    InterpolationMode,
18
    pil_modes_mapping,
19
20
    pil_to_tensor,
    to_pil_image,
21
)
22

23
24
from torchvision.utils import _log_api_usage_once

Nicolas Hug's avatar
Nicolas Hug committed
25
from ._meta import _get_size_image_pil, clamp_bounding_boxes, convert_bounding_box_format
26

27
from ._utils import _FillTypeJIT, _get_kernel, _register_five_ten_crop_kernel_internal, _register_kernel_internal
28

29

30
31
32
33
34
35
36
37
38
39
40
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


41
def horizontal_flip(inpt: torch.Tensor) -> torch.Tensor:
42
    """See :class:`~torchvision.transforms.v2.RandomHorizontalFlip` for details."""
43
    if torch.jit.is_scripting():
44
        return horizontal_flip_image(inpt)
45
46
47
48
49

    _log_api_usage_once(horizontal_flip)

    kernel = _get_kernel(horizontal_flip, type(inpt))
    return kernel(inpt)
50
51


52
@_register_kernel_internal(horizontal_flip, torch.Tensor)
53
@_register_kernel_internal(horizontal_flip, tv_tensors.Image)
54
def horizontal_flip_image(image: torch.Tensor) -> torch.Tensor:
55
56
57
    return image.flip(-1)


58
@_register_kernel_internal(horizontal_flip, PIL.Image.Image)
59
def _horizontal_flip_image_pil(image: PIL.Image.Image) -> PIL.Image.Image:
60
    return _FP.hflip(image)
61
62


63
@_register_kernel_internal(horizontal_flip, tv_tensors.Mask)
64
def horizontal_flip_mask(mask: torch.Tensor) -> torch.Tensor:
65
    return horizontal_flip_image(mask)
66
67


68
def horizontal_flip_bounding_boxes(
69
    bounding_boxes: torch.Tensor, format: tv_tensors.BoundingBoxFormat, canvas_size: Tuple[int, int]
70
) -> torch.Tensor:
71
    shape = bounding_boxes.shape
72

73
    bounding_boxes = bounding_boxes.clone().reshape(-1, 4)
74

75
    if format == tv_tensors.BoundingBoxFormat.XYXY:
Philip Meier's avatar
Philip Meier committed
76
        bounding_boxes[:, [2, 0]] = bounding_boxes[:, [0, 2]].sub_(canvas_size[1]).neg_()
77
    elif format == tv_tensors.BoundingBoxFormat.XYWH:
Philip Meier's avatar
Philip Meier committed
78
        bounding_boxes[:, 0].add_(bounding_boxes[:, 2]).sub_(canvas_size[1]).neg_()
79
    else:  # format == tv_tensors.BoundingBoxFormat.CXCYWH:
Philip Meier's avatar
Philip Meier committed
80
        bounding_boxes[:, 0].sub_(canvas_size[1]).neg_()
81

82
    return bounding_boxes.reshape(shape)
83
84


85
86
@_register_kernel_internal(horizontal_flip, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False)
def _horizontal_flip_bounding_boxes_dispatch(inpt: tv_tensors.BoundingBoxes) -> tv_tensors.BoundingBoxes:
87
88
89
    output = horizontal_flip_bounding_boxes(
        inpt.as_subclass(torch.Tensor), format=inpt.format, canvas_size=inpt.canvas_size
    )
90
    return tv_tensors.wrap(output, like=inpt)
91
92


93
@_register_kernel_internal(horizontal_flip, tv_tensors.Video)
94
def horizontal_flip_video(video: torch.Tensor) -> torch.Tensor:
95
    return horizontal_flip_image(video)
96
97


98
def vertical_flip(inpt: torch.Tensor) -> torch.Tensor:
99
    """See :class:`~torchvision.transforms.v2.RandomVerticalFlip` for details."""
100
    if torch.jit.is_scripting():
101
        return vertical_flip_image(inpt)
102
103
104
105
106

    _log_api_usage_once(vertical_flip)

    kernel = _get_kernel(vertical_flip, type(inpt))
    return kernel(inpt)
107
108


109
@_register_kernel_internal(vertical_flip, torch.Tensor)
110
@_register_kernel_internal(vertical_flip, tv_tensors.Image)
111
def vertical_flip_image(image: torch.Tensor) -> torch.Tensor:
112
113
114
    return image.flip(-2)


115
@_register_kernel_internal(vertical_flip, PIL.Image.Image)
Nicolas Hug's avatar
Nicolas Hug committed
116
def _vertical_flip_image_pil(image: PIL.Image.Image) -> PIL.Image.Image:
Philip Meier's avatar
Philip Meier committed
117
    return _FP.vflip(image)
118
119


120
@_register_kernel_internal(vertical_flip, tv_tensors.Mask)
121
def vertical_flip_mask(mask: torch.Tensor) -> torch.Tensor:
122
    return vertical_flip_image(mask)
123
124


125
def vertical_flip_bounding_boxes(
126
    bounding_boxes: torch.Tensor, format: tv_tensors.BoundingBoxFormat, canvas_size: Tuple[int, int]
127
) -> torch.Tensor:
128
    shape = bounding_boxes.shape
129

130
    bounding_boxes = bounding_boxes.clone().reshape(-1, 4)
131

132
    if format == tv_tensors.BoundingBoxFormat.XYXY:
Philip Meier's avatar
Philip Meier committed
133
        bounding_boxes[:, [1, 3]] = bounding_boxes[:, [3, 1]].sub_(canvas_size[0]).neg_()
134
    elif format == tv_tensors.BoundingBoxFormat.XYWH:
Philip Meier's avatar
Philip Meier committed
135
        bounding_boxes[:, 1].add_(bounding_boxes[:, 3]).sub_(canvas_size[0]).neg_()
136
    else:  # format == tv_tensors.BoundingBoxFormat.CXCYWH:
Philip Meier's avatar
Philip Meier committed
137
        bounding_boxes[:, 1].sub_(canvas_size[0]).neg_()
138

139
    return bounding_boxes.reshape(shape)
140
141


142
143
@_register_kernel_internal(vertical_flip, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False)
def _vertical_flip_bounding_boxes_dispatch(inpt: tv_tensors.BoundingBoxes) -> tv_tensors.BoundingBoxes:
144
145
146
    output = vertical_flip_bounding_boxes(
        inpt.as_subclass(torch.Tensor), format=inpt.format, canvas_size=inpt.canvas_size
    )
147
    return tv_tensors.wrap(output, like=inpt)
148

149

150
@_register_kernel_internal(vertical_flip, tv_tensors.Video)
151
def vertical_flip_video(video: torch.Tensor) -> torch.Tensor:
152
    return vertical_flip_image(video)
153
154


155
156
157
158
159
160
# 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


161
def _compute_resized_output_size(
162
    canvas_size: Tuple[int, int], size: Optional[List[int]], max_size: Optional[int] = None
163
164
165
) -> List[int]:
    if isinstance(size, int):
        size = [size]
166
    elif max_size is not None and size is not None and len(size) != 1:
167
        raise ValueError(
168
            "max_size should only be passed if size is None or specifies the length of the smaller edge, "
169
170
            "i.e. size should be an int or a sequence of length 1 in torchscript mode."
        )
171
    return __compute_resized_output_size(canvas_size, size=size, max_size=max_size, allow_size_none=True)
172
173


174
def resize(
175
    inpt: torch.Tensor,
176
    size: Optional[List[int]],
177
178
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
    max_size: Optional[int] = None,
179
    antialias: Optional[bool] = True,
180
) -> torch.Tensor:
181
    """See :class:`~torchvision.transforms.v2.Resize` for details."""
182
    if torch.jit.is_scripting():
183
        return resize_image(inpt, size=size, interpolation=interpolation, max_size=max_size, antialias=antialias)
184
185
186
187
188

    _log_api_usage_once(resize)

    kernel = _get_kernel(resize, type(inpt))
    return kernel(inpt, size=size, interpolation=interpolation, max_size=max_size, antialias=antialias)
189
190


191
192
193
194
195
196
# This is an internal helper method for resize_image. We should put it here instead of keeping it
# inside resize_image due to torchscript.
# uint8 dtype support for bilinear and bicubic is limited to cpu and
# according to our benchmarks on eager, non-AVX CPUs should still prefer u8->f32->interpolate->u8 path for bilinear
def _do_native_uint8_resize_on_cpu(interpolation: InterpolationMode) -> bool:
    if interpolation == InterpolationMode.BILINEAR:
197
        if torch.compiler.is_compiling():
198
199
200
201
202
203
204
            return True
        else:
            return "AVX2" in torch.backends.cpu.get_cpu_capability()

    return interpolation == InterpolationMode.BICUBIC


205
@_register_kernel_internal(resize, torch.Tensor)
206
@_register_kernel_internal(resize, tv_tensors.Image)
207
def resize_image(
208
    image: torch.Tensor,
209
    size: Optional[List[int]],
210
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
211
    max_size: Optional[int] = None,
212
    antialias: Optional[bool] = True,
213
) -> torch.Tensor:
214
    interpolation = _check_interpolation(interpolation)
215
    antialias = False if antialias is None else antialias
216
217
218
    align_corners: Optional[bool] = None
    if interpolation == InterpolationMode.BILINEAR or interpolation == InterpolationMode.BICUBIC:
        align_corners = False
219
    else:
220
        # The default of antialias is True from 0.17, so we don't warn or
221
222
        # error if other interpolation modes are used. This is documented.
        antialias = False
223

224
    shape = image.shape
225
    numel = image.numel()
226
    num_channels, old_height, old_width = shape[-3:]
vfdev's avatar
vfdev committed
227
    new_height, new_width = _compute_resized_output_size((old_height, old_width), size=size, max_size=max_size)
228

229
230
    if (new_height, new_width) == (old_height, old_width):
        return image
231
    elif numel > 0:
232
        dtype = image.dtype
233
234
235
236
        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)
237
        elif image.device.type == "cpu":
238
            if _do_native_uint8_resize_on_cpu(interpolation):
239
                acceptable_dtypes.append(torch.uint8)
240

241
        image = image.reshape(-1, num_channels, old_height, old_width)
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
        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
257
258
259
260
        if need_cast:
            image = image.to(dtype=torch.float32)

        image = interpolate(
261
262
            image,
            size=[new_height, new_width],
263
264
            mode=interpolation.value,
            align_corners=align_corners,
265
266
            antialias=antialias,
        )
267

268
269
        if need_cast:
            if interpolation == InterpolationMode.BICUBIC and dtype == torch.uint8:
270
                # This path is hit on non-AVX archs, or on GPU.
271
                image = image.clamp_(min=0, max=255)
272
273
274
            if dtype in (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64):
                image = image.round_()
            image = image.to(dtype=dtype)
275

276
    return image.reshape(shape[:-3] + (num_channels, new_height, new_width))
277
278


279
def _resize_image_pil(
280
    image: PIL.Image.Image,
281
    size: Union[Sequence[int], int],
282
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
283
284
    max_size: Optional[int] = None,
) -> PIL.Image.Image:
285
286
287
288
289
290
291
    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,
    )

292
    interpolation = _check_interpolation(interpolation)
293
294
295
296
297

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

    return image.resize((new_width, new_height), resample=pil_modes_mapping[interpolation])
298
299


300
@_register_kernel_internal(resize, PIL.Image.Image)
301
def __resize_image_pil_dispatch(
302
303
304
305
    image: PIL.Image.Image,
    size: Union[Sequence[int], int],
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
    max_size: Optional[int] = None,
306
    antialias: Optional[bool] = True,
307
308
309
) -> PIL.Image.Image:
    if antialias is False:
        warnings.warn("Anti-alias option is always applied for PIL Image input. Argument antialias is ignored.")
310
    return _resize_image_pil(image, size=size, interpolation=interpolation, max_size=max_size)
311
312


313
def resize_mask(mask: torch.Tensor, size: Optional[List[int]], max_size: Optional[int] = None) -> torch.Tensor:
314
315
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
316
317
318
319
        needs_squeeze = True
    else:
        needs_squeeze = False

320
    output = resize_image(mask, size=size, interpolation=InterpolationMode.NEAREST, max_size=max_size)
321
322
323
324
325

    if needs_squeeze:
        output = output.squeeze(0)

    return output
326
327


328
@_register_kernel_internal(resize, tv_tensors.Mask, tv_tensor_wrapper=False)
329
def _resize_mask_dispatch(
330
331
    inpt: tv_tensors.Mask, size: List[int], max_size: Optional[int] = None, **kwargs: Any
) -> tv_tensors.Mask:
332
    output = resize_mask(inpt.as_subclass(torch.Tensor), size, max_size=max_size)
333
    return tv_tensors.wrap(output, like=inpt)
334
335


336
def resize_bounding_boxes(
337
338
339
340
    bounding_boxes: torch.Tensor,
    canvas_size: Tuple[int, int],
    size: Optional[List[int]],
    max_size: Optional[int] = None,
341
) -> Tuple[torch.Tensor, Tuple[int, int]]:
Philip Meier's avatar
Philip Meier committed
342
343
    old_height, old_width = canvas_size
    new_height, new_width = _compute_resized_output_size(canvas_size, size=size, max_size=max_size)
344
345

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

348
349
    w_ratio = new_width / old_width
    h_ratio = new_height / old_height
350
    ratios = torch.tensor([w_ratio, h_ratio, w_ratio, h_ratio], device=bounding_boxes.device)
351
    return (
352
        bounding_boxes.mul(ratios).to(bounding_boxes.dtype),
353
354
        (new_height, new_width),
    )
355
356


357
@_register_kernel_internal(resize, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False)
358
def _resize_bounding_boxes_dispatch(
359
    inpt: tv_tensors.BoundingBoxes, size: Optional[List[int]], max_size: Optional[int] = None, **kwargs: Any
360
) -> tv_tensors.BoundingBoxes:
361
362
363
    output, canvas_size = resize_bounding_boxes(
        inpt.as_subclass(torch.Tensor), inpt.canvas_size, size, max_size=max_size
    )
364
    return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size)
365
366


367
@_register_kernel_internal(resize, tv_tensors.Video)
368
369
def resize_video(
    video: torch.Tensor,
370
    size: Optional[List[int]],
371
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
372
    max_size: Optional[int] = None,
373
    antialias: Optional[bool] = True,
374
) -> torch.Tensor:
375
    return resize_image(video, size=size, interpolation=interpolation, max_size=max_size, antialias=antialias)
376
377


378
def affine(
379
    inpt: torch.Tensor,
380
381
382
383
384
    angle: Union[int, float],
    translate: List[float],
    scale: float,
    shear: List[float],
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
385
    fill: _FillTypeJIT = None,
386
    center: Optional[List[float]] = None,
387
) -> torch.Tensor:
388
    """See :class:`~torchvision.transforms.v2.RandomAffine` for details."""
389
    if torch.jit.is_scripting():
390
        return affine_image(
391
            inpt,
392
            angle=angle,
393
394
395
396
397
398
399
            translate=translate,
            scale=scale,
            shear=shear,
            interpolation=interpolation,
            fill=fill,
            center=center,
        )
400
401
402
403
404
405
406
407
408
409
410
411
412
413

    _log_api_usage_once(affine)

    kernel = _get_kernel(affine, type(inpt))
    return kernel(
        inpt,
        angle=angle,
        translate=translate,
        scale=scale,
        shear=shear,
        interpolation=interpolation,
        fill=fill,
        center=center,
    )
414
415


416
def _affine_parse_args(
417
    angle: Union[int, float],
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
    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}")

460
461
462
463
464
    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]
465
466
467
468

    return angle, translate, shear, center


469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
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
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]:
531
    if torch.compiler.is_compiling() and not torch.jit.is_scripting():
532
533
534
535
536
537
        return _compute_affine_output_size_python(matrix, w, h)
    else:
        return _compute_affine_output_size_tensor(matrix, w, h)


def _compute_affine_output_size_tensor(matrix: List[float], w: int, h: int) -> Tuple[int, int]:
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
    # 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


572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
def _compute_affine_output_size_python(matrix: List[float], w: int, h: int) -> Tuple[int, int]:
    # Mostly copied from PIL implementation:
    # The only difference is with transformed points as input matrix has zero translation part here and
    # PIL has a centered translation part.
    # https://github.com/python-pillow/Pillow/blob/11de3318867e4398057373ee9f12dcb33db7335c/src/PIL/Image.py#L2054

    a, b, c, d, e, f = matrix
    xx = []
    yy = []

    half_w = 0.5 * w
    half_h = 0.5 * h
    for x, y in ((-half_w, -half_h), (half_w, -half_h), (half_w, half_h), (-half_w, half_h)):
        nx = a * x + b * y + c
        ny = d * x + e * y + f
        xx.append(nx + half_w)
        yy.append(ny + half_h)

    nw = math.ceil(max(xx)) - math.floor(min(xx))
    nh = math.ceil(max(yy)) - math.floor(min(yy))
    return int(nw), int(nh)  # w, h


595
def _apply_grid_transform(img: torch.Tensor, grid: torch.Tensor, mode: str, fill: _FillTypeJIT) -> torch.Tensor:
596
597
598
599
600
601
602
603
604
605
    input_shape = img.shape
    output_height, output_width = grid.shape[1], grid.shape[2]
    num_channels, input_height, input_width = input_shape[-3:]
    output_shape = input_shape[:-3] + (num_channels, output_height, output_width)

    if img.numel() == 0:
        return img.reshape(output_shape)

    img = img.reshape(-1, num_channels, input_height, input_width)
    squashed_batch_size = img.shape[0]
606

607
608
609
610
    # 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)

611
    if squashed_batch_size > 1:
612
        # Apply same grid to a batch of images
613
        grid = grid.expand(squashed_batch_size, -1, -1, -1)
614
615
616

    # Append a dummy mask for customized fill colors, should be faster than grid_sample() twice
    if fill is not None:
617
618
619
        mask = torch.ones(
            (squashed_batch_size, 1, input_height, input_width), dtype=float_img.dtype, device=float_img.device
        )
620
621
622
623
624
625
626
627
        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)
628
        fill_list = fill if isinstance(fill, (tuple, list)) else [float(fill)]  # type: ignore[arg-type]
629
630
        fill_img = torch.tensor(fill_list, dtype=float_img.dtype, device=float_img.device).view(1, -1, 1, 1)
        if mode == "nearest":
631
            float_img = torch.where(mask < 0.5, fill_img.expand_as(float_img), float_img)
632
633
634
635
636
        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)

637
638
    img = float_img.round_().to(img.dtype) if not fp else float_img

639
    return img.reshape(output_shape)
640
641
642
643
644
645


def _assert_grid_transform_inputs(
    image: torch.Tensor,
    matrix: Optional[List[float]],
    interpolation: str,
646
    fill: _FillTypeJIT,
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
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
    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)


702
@_register_kernel_internal(affine, torch.Tensor)
703
@_register_kernel_internal(affine, tv_tensors.Image)
704
def affine_image(
705
    image: torch.Tensor,
706
    angle: Union[int, float],
707
708
709
    translate: List[float],
    scale: float,
    shear: List[float],
710
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
711
    fill: _FillTypeJIT = None,
712
713
    center: Optional[List[float]] = None,
) -> torch.Tensor:
714
715
    interpolation = _check_interpolation(interpolation)

716
717
    angle, translate, shear, center = _affine_parse_args(angle, translate, scale, shear, interpolation, center)

718
719
    height, width = image.shape[-2:]

720
721
722
    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.
723
        center_f = [(c - s * 0.5) for c, s in zip(center, [width, height])]
724

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

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

730
    dtype = image.dtype if torch.is_floating_point(image) else torch.float32
731
732
    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)
733
    return _apply_grid_transform(image, grid, interpolation.value, fill=fill)
734
735


736
@_register_kernel_internal(affine, PIL.Image.Image)
737
def _affine_image_pil(
738
    image: PIL.Image.Image,
739
    angle: Union[int, float],
740
741
742
    translate: List[float],
    scale: float,
    shear: List[float],
743
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
744
    fill: _FillTypeJIT = None,
745
746
    center: Optional[List[float]] = None,
) -> PIL.Image.Image:
747
    interpolation = _check_interpolation(interpolation)
748
749
750
751
752
753
    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:
754
        height, width = _get_size_image_pil(image)
755
756
757
        center = [width * 0.5, height * 0.5]
    matrix = _get_inverse_affine_matrix(center, angle, translate, scale, shear)

758
    return _FP.affine(image, matrix, interpolation=pil_modes_mapping[interpolation], fill=fill)
759
760


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

    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 = (
Nicolas Hug's avatar
Nicolas Hug committed
781
        convert_bounding_box_format(
782
            bounding_boxes, old_format=format, new_format=tv_tensors.BoundingBoxFormat.XYXY, inplace=True
783
784
785
        )
    ).reshape(-1, 4)

786
787
788
    angle, translate, shear, center = _affine_parse_args(
        angle, translate, scale, shear, InterpolationMode.NEAREST, center
    )
789

790
    if center is None:
Philip Meier's avatar
Philip Meier committed
791
        height, width = canvas_size
792
793
        center = [width * 0.5, height * 0.5]

794
795
796
797
798
799
800
    affine_vector = _get_inverse_affine_matrix(center, angle, translate, scale, shear, inverted=False)
    transposed_affine_matrix = (
        torch.tensor(
            affine_vector,
            dtype=dtype,
            device=device,
        )
801
        .reshape(2, 3)
802
803
        .T
    )
804
805
806
807
    # 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)]
808
    points = bounding_boxes[:, [[0, 1], [2, 1], [2, 3], [0, 3]]].reshape(-1, 2)
809
    points = torch.cat([points, torch.ones(points.shape[0], 1, device=device, dtype=dtype)], dim=-1)
810
    # 2) Now let's transform the points using affine matrix
811
    transformed_points = torch.matmul(points, transposed_affine_matrix)
812
813
    # 3) Reshape transformed points to [N boxes, 4 points, x/y coords]
    # and compute bounding box from 4 transformed points:
814
    transformed_points = transformed_points.reshape(-1, 4, 2)
815
    out_bbox_mins, out_bbox_maxs = torch.aminmax(transformed_points, dim=1)
816
    out_bboxes = torch.cat([out_bbox_mins, out_bbox_maxs], dim=1)
817
818
819
820

    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
821
        height, width = canvas_size
822
823
824
        points = torch.tensor(
            [
                [0.0, 0.0, 1.0],
825
826
827
                [0.0, float(height), 1.0],
                [float(width), float(height), 1.0],
                [float(width), 0.0, 1.0],
828
829
830
831
            ],
            dtype=dtype,
            device=device,
        )
832
        new_points = torch.matmul(points, transposed_affine_matrix)
833
        tr = torch.amin(new_points, dim=0, keepdim=True)
834
        # Translate bounding boxes
835
        out_bboxes.sub_(tr.repeat((1, 2)))
836
837
        # Estimate meta-data for image with inverted=True
        affine_vector = _get_inverse_affine_matrix(center, angle, translate, scale, shear)
838
        new_width, new_height = _compute_affine_output_size(affine_vector, width, height)
Philip Meier's avatar
Philip Meier committed
839
        canvas_size = (new_height, new_width)
840

841
    out_bboxes = clamp_bounding_boxes(out_bboxes, format=tv_tensors.BoundingBoxFormat.XYXY, canvas_size=canvas_size)
Nicolas Hug's avatar
Nicolas Hug committed
842
    out_bboxes = convert_bounding_box_format(
843
        out_bboxes, old_format=tv_tensors.BoundingBoxFormat.XYXY, new_format=format, inplace=True
844
845
846
    ).reshape(original_shape)

    out_bboxes = out_bboxes.to(original_dtype)
Philip Meier's avatar
Philip Meier committed
847
    return out_bboxes, canvas_size
848
849


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


874
@_register_kernel_internal(affine, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False)
875
def _affine_bounding_boxes_dispatch(
876
    inpt: tv_tensors.BoundingBoxes,
877
878
879
880
881
882
    angle: Union[int, float],
    translate: List[float],
    scale: float,
    shear: List[float],
    center: Optional[List[float]] = None,
    **kwargs,
883
) -> tv_tensors.BoundingBoxes:
884
885
886
887
888
889
890
891
892
893
    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,
    )
894
    return tv_tensors.wrap(output, like=inpt)
895
896


897
898
def affine_mask(
    mask: torch.Tensor,
899
    angle: Union[int, float],
900
901
902
    translate: List[float],
    scale: float,
    shear: List[float],
903
    fill: _FillTypeJIT = None,
904
905
    center: Optional[List[float]] = None,
) -> torch.Tensor:
906
907
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
908
909
910
911
        needs_squeeze = True
    else:
        needs_squeeze = False

912
    output = affine_image(
913
        mask,
914
915
916
917
918
        angle=angle,
        translate=translate,
        scale=scale,
        shear=shear,
        interpolation=InterpolationMode.NEAREST,
919
        fill=fill,
920
921
922
        center=center,
    )

923
924
925
926
927
    if needs_squeeze:
        output = output.squeeze(0)

    return output

928

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


952
@_register_kernel_internal(affine, tv_tensors.Video)
953
954
955
956
957
958
def affine_video(
    video: torch.Tensor,
    angle: Union[int, float],
    translate: List[float],
    scale: float,
    shear: List[float],
959
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
960
    fill: _FillTypeJIT = None,
961
962
    center: Optional[List[float]] = None,
) -> torch.Tensor:
963
    return affine_image(
964
965
966
967
968
969
970
971
972
973
974
        video,
        angle=angle,
        translate=translate,
        scale=scale,
        shear=shear,
        interpolation=interpolation,
        fill=fill,
        center=center,
    )


975
def rotate(
976
    inpt: torch.Tensor,
977
    angle: float,
978
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
979
    expand: bool = False,
980
    center: Optional[List[float]] = None,
981
982
    fill: _FillTypeJIT = None,
) -> torch.Tensor:
983
    """See :class:`~torchvision.transforms.v2.RandomRotation` for details."""
984
    if torch.jit.is_scripting():
985
        return rotate_image(inpt, angle=angle, interpolation=interpolation, expand=expand, fill=fill, center=center)
986

987
    _log_api_usage_once(rotate)
988

989
990
991
992
993
    kernel = _get_kernel(rotate, type(inpt))
    return kernel(inpt, angle=angle, interpolation=interpolation, expand=expand, fill=fill, center=center)


@_register_kernel_internal(rotate, torch.Tensor)
994
@_register_kernel_internal(rotate, tv_tensors.Image)
995
def rotate_image(
996
    image: torch.Tensor,
997
    angle: float,
998
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
999
1000
    expand: bool = False,
    center: Optional[List[float]] = None,
1001
    fill: _FillTypeJIT = None,
1002
) -> torch.Tensor:
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
    angle = angle % 360  # shift angle to [0, 360) range

    # fast path: transpose without affine transform
    if center is None:
        if angle == 0:
            return image.clone()
        if angle == 180:
            return torch.rot90(image, k=2, dims=(-2, -1))

        if expand or image.shape[-1] == image.shape[-2]:
            if angle == 90:
                return torch.rot90(image, k=1, dims=(-2, -1))
            if angle == 270:
                return torch.rot90(image, k=3, dims=(-2, -1))

1018
1019
    interpolation = _check_interpolation(interpolation)

1020
    input_height, input_width = image.shape[-2:]
1021

1022
1023
    center_f = [0.0, 0.0]
    if center is not None:
1024
        # Center values should be in pixel coordinates but translated such that (0, 0) corresponds to image center.
1025
        center_f = [(c - s * 0.5) for c, s in zip(center, [input_width, input_height])]
1026
1027
1028
1029

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

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

1033
1034
1035
1036
1037
1038
1039
    output_width, output_height = (
        _compute_affine_output_size(matrix, input_width, input_height) if expand else (input_width, input_height)
    )
    dtype = image.dtype if torch.is_floating_point(image) else torch.float32
    theta = torch.tensor(matrix, dtype=dtype, device=image.device).reshape(1, 2, 3)
    grid = _affine_grid(theta, w=input_width, h=input_height, ow=output_width, oh=output_height)
    return _apply_grid_transform(image, grid, interpolation.value, fill=fill)
1040
1041


1042
@_register_kernel_internal(rotate, PIL.Image.Image)
1043
def _rotate_image_pil(
1044
    image: PIL.Image.Image,
1045
    angle: float,
1046
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
1047
1048
    expand: bool = False,
    center: Optional[List[float]] = None,
1049
    fill: _FillTypeJIT = None,
1050
) -> PIL.Image.Image:
1051
1052
    interpolation = _check_interpolation(interpolation)

1053
    return _FP.rotate(
1054
        image, angle, interpolation=pil_modes_mapping[interpolation], expand=expand, fill=fill, center=center
1055
1056
1057
    )


1058
1059
def rotate_bounding_boxes(
    bounding_boxes: torch.Tensor,
1060
    format: tv_tensors.BoundingBoxFormat,
Philip Meier's avatar
Philip Meier committed
1061
    canvas_size: Tuple[int, int],
1062
1063
1064
    angle: float,
    expand: bool = False,
    center: Optional[List[float]] = None,
1065
) -> Tuple[torch.Tensor, Tuple[int, int]]:
1066
1067
    return _affine_bounding_boxes_with_expand(
        bounding_boxes,
1068
        format=format,
Philip Meier's avatar
Philip Meier committed
1069
        canvas_size=canvas_size,
1070
1071
1072
1073
1074
1075
1076
        angle=-angle,
        translate=[0.0, 0.0],
        scale=1.0,
        shear=[0.0, 0.0],
        center=center,
        expand=expand,
    )
1077
1078


1079
@_register_kernel_internal(rotate, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False)
1080
def _rotate_bounding_boxes_dispatch(
1081
1082
    inpt: tv_tensors.BoundingBoxes, angle: float, expand: bool = False, center: Optional[List[float]] = None, **kwargs
) -> tv_tensors.BoundingBoxes:
1083
1084
1085
1086
1087
1088
1089
1090
    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,
    )
1091
    return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size)
1092
1093


1094
1095
def rotate_mask(
    mask: torch.Tensor,
1096
1097
1098
    angle: float,
    expand: bool = False,
    center: Optional[List[float]] = None,
1099
    fill: _FillTypeJIT = None,
1100
) -> torch.Tensor:
1101
1102
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
1103
1104
1105
1106
        needs_squeeze = True
    else:
        needs_squeeze = False

1107
    output = rotate_image(
1108
        mask,
1109
1110
1111
        angle=angle,
        expand=expand,
        interpolation=InterpolationMode.NEAREST,
1112
        fill=fill,
1113
1114
1115
        center=center,
    )

1116
1117
1118
1119
1120
    if needs_squeeze:
        output = output.squeeze(0)

    return output

1121

1122
@_register_kernel_internal(rotate, tv_tensors.Mask, tv_tensor_wrapper=False)
1123
def _rotate_mask_dispatch(
1124
    inpt: tv_tensors.Mask,
1125
1126
1127
    angle: float,
    expand: bool = False,
    center: Optional[List[float]] = None,
1128
    fill: _FillTypeJIT = None,
1129
    **kwargs,
1130
) -> tv_tensors.Mask:
1131
    output = rotate_mask(inpt.as_subclass(torch.Tensor), angle=angle, expand=expand, fill=fill, center=center)
1132
    return tv_tensors.wrap(output, like=inpt)
1133
1134


1135
@_register_kernel_internal(rotate, tv_tensors.Video)
1136
1137
1138
def rotate_video(
    video: torch.Tensor,
    angle: float,
1139
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
1140
1141
    expand: bool = False,
    center: Optional[List[float]] = None,
1142
    fill: _FillTypeJIT = None,
1143
) -> torch.Tensor:
1144
    return rotate_image(video, angle, interpolation=interpolation, expand=expand, fill=fill, center=center)
1145
1146


1147
def pad(
1148
    inpt: torch.Tensor,
1149
1150
1151
    padding: List[int],
    fill: Optional[Union[int, float, List[float]]] = None,
    padding_mode: str = "constant",
1152
) -> torch.Tensor:
1153
    """See :class:`~torchvision.transforms.v2.Pad` for details."""
1154
    if torch.jit.is_scripting():
1155
        return pad_image(inpt, padding=padding, fill=fill, padding_mode=padding_mode)
1156

1157
    _log_api_usage_once(pad)
1158

1159
1160
    kernel = _get_kernel(pad, type(inpt))
    return kernel(inpt, padding=padding, fill=fill, padding_mode=padding_mode)
1161
1162


1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
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]
1185

1186

1187
@_register_kernel_internal(pad, torch.Tensor)
1188
@_register_kernel_internal(pad, tv_tensors.Image)
1189
def pad_image(
1190
    image: torch.Tensor,
1191
1192
    padding: List[int],
    fill: Optional[Union[int, float, List[float]]] = None,
1193
1194
    padding_mode: str = "constant",
) -> torch.Tensor:
1195
    # Be aware that while `padding` has order `[left, top, right, bottom]`, `torch_padding` uses
1196
1197
1198
1199
    # `[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)

1200
    if padding_mode not in ("constant", "edge", "reflect", "symmetric"):
1201
1202
1203
1204
1205
        raise ValueError(
            f"`padding_mode` should be either `'constant'`, `'edge'`, `'reflect'` or `'symmetric'`, "
            f"but got `'{padding_mode}'`."
        )

1206
    if fill is None:
1207
1208
1209
1210
1211
1212
        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)
1213
    else:
1214
        return _pad_with_vector_fill(image, torch_padding, fill=fill, padding_mode=padding_mode)
1215
1216
1217


def _pad_with_scalar_fill(
1218
    image: torch.Tensor,
1219
1220
1221
    torch_padding: List[int],
    fill: Union[int, float],
    padding_mode: str,
1222
) -> torch.Tensor:
1223
1224
    shape = image.shape
    num_channels, height, width = shape[-3:]
1225

1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
    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
1249

1250
1251
1252
1253
1254
        image = torch_pad(image, torch_padding, mode=padding_mode)

        if needs_cast:
            image = image.to(dtype)
    else:  # padding_mode == "symmetric"
1255
        image = _pad_symmetric(image, torch_padding)
1256
1257

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

1259
    return image.reshape(shape[:-3] + (num_channels, new_height, new_width))
1260
1261


1262
# TODO: This should be removed once torch_pad supports non-scalar padding values
1263
def _pad_with_vector_fill(
1264
    image: torch.Tensor,
1265
    torch_padding: List[int],
1266
    fill: List[float],
1267
    padding_mode: str,
1268
1269
1270
1271
) -> torch.Tensor:
    if padding_mode != "constant":
        raise ValueError(f"Padding mode '{padding_mode}' is not supported if fill is not scalar")

1272
1273
    output = _pad_with_scalar_fill(image, torch_padding, fill=0, padding_mode="constant")
    left, right, top, bottom = torch_padding
1274
1275
1276
1277
1278

    # We are creating the tensor in the autodetected dtype first and convert to the right one after to avoid an implicit
    # float -> int conversion. That happens for example for the valid input of a uint8 image with floating point fill
    # value.
    fill = torch.tensor(fill, device=image.device).to(dtype=image.dtype).reshape(-1, 1, 1)
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290

    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


1291
_pad_image_pil = _register_kernel_internal(pad, PIL.Image.Image)(_FP.pad)
1292
1293


1294
@_register_kernel_internal(pad, tv_tensors.Mask)
1295
1296
def pad_mask(
    mask: torch.Tensor,
1297
1298
    padding: List[int],
    fill: Optional[Union[int, float, List[float]]] = None,
1299
1300
    padding_mode: str = "constant",
) -> torch.Tensor:
1301
1302
1303
    if fill is None:
        fill = 0

1304
    if isinstance(fill, (tuple, list)):
1305
1306
        raise ValueError("Non-scalar fill value is not supported")

1307
1308
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
1309
1310
1311
1312
        needs_squeeze = True
    else:
        needs_squeeze = False

1313
    output = pad_image(mask, padding=padding, fill=fill, padding_mode=padding_mode)
1314
1315
1316
1317
1318

    if needs_squeeze:
        output = output.squeeze(0)

    return output
1319
1320


1321
1322
def pad_bounding_boxes(
    bounding_boxes: torch.Tensor,
1323
    format: tv_tensors.BoundingBoxFormat,
Philip Meier's avatar
Philip Meier committed
1324
    canvas_size: Tuple[int, int],
1325
    padding: List[int],
vfdev's avatar
vfdev committed
1326
    padding_mode: str = "constant",
1327
) -> Tuple[torch.Tensor, Tuple[int, int]]:
vfdev's avatar
vfdev committed
1328
1329
1330
1331
    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")

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

1334
    if format == tv_tensors.BoundingBoxFormat.XYXY:
1335
1336
1337
        pad = [left, top, left, top]
    else:
        pad = [left, top, 0, 0]
1338
    bounding_boxes = bounding_boxes + torch.tensor(pad, dtype=bounding_boxes.dtype, device=bounding_boxes.device)
1339

Philip Meier's avatar
Philip Meier committed
1340
    height, width = canvas_size
1341
1342
    height += top + bottom
    width += left + right
Philip Meier's avatar
Philip Meier committed
1343
    canvas_size = (height, width)
1344

Philip Meier's avatar
Philip Meier committed
1345
    return clamp_bounding_boxes(bounding_boxes, format=format, canvas_size=canvas_size), canvas_size
1346
1347


1348
@_register_kernel_internal(pad, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False)
1349
def _pad_bounding_boxes_dispatch(
1350
1351
    inpt: tv_tensors.BoundingBoxes, padding: List[int], padding_mode: str = "constant", **kwargs
) -> tv_tensors.BoundingBoxes:
1352
1353
1354
1355
1356
1357
1358
    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,
    )
1359
    return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size)
1360
1361


1362
@_register_kernel_internal(pad, tv_tensors.Video)
1363
1364
def pad_video(
    video: torch.Tensor,
1365
1366
    padding: List[int],
    fill: Optional[Union[int, float, List[float]]] = None,
1367
1368
    padding_mode: str = "constant",
) -> torch.Tensor:
1369
    return pad_image(video, padding, fill=fill, padding_mode=padding_mode)
1370
1371


1372
def crop(inpt: torch.Tensor, top: int, left: int, height: int, width: int) -> torch.Tensor:
1373
    """See :class:`~torchvision.transforms.v2.RandomCrop` for details."""
1374
    if torch.jit.is_scripting():
1375
        return crop_image(inpt, top=top, left=left, height=height, width=width)
1376
1377

    _log_api_usage_once(crop)
1378

1379
1380
    kernel = _get_kernel(crop, type(inpt))
    return kernel(inpt, top=top, left=left, height=height, width=width)
1381

1382
1383

@_register_kernel_internal(crop, torch.Tensor)
1384
@_register_kernel_internal(crop, tv_tensors.Image)
1385
def crop_image(image: torch.Tensor, top: int, left: int, height: int, width: int) -> torch.Tensor:
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
    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]


1403
1404
_crop_image_pil = _FP.crop
_register_kernel_internal(crop, PIL.Image.Image)(_crop_image_pil)
1405
1406


1407
1408
def crop_bounding_boxes(
    bounding_boxes: torch.Tensor,
1409
    format: tv_tensors.BoundingBoxFormat,
1410
1411
    top: int,
    left: int,
1412
1413
1414
    height: int,
    width: int,
) -> Tuple[torch.Tensor, Tuple[int, int]]:
1415

1416
    # Crop or implicit pad if left and/or top have negative values:
1417
    if format == tv_tensors.BoundingBoxFormat.XYXY:
1418
        sub = [left, top, left, top]
1419
    else:
1420
1421
        sub = [left, top, 0, 0]

1422
    bounding_boxes = bounding_boxes - torch.tensor(sub, dtype=bounding_boxes.dtype, device=bounding_boxes.device)
Philip Meier's avatar
Philip Meier committed
1423
    canvas_size = (height, width)
1424

Philip Meier's avatar
Philip Meier committed
1425
    return clamp_bounding_boxes(bounding_boxes, format=format, canvas_size=canvas_size), canvas_size
1426
1427


1428
@_register_kernel_internal(crop, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False)
1429
def _crop_bounding_boxes_dispatch(
1430
1431
    inpt: tv_tensors.BoundingBoxes, top: int, left: int, height: int, width: int
) -> tv_tensors.BoundingBoxes:
1432
1433
1434
    output, canvas_size = crop_bounding_boxes(
        inpt.as_subclass(torch.Tensor), format=inpt.format, top=top, left=left, height=height, width=width
    )
1435
    return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size)
1436
1437


1438
@_register_kernel_internal(crop, tv_tensors.Mask)
1439
def crop_mask(mask: torch.Tensor, top: int, left: int, height: int, width: int) -> torch.Tensor:
1440
1441
1442
1443
1444
1445
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
        needs_squeeze = True
    else:
        needs_squeeze = False

1446
    output = crop_image(mask, top, left, height, width)
1447
1448
1449
1450
1451

    if needs_squeeze:
        output = output.squeeze(0)

    return output
1452
1453


1454
@_register_kernel_internal(crop, tv_tensors.Video)
1455
def crop_video(video: torch.Tensor, top: int, left: int, height: int, width: int) -> torch.Tensor:
1456
    return crop_image(video, top, left, height, width)
1457
1458


1459
def perspective(
1460
    inpt: torch.Tensor,
1461
1462
1463
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1464
    fill: _FillTypeJIT = None,
1465
    coefficients: Optional[List[float]] = None,
1466
) -> torch.Tensor:
1467
    """See :class:`~torchvision.transforms.v2.RandomPerspective` for details."""
1468
    if torch.jit.is_scripting():
1469
        return perspective_image(
1470
1471
1472
1473
1474
1475
            inpt,
            startpoints=startpoints,
            endpoints=endpoints,
            interpolation=interpolation,
            fill=fill,
            coefficients=coefficients,
1476
        )
1477

1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
    _log_api_usage_once(perspective)

    kernel = _get_kernel(perspective, type(inpt))
    return kernel(
        inpt,
        startpoints=startpoints,
        endpoints=endpoints,
        interpolation=interpolation,
        fill=fill,
        coefficients=coefficients,
    )

1490

1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
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)
1506
    x_grid = torch.linspace(d, ow + d - 1.0, steps=ow, device=device, dtype=dtype)
1507
    base_grid[..., 0].copy_(x_grid)
1508
    y_grid = torch.linspace(d, oh + d - 1.0, steps=oh, device=device, dtype=dtype).unsqueeze_(-1)
1509
1510
1511
1512
    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))
1513
1514
1515
    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))
1516
1517
1518
1519
1520

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


1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
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.")


1538
@_register_kernel_internal(perspective, torch.Tensor)
1539
@_register_kernel_internal(perspective, tv_tensors.Image)
1540
def perspective_image(
1541
    image: torch.Tensor,
1542
1543
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
1544
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1545
    fill: _FillTypeJIT = None,
1546
    coefficients: Optional[List[float]] = None,
1547
) -> torch.Tensor:
1548
    perspective_coeffs = _perspective_coefficients(startpoints, endpoints, coefficients)
1549
1550
    interpolation = _check_interpolation(interpolation)

1551
    _assert_grid_transform_inputs(
1552
1553
1554
1555
1556
1557
1558
1559
        image,
        matrix=None,
        interpolation=interpolation.value,
        fill=fill,
        supported_interpolation_modes=["nearest", "bilinear"],
        coeffs=perspective_coeffs,
    )

1560
    oh, ow = image.shape[-2:]
1561
    dtype = image.dtype if torch.is_floating_point(image) else torch.float32
1562
    grid = _perspective_grid(perspective_coeffs, ow=ow, oh=oh, dtype=dtype, device=image.device)
1563
    return _apply_grid_transform(image, grid, interpolation.value, fill=fill)
1564
1565


1566
@_register_kernel_internal(perspective, PIL.Image.Image)
1567
def _perspective_image_pil(
1568
    image: PIL.Image.Image,
1569
1570
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
1571
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1572
    fill: _FillTypeJIT = None,
1573
    coefficients: Optional[List[float]] = None,
1574
) -> PIL.Image.Image:
1575
    perspective_coeffs = _perspective_coefficients(startpoints, endpoints, coefficients)
1576
    interpolation = _check_interpolation(interpolation)
1577
    return _FP.perspective(image, perspective_coeffs, interpolation=pil_modes_mapping[interpolation], fill=fill)
1578
1579


1580
1581
def perspective_bounding_boxes(
    bounding_boxes: torch.Tensor,
1582
    format: tv_tensors.BoundingBoxFormat,
Philip Meier's avatar
Philip Meier committed
1583
    canvas_size: Tuple[int, int],
1584
1585
1586
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
    coefficients: Optional[List[float]] = None,
1587
) -> torch.Tensor:
1588
1589
    if bounding_boxes.numel() == 0:
        return bounding_boxes
1590

1591
    perspective_coeffs = _perspective_coefficients(startpoints, endpoints, coefficients)
1592

1593
    original_shape = bounding_boxes.shape
Nicolas Hug's avatar
Nicolas Hug committed
1594
    # TODO: first cast to float if bbox is int64 before convert_bounding_box_format
1595
    bounding_boxes = (
1596
        convert_bounding_box_format(bounding_boxes, old_format=format, new_format=tv_tensors.BoundingBoxFormat.XYXY)
1597
    ).reshape(-1, 4)
1598

1599
1600
    dtype = bounding_boxes.dtype if torch.is_floating_point(bounding_boxes) else torch.float32
    device = bounding_boxes.device
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631

    # 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,
    ]

1632
1633
    theta1 = torch.tensor(
        [[inv_coeffs[0], inv_coeffs[1], inv_coeffs[2]], [inv_coeffs[3], inv_coeffs[4], inv_coeffs[5]]],
1634
1635
1636
1637
        dtype=dtype,
        device=device,
    )

1638
1639
1640
1641
    theta2 = torch.tensor(
        [[inv_coeffs[6], inv_coeffs[7], 1.0], [inv_coeffs[6], inv_coeffs[7], 1.0]], dtype=dtype, device=device
    )

1642
1643
1644
1645
    # 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)]
1646
    points = bounding_boxes[:, [[0, 1], [2, 1], [2, 3], [0, 3]]].reshape(-1, 2)
1647
1648
1649
1650
1651
    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)

1652
1653
    numer_points = torch.matmul(points, theta1.T)
    denom_points = torch.matmul(points, theta2.T)
1654
    transformed_points = numer_points.div_(denom_points)
1655
1656
1657

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

1661
1662
    out_bboxes = clamp_bounding_boxes(
        torch.cat([out_bbox_mins, out_bbox_maxs], dim=1).to(bounding_boxes.dtype),
1663
        format=tv_tensors.BoundingBoxFormat.XYXY,
Philip Meier's avatar
Philip Meier committed
1664
        canvas_size=canvas_size,
1665
    )
1666
1667
1668

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

Nicolas Hug's avatar
Nicolas Hug committed
1669
    return convert_bounding_box_format(
1670
        out_bboxes, old_format=tv_tensors.BoundingBoxFormat.XYXY, new_format=format, inplace=True
1671
    ).reshape(original_shape)
1672
1673


1674
@_register_kernel_internal(perspective, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False)
1675
def _perspective_bounding_boxes_dispatch(
1676
    inpt: tv_tensors.BoundingBoxes,
1677
1678
1679
1680
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
    coefficients: Optional[List[float]] = None,
    **kwargs,
1681
) -> tv_tensors.BoundingBoxes:
1682
1683
1684
1685
1686
1687
1688
1689
    output = perspective_bounding_boxes(
        inpt.as_subclass(torch.Tensor),
        format=inpt.format,
        canvas_size=inpt.canvas_size,
        startpoints=startpoints,
        endpoints=endpoints,
        coefficients=coefficients,
    )
1690
    return tv_tensors.wrap(output, like=inpt)
1691
1692


1693
1694
def perspective_mask(
    mask: torch.Tensor,
1695
1696
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
1697
    fill: _FillTypeJIT = None,
1698
    coefficients: Optional[List[float]] = None,
1699
) -> torch.Tensor:
1700
1701
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
1702
1703
1704
1705
        needs_squeeze = True
    else:
        needs_squeeze = False

1706
    output = perspective_image(
1707
        mask, startpoints, endpoints, interpolation=InterpolationMode.NEAREST, fill=fill, coefficients=coefficients
1708
    )
1709

1710
1711
1712
1713
1714
    if needs_squeeze:
        output = output.squeeze(0)

    return output

1715

1716
@_register_kernel_internal(perspective, tv_tensors.Mask, tv_tensor_wrapper=False)
1717
def _perspective_mask_dispatch(
1718
    inpt: tv_tensors.Mask,
1719
1720
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
1721
    fill: _FillTypeJIT = None,
1722
1723
    coefficients: Optional[List[float]] = None,
    **kwargs,
1724
) -> tv_tensors.Mask:
1725
1726
1727
1728
1729
1730
1731
    output = perspective_mask(
        inpt.as_subclass(torch.Tensor),
        startpoints=startpoints,
        endpoints=endpoints,
        fill=fill,
        coefficients=coefficients,
    )
1732
    return tv_tensors.wrap(output, like=inpt)
1733
1734


1735
@_register_kernel_internal(perspective, tv_tensors.Video)
1736
1737
def perspective_video(
    video: torch.Tensor,
1738
1739
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
1740
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1741
    fill: _FillTypeJIT = None,
1742
    coefficients: Optional[List[float]] = None,
1743
) -> torch.Tensor:
1744
    return perspective_image(
1745
1746
        video, startpoints, endpoints, interpolation=interpolation, fill=fill, coefficients=coefficients
    )
1747
1748


1749
def elastic(
1750
    inpt: torch.Tensor,
1751
    displacement: torch.Tensor,
1752
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1753
1754
    fill: _FillTypeJIT = None,
) -> torch.Tensor:
1755
    """See :class:`~torchvision.transforms.v2.ElasticTransform` for details."""
1756
    if torch.jit.is_scripting():
1757
        return elastic_image(inpt, displacement=displacement, interpolation=interpolation, fill=fill)
1758
1759
1760
1761
1762

    _log_api_usage_once(elastic)

    kernel = _get_kernel(elastic, type(inpt))
    return kernel(inpt, displacement=displacement, interpolation=interpolation, fill=fill)
1763
1764


1765
1766
1767
elastic_transform = elastic


1768
@_register_kernel_internal(elastic, torch.Tensor)
1769
@_register_kernel_internal(elastic, tv_tensors.Image)
1770
def elastic_image(
1771
    image: torch.Tensor,
1772
    displacement: torch.Tensor,
1773
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1774
    fill: _FillTypeJIT = None,
1775
) -> torch.Tensor:
Philip Meier's avatar
Philip Meier committed
1776
1777
1778
    if not isinstance(displacement, torch.Tensor):
        raise TypeError("Argument displacement should be a Tensor")

1779
1780
    interpolation = _check_interpolation(interpolation)

1781
    height, width = image.shape[-2:]
1782
    device = image.device
1783
    dtype = image.dtype if torch.is_floating_point(image) else torch.float32
1784
1785
1786
1787
1788
1789
1790

    # 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

1791
    # We are aware that if input image dtype is uint8 and displacement is float64 then
1792
    # displacement will be cast to float32 and all computations will be done with float32
1793
    # We can fix this later if needed
1794

1795
    expected_shape = (1, height, width, 2)
1796
1797
1798
    if expected_shape != displacement.shape:
        raise ValueError(f"Argument displacement shape should be {expected_shape}, but given {displacement.shape}")

1799
1800
1801
    grid = _create_identity_grid((height, width), device=device, dtype=dtype).add_(
        displacement.to(dtype=dtype, device=device)
    )
1802
    output = _apply_grid_transform(image, grid, interpolation.value, fill=fill)
1803

1804
1805
1806
    if is_cpu_half:
        output = output.to(torch.float16)

1807
    return output
1808
1809


1810
@_register_kernel_internal(elastic, PIL.Image.Image)
1811
def _elastic_image_pil(
1812
    image: PIL.Image.Image,
1813
    displacement: torch.Tensor,
1814
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1815
    fill: _FillTypeJIT = None,
1816
) -> PIL.Image.Image:
1817
    t_img = pil_to_tensor(image)
1818
    output = elastic_image(t_img, displacement, interpolation=interpolation, fill=fill)
1819
    return to_pil_image(output, mode=image.mode)
1820
1821


1822
def _create_identity_grid(size: Tuple[int, int], device: torch.device, dtype: torch.dtype) -> torch.Tensor:
1823
    sy, sx = size
1824
1825
    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)
1826
1827
    base_grid[..., 0].copy_(x_grid)

1828
    y_grid = torch.linspace((-sy + 1) / sy, (sy - 1) / sy, sy, device=device, dtype=dtype).unsqueeze_(-1)
1829
1830
1831
1832
1833
    base_grid[..., 1].copy_(y_grid)

    return base_grid


1834
1835
def elastic_bounding_boxes(
    bounding_boxes: torch.Tensor,
1836
    format: tv_tensors.BoundingBoxFormat,
Philip Meier's avatar
Philip Meier committed
1837
    canvas_size: Tuple[int, int],
1838
1839
    displacement: torch.Tensor,
) -> torch.Tensor:
Philip Meier's avatar
Philip Meier committed
1840
1841
1842
1843
1844
1845
    expected_shape = (1, canvas_size[0], canvas_size[1], 2)
    if not isinstance(displacement, torch.Tensor):
        raise TypeError("Argument displacement should be a Tensor")
    elif displacement.shape != expected_shape:
        raise ValueError(f"Argument displacement shape should be {expected_shape}, but given {displacement.shape}")

1846
1847
    if bounding_boxes.numel() == 0:
        return bounding_boxes
1848

1849
    # TODO: add in docstring about approximation we are doing for grid inversion
1850
1851
    device = bounding_boxes.device
    dtype = bounding_boxes.dtype if torch.is_floating_point(bounding_boxes) else torch.float32
1852
1853
1854

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

1856
    original_shape = bounding_boxes.shape
Nicolas Hug's avatar
Nicolas Hug committed
1857
    # TODO: first cast to float if bbox is int64 before convert_bounding_box_format
1858
    bounding_boxes = (
1859
        convert_bounding_box_format(bounding_boxes, old_format=format, new_format=tv_tensors.BoundingBoxFormat.XYXY)
1860
    ).reshape(-1, 4)
1861

Philip Meier's avatar
Philip Meier committed
1862
    id_grid = _create_identity_grid(canvas_size, device=device, dtype=dtype)
1863
1864
    # We construct an approximation of inverse grid as inv_grid = id_grid - displacement
    # This is not an exact inverse of the grid
1865
    inv_grid = id_grid.sub_(displacement)
1866
1867

    # Get points from bboxes
1868
    points = bounding_boxes[:, [[0, 1], [2, 1], [2, 3], [0, 3]]].reshape(-1, 2)
1869
1870
1871
1872
1873
    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]

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

1878
    transformed_points = transformed_points.reshape(-1, 4, 2)
1879
    out_bbox_mins, out_bbox_maxs = torch.aminmax(transformed_points, dim=1)
1880
1881
    out_bboxes = clamp_bounding_boxes(
        torch.cat([out_bbox_mins, out_bbox_maxs], dim=1).to(bounding_boxes.dtype),
1882
        format=tv_tensors.BoundingBoxFormat.XYXY,
Philip Meier's avatar
Philip Meier committed
1883
        canvas_size=canvas_size,
1884
    )
1885

Nicolas Hug's avatar
Nicolas Hug committed
1886
    return convert_bounding_box_format(
1887
        out_bboxes, old_format=tv_tensors.BoundingBoxFormat.XYXY, new_format=format, inplace=True
1888
    ).reshape(original_shape)
1889
1890


1891
@_register_kernel_internal(elastic, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False)
1892
def _elastic_bounding_boxes_dispatch(
1893
1894
    inpt: tv_tensors.BoundingBoxes, displacement: torch.Tensor, **kwargs
) -> tv_tensors.BoundingBoxes:
1895
1896
1897
    output = elastic_bounding_boxes(
        inpt.as_subclass(torch.Tensor), format=inpt.format, canvas_size=inpt.canvas_size, displacement=displacement
    )
1898
    return tv_tensors.wrap(output, like=inpt)
1899
1900


1901
1902
1903
def elastic_mask(
    mask: torch.Tensor,
    displacement: torch.Tensor,
1904
    fill: _FillTypeJIT = None,
1905
) -> torch.Tensor:
1906
1907
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
1908
1909
1910
1911
        needs_squeeze = True
    else:
        needs_squeeze = False

1912
    output = elastic_image(mask, displacement=displacement, interpolation=InterpolationMode.NEAREST, fill=fill)
1913
1914
1915
1916
1917

    if needs_squeeze:
        output = output.squeeze(0)

    return output
1918
1919


1920
@_register_kernel_internal(elastic, tv_tensors.Mask, tv_tensor_wrapper=False)
1921
def _elastic_mask_dispatch(
1922
1923
    inpt: tv_tensors.Mask, displacement: torch.Tensor, fill: _FillTypeJIT = None, **kwargs
) -> tv_tensors.Mask:
1924
    output = elastic_mask(inpt.as_subclass(torch.Tensor), displacement=displacement, fill=fill)
1925
    return tv_tensors.wrap(output, like=inpt)
1926
1927


1928
@_register_kernel_internal(elastic, tv_tensors.Video)
1929
1930
1931
def elastic_video(
    video: torch.Tensor,
    displacement: torch.Tensor,
1932
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1933
    fill: _FillTypeJIT = None,
1934
) -> torch.Tensor:
1935
    return elastic_image(video, displacement, interpolation=interpolation, fill=fill)
1936
1937


1938
def center_crop(inpt: torch.Tensor, output_size: List[int]) -> torch.Tensor:
1939
    """See :class:`~torchvision.transforms.v2.RandomCrop` for details."""
1940
    if torch.jit.is_scripting():
1941
        return center_crop_image(inpt, output_size=output_size)
1942
1943
1944
1945
1946

    _log_api_usage_once(center_crop)

    kernel = _get_kernel(center_crop, type(inpt))
    return kernel(inpt, output_size=output_size)
1947
1948


1949
1950
def _center_crop_parse_output_size(output_size: List[int]) -> List[int]:
    if isinstance(output_size, numbers.Number):
1951
1952
        s = int(output_size)
        return [s, s]
1953
    elif isinstance(output_size, (tuple, list)) and len(output_size) == 1:
1954
        return [output_size[0], output_size[0]]
1955
1956
    else:
        return list(output_size)
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975


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


1976
@_register_kernel_internal(center_crop, torch.Tensor)
1977
@_register_kernel_internal(center_crop, tv_tensors.Image)
1978
def center_crop_image(image: torch.Tensor, output_size: List[int]) -> torch.Tensor:
1979
    crop_height, crop_width = _center_crop_parse_output_size(output_size)
1980
1981
1982
1983
    shape = image.shape
    if image.numel() == 0:
        return image.reshape(shape[:-2] + (crop_height, crop_width))
    image_height, image_width = shape[-2:]
1984
1985
1986

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

1989
        image_height, image_width = image.shape[-2:]
1990
        if crop_width == image_width and crop_height == image_height:
1991
            return image
1992
1993

    crop_top, crop_left = _center_crop_compute_crop_anchor(crop_height, crop_width, image_height, image_width)
1994
    return image[..., crop_top : (crop_top + crop_height), crop_left : (crop_left + crop_width)]
1995
1996


1997
@_register_kernel_internal(center_crop, PIL.Image.Image)
1998
def _center_crop_image_pil(image: PIL.Image.Image, output_size: List[int]) -> PIL.Image.Image:
1999
    crop_height, crop_width = _center_crop_parse_output_size(output_size)
2000
    image_height, image_width = _get_size_image_pil(image)
2001
2002
2003

    if crop_height > image_height or crop_width > image_width:
        padding_ltrb = _center_crop_compute_padding(crop_height, crop_width, image_height, image_width)
2004
        image = _pad_image_pil(image, padding_ltrb, fill=0)
2005

2006
        image_height, image_width = _get_size_image_pil(image)
2007
        if crop_width == image_width and crop_height == image_height:
2008
            return image
2009
2010

    crop_top, crop_left = _center_crop_compute_crop_anchor(crop_height, crop_width, image_height, image_width)
2011
    return _crop_image_pil(image, crop_top, crop_left, crop_height, crop_width)
2012
2013


2014
2015
def center_crop_bounding_boxes(
    bounding_boxes: torch.Tensor,
2016
    format: tv_tensors.BoundingBoxFormat,
Philip Meier's avatar
Philip Meier committed
2017
    canvas_size: Tuple[int, int],
2018
    output_size: List[int],
2019
) -> Tuple[torch.Tensor, Tuple[int, int]]:
2020
    crop_height, crop_width = _center_crop_parse_output_size(output_size)
Philip Meier's avatar
Philip Meier committed
2021
    crop_top, crop_left = _center_crop_compute_crop_anchor(crop_height, crop_width, *canvas_size)
2022
2023
2024
    return crop_bounding_boxes(
        bounding_boxes, format, top=crop_top, left=crop_left, height=crop_height, width=crop_width
    )
2025
2026


2027
@_register_kernel_internal(center_crop, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False)
2028
def _center_crop_bounding_boxes_dispatch(
2029
2030
    inpt: tv_tensors.BoundingBoxes, output_size: List[int]
) -> tv_tensors.BoundingBoxes:
2031
2032
2033
    output, canvas_size = center_crop_bounding_boxes(
        inpt.as_subclass(torch.Tensor), format=inpt.format, canvas_size=inpt.canvas_size, output_size=output_size
    )
2034
    return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size)
2035
2036


2037
@_register_kernel_internal(center_crop, tv_tensors.Mask)
2038
2039
2040
def center_crop_mask(mask: torch.Tensor, output_size: List[int]) -> torch.Tensor:
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
2041
2042
2043
2044
        needs_squeeze = True
    else:
        needs_squeeze = False

2045
    output = center_crop_image(image=mask, output_size=output_size)
2046
2047
2048
2049
2050

    if needs_squeeze:
        output = output.squeeze(0)

    return output
2051
2052


2053
@_register_kernel_internal(center_crop, tv_tensors.Video)
2054
def center_crop_video(video: torch.Tensor, output_size: List[int]) -> torch.Tensor:
2055
    return center_crop_image(video, output_size)
2056
2057


2058
def resized_crop(
2059
    inpt: torch.Tensor,
2060
2061
2062
2063
2064
2065
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
2066
    antialias: Optional[bool] = True,
2067
) -> torch.Tensor:
2068
    """See :class:`~torchvision.transforms.v2.RandomResizedCrop` for details."""
2069
    if torch.jit.is_scripting():
2070
        return resized_crop_image(
2071
2072
2073
2074
2075
2076
2077
2078
            inpt,
            top=top,
            left=left,
            height=height,
            width=width,
            size=size,
            interpolation=interpolation,
            antialias=antialias,
2079
        )
2080

2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
    _log_api_usage_once(resized_crop)

    kernel = _get_kernel(resized_crop, type(inpt))
    return kernel(
        inpt,
        top=top,
        left=left,
        height=height,
        width=width,
        size=size,
        interpolation=interpolation,
        antialias=antialias,
    )
2094

2095
2096

@_register_kernel_internal(resized_crop, torch.Tensor)
2097
@_register_kernel_internal(resized_crop, tv_tensors.Image)
2098
def resized_crop_image(
2099
    image: torch.Tensor,
2100
2101
2102
2103
2104
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
2105
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
2106
    antialias: Optional[bool] = True,
2107
) -> torch.Tensor:
2108
2109
    image = crop_image(image, top, left, height, width)
    return resize_image(image, size, interpolation=interpolation, antialias=antialias)
2110
2111


2112
def _resized_crop_image_pil(
2113
    image: PIL.Image.Image,
2114
2115
2116
2117
2118
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
2119
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
2120
) -> PIL.Image.Image:
2121
2122
    image = _crop_image_pil(image, top, left, height, width)
    return _resize_image_pil(image, size, interpolation=interpolation)
2123
2124


2125
@_register_kernel_internal(resized_crop, PIL.Image.Image)
2126
def _resized_crop_image_pil_dispatch(
2127
2128
2129
2130
2131
2132
2133
    image: PIL.Image.Image,
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
2134
    antialias: Optional[bool] = True,
2135
2136
2137
) -> PIL.Image.Image:
    if antialias is False:
        warnings.warn("Anti-alias option is always applied for PIL Image input. Argument antialias is ignored.")
2138
    return _resized_crop_image_pil(
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
        image,
        top=top,
        left=left,
        height=height,
        width=width,
        size=size,
        interpolation=interpolation,
    )


2149
2150
def resized_crop_bounding_boxes(
    bounding_boxes: torch.Tensor,
2151
    format: tv_tensors.BoundingBoxFormat,
2152
2153
2154
2155
2156
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
2157
) -> Tuple[torch.Tensor, Tuple[int, int]]:
2158
2159
2160
2161
    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)


2162
@_register_kernel_internal(resized_crop, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False)
2163
def _resized_crop_bounding_boxes_dispatch(
2164
2165
    inpt: tv_tensors.BoundingBoxes, top: int, left: int, height: int, width: int, size: List[int], **kwargs
) -> tv_tensors.BoundingBoxes:
2166
2167
2168
    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
    )
2169
    return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size)
2170
2171


2172
def resized_crop_mask(
2173
2174
2175
2176
2177
2178
2179
    mask: torch.Tensor,
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
) -> torch.Tensor:
2180
2181
    mask = crop_mask(mask, top, left, height, width)
    return resize_mask(mask, size)
2182
2183


2184
@_register_kernel_internal(resized_crop, tv_tensors.Mask, tv_tensor_wrapper=False)
2185
def _resized_crop_mask_dispatch(
2186
2187
    inpt: tv_tensors.Mask, top: int, left: int, height: int, width: int, size: List[int], **kwargs
) -> tv_tensors.Mask:
2188
2189
2190
    output = resized_crop_mask(
        inpt.as_subclass(torch.Tensor), top=top, left=left, height=height, width=width, size=size
    )
2191
    return tv_tensors.wrap(output, like=inpt)
2192
2193


2194
@_register_kernel_internal(resized_crop, tv_tensors.Video)
2195
2196
2197
2198
2199
2200
2201
def resized_crop_video(
    video: torch.Tensor,
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
2202
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
2203
    antialias: Optional[bool] = True,
2204
) -> torch.Tensor:
2205
    return resized_crop_image(
2206
2207
2208
2209
        video, top, left, height, width, antialias=antialias, size=size, interpolation=interpolation
    )


2210
def five_crop(
2211
2212
    inpt: torch.Tensor, size: List[int]
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
2213
    """See :class:`~torchvision.transforms.v2.FiveCrop` for details."""
2214
    if torch.jit.is_scripting():
2215
        return five_crop_image(inpt, size=size)
2216
2217
2218
2219
2220

    _log_api_usage_once(five_crop)

    kernel = _get_kernel(five_crop, type(inpt))
    return kernel(inpt, size=size)
2221
2222


2223
2224
def _parse_five_crop_size(size: List[int]) -> List[int]:
    if isinstance(size, numbers.Number):
2225
2226
        s = int(size)
        size = [s, s]
2227
    elif isinstance(size, (tuple, list)) and len(size) == 1:
2228
2229
        s = size[0]
        size = [s, s]
2230
2231
2232
2233
2234
2235
2236

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

    return size


2237
@_register_five_ten_crop_kernel_internal(five_crop, torch.Tensor)
2238
@_register_five_ten_crop_kernel_internal(five_crop, tv_tensors.Image)
2239
def five_crop_image(
2240
    image: torch.Tensor, size: List[int]
2241
2242
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    crop_height, crop_width = _parse_five_crop_size(size)
2243
    image_height, image_width = image.shape[-2:]
2244
2245

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

2248
2249
2250
2251
2252
    tl = crop_image(image, 0, 0, crop_height, crop_width)
    tr = crop_image(image, 0, image_width - crop_width, crop_height, crop_width)
    bl = crop_image(image, image_height - crop_height, 0, crop_height, crop_width)
    br = crop_image(image, image_height - crop_height, image_width - crop_width, crop_height, crop_width)
    center = center_crop_image(image, [crop_height, crop_width])
2253
2254
2255
2256

    return tl, tr, bl, br, center


2257
@_register_five_ten_crop_kernel_internal(five_crop, PIL.Image.Image)
2258
def _five_crop_image_pil(
2259
    image: PIL.Image.Image, size: List[int]
2260
2261
) -> 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)
2262
    image_height, image_width = _get_size_image_pil(image)
2263
2264

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

2267
2268
2269
2270
2271
    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])
2272
2273
2274
2275

    return tl, tr, bl, br, center


2276
@_register_five_ten_crop_kernel_internal(five_crop, tv_tensors.Video)
2277
2278
2279
def five_crop_video(
    video: torch.Tensor, size: List[int]
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
2280
    return five_crop_image(video, size)
2281
2282


2283
def ten_crop(
2284
    inpt: torch.Tensor, size: List[int], vertical_flip: bool = False
2285
) -> Tuple[
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
2296
]:
2297
    """See :class:`~torchvision.transforms.v2.TenCrop` for details."""
2298
    if torch.jit.is_scripting():
2299
        return ten_crop_image(inpt, size=size, vertical_flip=vertical_flip)
2300
2301
2302
2303
2304

    _log_api_usage_once(ten_crop)

    kernel = _get_kernel(ten_crop, type(inpt))
    return kernel(inpt, size=size, vertical_flip=vertical_flip)
2305
2306


2307
@_register_five_ten_crop_kernel_internal(ten_crop, torch.Tensor)
2308
@_register_five_ten_crop_kernel_internal(ten_crop, tv_tensors.Image)
2309
def ten_crop_image(
Philip Meier's avatar
Philip Meier committed
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
    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,
]:
2323
    non_flipped = five_crop_image(image, size)
2324
2325

    if vertical_flip:
2326
        image = vertical_flip_image(image)
2327
    else:
2328
        image = horizontal_flip_image(image)
2329

2330
    flipped = five_crop_image(image, size)
2331

Philip Meier's avatar
Philip Meier committed
2332
    return non_flipped + flipped
2333
2334


2335
@_register_five_ten_crop_kernel_internal(ten_crop, PIL.Image.Image)
2336
def _ten_crop_image_pil(
Philip Meier's avatar
Philip Meier committed
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
    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,
]:
2350
    non_flipped = _five_crop_image_pil(image, size)
2351
2352

    if vertical_flip:
2353
        image = _vertical_flip_image_pil(image)
2354
    else:
2355
        image = _horizontal_flip_image_pil(image)
2356

2357
    flipped = _five_crop_image_pil(image, size)
Philip Meier's avatar
Philip Meier committed
2358
2359
2360
2361

    return non_flipped + flipped


2362
@_register_five_ten_crop_kernel_internal(ten_crop, tv_tensors.Video)
Philip Meier's avatar
Philip Meier committed
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372
2373
2374
2375
2376
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,
]:
2377
    return ten_crop_image(video, size, vertical_flip=vertical_flip)