_geometry.py 84.8 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)
116
def _vertical_flip_image_pil(image: PIL.Image) -> PIL.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(
Philip Meier's avatar
Philip Meier committed
162
    canvas_size: Tuple[int, int], size: List[int], max_size: Optional[int] = None
163
164
165
) -> List[int]:
    if isinstance(size, int):
        size = [size]
166
167
168
169
170
    elif max_size is not None and len(size) != 1:
        raise ValueError(
            "max_size should only be passed if size specifies the length of the smaller edge, "
            "i.e. size should be an int or a sequence of length 1 in torchscript mode."
        )
Philip Meier's avatar
Philip Meier committed
171
    return __compute_resized_output_size(canvas_size, size=size, max_size=max_size)
172
173


174
def resize(
175
    inpt: torch.Tensor,
176
177
178
    size: List[int],
    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
197
198
199
200
201
202
203
204
# 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:
        if torch._dynamo.is_compiling():
            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
209
    image: torch.Tensor,
    size: 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
314
315
def resize_mask(mask: torch.Tensor, size: List[int], max_size: Optional[int] = None) -> torch.Tensor:
    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(
Philip Meier's avatar
Philip Meier committed
337
    bounding_boxes: torch.Tensor, canvas_size: Tuple[int, int], size: List[int], max_size: Optional[int] = None
338
) -> Tuple[torch.Tensor, Tuple[int, int]]:
Philip Meier's avatar
Philip Meier committed
339
340
    old_height, old_width = canvas_size
    new_height, new_width = _compute_resized_output_size(canvas_size, size=size, max_size=max_size)
341
342

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

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


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


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


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

    _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,
    )
411
412


413
def _affine_parse_args(
414
    angle: Union[int, float],
415
416
417
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
    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}")

457
458
459
460
461
    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]
462
463
464
465

    return angle, translate, shear, center


466
467
468
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
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]:
528
529
530
531
532
533
534
    if torch._dynamo.is_compiling() and not torch.jit.is_scripting():
        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]:
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
    # 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


569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
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


592
def _apply_grid_transform(img: torch.Tensor, grid: torch.Tensor, mode: str, fill: _FillTypeJIT) -> torch.Tensor:
593
594
595
596
597
598
599
600
601
602
    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]
603

604
605
606
607
    # 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)

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

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

634
635
    img = float_img.round_().to(img.dtype) if not fp else float_img

636
    return img.reshape(output_shape)
637
638
639
640
641
642


def _assert_grid_transform_inputs(
    image: torch.Tensor,
    matrix: Optional[List[float]],
    interpolation: str,
643
    fill: _FillTypeJIT,
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
    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)


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

713
714
    angle, translate, shear, center = _affine_parse_args(angle, translate, scale, shear, interpolation, center)

715
716
    height, width = image.shape[-2:]

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

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

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

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


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

755
    return _FP.affine(image, matrix, interpolation=pil_modes_mapping[interpolation], fill=fill)
756
757


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

    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
778
        convert_bounding_box_format(
779
            bounding_boxes, old_format=format, new_format=tv_tensors.BoundingBoxFormat.XYXY, inplace=True
780
781
782
        )
    ).reshape(-1, 4)

783
784
785
    angle, translate, shear, center = _affine_parse_args(
        angle, translate, scale, shear, InterpolationMode.NEAREST, center
    )
786

787
    if center is None:
Philip Meier's avatar
Philip Meier committed
788
        height, width = canvas_size
789
790
        center = [width * 0.5, height * 0.5]

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

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

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

    out_bboxes = out_bboxes.to(original_dtype)
Philip Meier's avatar
Philip Meier committed
844
    return out_bboxes, canvas_size
845
846


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


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


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

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

920
921
922
923
924
    if needs_squeeze:
        output = output.squeeze(0)

    return output

925

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


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


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

984
    _log_api_usage_once(rotate)
985

986
987
988
989
990
    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)
991
@_register_kernel_internal(rotate, tv_tensors.Image)
992
def rotate_image(
993
    image: torch.Tensor,
994
    angle: float,
995
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
996
997
    expand: bool = False,
    center: Optional[List[float]] = None,
998
    fill: _FillTypeJIT = None,
999
) -> torch.Tensor:
1000
1001
    interpolation = _check_interpolation(interpolation)

1002
    input_height, input_width = image.shape[-2:]
1003

1004
1005
    center_f = [0.0, 0.0]
    if center is not None:
1006
        # Center values should be in pixel coordinates but translated such that (0, 0) corresponds to image center.
1007
        center_f = [(c - s * 0.5) for c, s in zip(center, [input_width, input_height])]
1008
1009
1010
1011

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

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

1015
1016
1017
1018
1019
1020
1021
    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)
1022
1023


1024
@_register_kernel_internal(rotate, PIL.Image.Image)
1025
def _rotate_image_pil(
1026
    image: PIL.Image.Image,
1027
    angle: float,
1028
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
1029
1030
    expand: bool = False,
    center: Optional[List[float]] = None,
1031
    fill: _FillTypeJIT = None,
1032
) -> PIL.Image.Image:
1033
1034
    interpolation = _check_interpolation(interpolation)

1035
    return _FP.rotate(
1036
        image, angle, interpolation=pil_modes_mapping[interpolation], expand=expand, fill=fill, center=center
1037
1038
1039
    )


1040
1041
def rotate_bounding_boxes(
    bounding_boxes: torch.Tensor,
1042
    format: tv_tensors.BoundingBoxFormat,
Philip Meier's avatar
Philip Meier committed
1043
    canvas_size: Tuple[int, int],
1044
1045
1046
    angle: float,
    expand: bool = False,
    center: Optional[List[float]] = None,
1047
) -> Tuple[torch.Tensor, Tuple[int, int]]:
1048
1049
    return _affine_bounding_boxes_with_expand(
        bounding_boxes,
1050
        format=format,
Philip Meier's avatar
Philip Meier committed
1051
        canvas_size=canvas_size,
1052
1053
1054
1055
1056
1057
1058
        angle=-angle,
        translate=[0.0, 0.0],
        scale=1.0,
        shear=[0.0, 0.0],
        center=center,
        expand=expand,
    )
1059
1060


1061
@_register_kernel_internal(rotate, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False)
1062
def _rotate_bounding_boxes_dispatch(
1063
1064
    inpt: tv_tensors.BoundingBoxes, angle: float, expand: bool = False, center: Optional[List[float]] = None, **kwargs
) -> tv_tensors.BoundingBoxes:
1065
1066
1067
1068
1069
1070
1071
1072
    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,
    )
1073
    return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size)
1074
1075


1076
1077
def rotate_mask(
    mask: torch.Tensor,
1078
1079
1080
    angle: float,
    expand: bool = False,
    center: Optional[List[float]] = None,
1081
    fill: _FillTypeJIT = None,
1082
) -> torch.Tensor:
1083
1084
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
1085
1086
1087
1088
        needs_squeeze = True
    else:
        needs_squeeze = False

1089
    output = rotate_image(
1090
        mask,
1091
1092
1093
        angle=angle,
        expand=expand,
        interpolation=InterpolationMode.NEAREST,
1094
        fill=fill,
1095
1096
1097
        center=center,
    )

1098
1099
1100
1101
1102
    if needs_squeeze:
        output = output.squeeze(0)

    return output

1103

1104
@_register_kernel_internal(rotate, tv_tensors.Mask, tv_tensor_wrapper=False)
1105
def _rotate_mask_dispatch(
1106
    inpt: tv_tensors.Mask,
1107
1108
1109
    angle: float,
    expand: bool = False,
    center: Optional[List[float]] = None,
1110
    fill: _FillTypeJIT = None,
1111
    **kwargs,
1112
) -> tv_tensors.Mask:
1113
    output = rotate_mask(inpt.as_subclass(torch.Tensor), angle=angle, expand=expand, fill=fill, center=center)
1114
    return tv_tensors.wrap(output, like=inpt)
1115
1116


1117
@_register_kernel_internal(rotate, tv_tensors.Video)
1118
1119
1120
def rotate_video(
    video: torch.Tensor,
    angle: float,
1121
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
1122
1123
    expand: bool = False,
    center: Optional[List[float]] = None,
1124
    fill: _FillTypeJIT = None,
1125
) -> torch.Tensor:
1126
    return rotate_image(video, angle, interpolation=interpolation, expand=expand, fill=fill, center=center)
1127
1128


1129
def pad(
1130
    inpt: torch.Tensor,
1131
1132
1133
    padding: List[int],
    fill: Optional[Union[int, float, List[float]]] = None,
    padding_mode: str = "constant",
1134
) -> torch.Tensor:
1135
    """See :class:`~torchvision.transforms.v2.Pad` for details."""
1136
    if torch.jit.is_scripting():
1137
        return pad_image(inpt, padding=padding, fill=fill, padding_mode=padding_mode)
1138

1139
    _log_api_usage_once(pad)
1140

1141
1142
    kernel = _get_kernel(pad, type(inpt))
    return kernel(inpt, padding=padding, fill=fill, padding_mode=padding_mode)
1143
1144


1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
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]
1167

1168

1169
@_register_kernel_internal(pad, torch.Tensor)
1170
@_register_kernel_internal(pad, tv_tensors.Image)
1171
def pad_image(
1172
    image: torch.Tensor,
1173
1174
    padding: List[int],
    fill: Optional[Union[int, float, List[float]]] = None,
1175
1176
    padding_mode: str = "constant",
) -> torch.Tensor:
1177
    # Be aware that while `padding` has order `[left, top, right, bottom]`, `torch_padding` uses
1178
1179
1180
1181
    # `[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)

1182
    if padding_mode not in ("constant", "edge", "reflect", "symmetric"):
1183
1184
1185
1186
1187
        raise ValueError(
            f"`padding_mode` should be either `'constant'`, `'edge'`, `'reflect'` or `'symmetric'`, "
            f"but got `'{padding_mode}'`."
        )

1188
    if fill is None:
1189
1190
1191
1192
1193
1194
        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)
1195
    else:
1196
        return _pad_with_vector_fill(image, torch_padding, fill=fill, padding_mode=padding_mode)
1197
1198
1199


def _pad_with_scalar_fill(
1200
    image: torch.Tensor,
1201
1202
1203
    torch_padding: List[int],
    fill: Union[int, float],
    padding_mode: str,
1204
) -> torch.Tensor:
1205
1206
    shape = image.shape
    num_channels, height, width = shape[-3:]
1207

1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
    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
1231

1232
1233
1234
1235
1236
        image = torch_pad(image, torch_padding, mode=padding_mode)

        if needs_cast:
            image = image.to(dtype)
    else:  # padding_mode == "symmetric"
1237
        image = _pad_symmetric(image, torch_padding)
1238
1239

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

1241
    return image.reshape(shape[:-3] + (num_channels, new_height, new_width))
1242
1243


1244
# TODO: This should be removed once torch_pad supports non-scalar padding values
1245
def _pad_with_vector_fill(
1246
    image: torch.Tensor,
1247
    torch_padding: List[int],
1248
    fill: List[float],
1249
    padding_mode: str,
1250
1251
1252
1253
) -> torch.Tensor:
    if padding_mode != "constant":
        raise ValueError(f"Padding mode '{padding_mode}' is not supported if fill is not scalar")

1254
1255
    output = _pad_with_scalar_fill(image, torch_padding, fill=0, padding_mode="constant")
    left, right, top, bottom = torch_padding
1256
1257
1258
1259
1260

    # 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)
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272

    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


1273
_pad_image_pil = _register_kernel_internal(pad, PIL.Image.Image)(_FP.pad)
1274
1275


1276
@_register_kernel_internal(pad, tv_tensors.Mask)
1277
1278
def pad_mask(
    mask: torch.Tensor,
1279
1280
    padding: List[int],
    fill: Optional[Union[int, float, List[float]]] = None,
1281
1282
    padding_mode: str = "constant",
) -> torch.Tensor:
1283
1284
1285
    if fill is None:
        fill = 0

1286
    if isinstance(fill, (tuple, list)):
1287
1288
        raise ValueError("Non-scalar fill value is not supported")

1289
1290
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
1291
1292
1293
1294
        needs_squeeze = True
    else:
        needs_squeeze = False

1295
    output = pad_image(mask, padding=padding, fill=fill, padding_mode=padding_mode)
1296
1297
1298
1299
1300

    if needs_squeeze:
        output = output.squeeze(0)

    return output
1301
1302


1303
1304
def pad_bounding_boxes(
    bounding_boxes: torch.Tensor,
1305
    format: tv_tensors.BoundingBoxFormat,
Philip Meier's avatar
Philip Meier committed
1306
    canvas_size: Tuple[int, int],
1307
    padding: List[int],
vfdev's avatar
vfdev committed
1308
    padding_mode: str = "constant",
1309
) -> Tuple[torch.Tensor, Tuple[int, int]]:
vfdev's avatar
vfdev committed
1310
1311
1312
1313
    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")

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

1316
    if format == tv_tensors.BoundingBoxFormat.XYXY:
1317
1318
1319
        pad = [left, top, left, top]
    else:
        pad = [left, top, 0, 0]
1320
    bounding_boxes = bounding_boxes + torch.tensor(pad, dtype=bounding_boxes.dtype, device=bounding_boxes.device)
1321

Philip Meier's avatar
Philip Meier committed
1322
    height, width = canvas_size
1323
1324
    height += top + bottom
    width += left + right
Philip Meier's avatar
Philip Meier committed
1325
    canvas_size = (height, width)
1326

Philip Meier's avatar
Philip Meier committed
1327
    return clamp_bounding_boxes(bounding_boxes, format=format, canvas_size=canvas_size), canvas_size
1328
1329


1330
@_register_kernel_internal(pad, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False)
1331
def _pad_bounding_boxes_dispatch(
1332
1333
    inpt: tv_tensors.BoundingBoxes, padding: List[int], padding_mode: str = "constant", **kwargs
) -> tv_tensors.BoundingBoxes:
1334
1335
1336
1337
1338
1339
1340
    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,
    )
1341
    return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size)
1342
1343


1344
@_register_kernel_internal(pad, tv_tensors.Video)
1345
1346
def pad_video(
    video: torch.Tensor,
1347
1348
    padding: List[int],
    fill: Optional[Union[int, float, List[float]]] = None,
1349
1350
    padding_mode: str = "constant",
) -> torch.Tensor:
1351
    return pad_image(video, padding, fill=fill, padding_mode=padding_mode)
1352
1353


1354
def crop(inpt: torch.Tensor, top: int, left: int, height: int, width: int) -> torch.Tensor:
1355
    """See :class:`~torchvision.transforms.v2.RandomCrop` for details."""
1356
    if torch.jit.is_scripting():
1357
        return crop_image(inpt, top=top, left=left, height=height, width=width)
1358
1359

    _log_api_usage_once(crop)
1360

1361
1362
    kernel = _get_kernel(crop, type(inpt))
    return kernel(inpt, top=top, left=left, height=height, width=width)
1363

1364
1365

@_register_kernel_internal(crop, torch.Tensor)
1366
@_register_kernel_internal(crop, tv_tensors.Image)
1367
def crop_image(image: torch.Tensor, top: int, left: int, height: int, width: int) -> torch.Tensor:
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
    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]


1385
1386
_crop_image_pil = _FP.crop
_register_kernel_internal(crop, PIL.Image.Image)(_crop_image_pil)
1387
1388


1389
1390
def crop_bounding_boxes(
    bounding_boxes: torch.Tensor,
1391
    format: tv_tensors.BoundingBoxFormat,
1392
1393
    top: int,
    left: int,
1394
1395
1396
    height: int,
    width: int,
) -> Tuple[torch.Tensor, Tuple[int, int]]:
1397

1398
    # Crop or implicit pad if left and/or top have negative values:
1399
    if format == tv_tensors.BoundingBoxFormat.XYXY:
1400
        sub = [left, top, left, top]
1401
    else:
1402
1403
        sub = [left, top, 0, 0]

1404
    bounding_boxes = bounding_boxes - torch.tensor(sub, dtype=bounding_boxes.dtype, device=bounding_boxes.device)
Philip Meier's avatar
Philip Meier committed
1405
    canvas_size = (height, width)
1406

Philip Meier's avatar
Philip Meier committed
1407
    return clamp_bounding_boxes(bounding_boxes, format=format, canvas_size=canvas_size), canvas_size
1408
1409


1410
@_register_kernel_internal(crop, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False)
1411
def _crop_bounding_boxes_dispatch(
1412
1413
    inpt: tv_tensors.BoundingBoxes, top: int, left: int, height: int, width: int
) -> tv_tensors.BoundingBoxes:
1414
1415
1416
    output, canvas_size = crop_bounding_boxes(
        inpt.as_subclass(torch.Tensor), format=inpt.format, top=top, left=left, height=height, width=width
    )
1417
    return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size)
1418
1419


1420
@_register_kernel_internal(crop, tv_tensors.Mask)
1421
def crop_mask(mask: torch.Tensor, top: int, left: int, height: int, width: int) -> torch.Tensor:
1422
1423
1424
1425
1426
1427
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
        needs_squeeze = True
    else:
        needs_squeeze = False

1428
    output = crop_image(mask, top, left, height, width)
1429
1430
1431
1432
1433

    if needs_squeeze:
        output = output.squeeze(0)

    return output
1434
1435


1436
@_register_kernel_internal(crop, tv_tensors.Video)
1437
def crop_video(video: torch.Tensor, top: int, left: int, height: int, width: int) -> torch.Tensor:
1438
    return crop_image(video, top, left, height, width)
1439
1440


1441
def perspective(
1442
    inpt: torch.Tensor,
1443
1444
1445
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1446
    fill: _FillTypeJIT = None,
1447
    coefficients: Optional[List[float]] = None,
1448
) -> torch.Tensor:
1449
    """See :class:`~torchvision.transforms.v2.RandomPerspective` for details."""
1450
    if torch.jit.is_scripting():
1451
        return perspective_image(
1452
1453
1454
1455
1456
1457
            inpt,
            startpoints=startpoints,
            endpoints=endpoints,
            interpolation=interpolation,
            fill=fill,
            coefficients=coefficients,
1458
        )
1459

1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
    _log_api_usage_once(perspective)

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

1472

1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
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)
1488
    x_grid = torch.linspace(d, ow + d - 1.0, steps=ow, device=device, dtype=dtype)
1489
    base_grid[..., 0].copy_(x_grid)
1490
    y_grid = torch.linspace(d, oh + d - 1.0, steps=oh, device=device, dtype=dtype).unsqueeze_(-1)
1491
1492
1493
1494
    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))
1495
1496
1497
    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))
1498
1499
1500
1501
1502

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


1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
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.")


1520
@_register_kernel_internal(perspective, torch.Tensor)
1521
@_register_kernel_internal(perspective, tv_tensors.Image)
1522
def perspective_image(
1523
    image: torch.Tensor,
1524
1525
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
1526
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1527
    fill: _FillTypeJIT = None,
1528
    coefficients: Optional[List[float]] = None,
1529
) -> torch.Tensor:
1530
    perspective_coeffs = _perspective_coefficients(startpoints, endpoints, coefficients)
1531
1532
    interpolation = _check_interpolation(interpolation)

1533
    _assert_grid_transform_inputs(
1534
1535
1536
1537
1538
1539
1540
1541
        image,
        matrix=None,
        interpolation=interpolation.value,
        fill=fill,
        supported_interpolation_modes=["nearest", "bilinear"],
        coeffs=perspective_coeffs,
    )

1542
    oh, ow = image.shape[-2:]
1543
    dtype = image.dtype if torch.is_floating_point(image) else torch.float32
1544
    grid = _perspective_grid(perspective_coeffs, ow=ow, oh=oh, dtype=dtype, device=image.device)
1545
    return _apply_grid_transform(image, grid, interpolation.value, fill=fill)
1546
1547


1548
@_register_kernel_internal(perspective, PIL.Image.Image)
1549
def _perspective_image_pil(
1550
    image: PIL.Image.Image,
1551
1552
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
1553
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1554
    fill: _FillTypeJIT = None,
1555
    coefficients: Optional[List[float]] = None,
1556
) -> PIL.Image.Image:
1557
    perspective_coeffs = _perspective_coefficients(startpoints, endpoints, coefficients)
1558
    interpolation = _check_interpolation(interpolation)
1559
    return _FP.perspective(image, perspective_coeffs, interpolation=pil_modes_mapping[interpolation], fill=fill)
1560
1561


1562
1563
def perspective_bounding_boxes(
    bounding_boxes: torch.Tensor,
1564
    format: tv_tensors.BoundingBoxFormat,
Philip Meier's avatar
Philip Meier committed
1565
    canvas_size: Tuple[int, int],
1566
1567
1568
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
    coefficients: Optional[List[float]] = None,
1569
) -> torch.Tensor:
1570
1571
    if bounding_boxes.numel() == 0:
        return bounding_boxes
1572

1573
    perspective_coeffs = _perspective_coefficients(startpoints, endpoints, coefficients)
1574

1575
    original_shape = bounding_boxes.shape
Nicolas Hug's avatar
Nicolas Hug committed
1576
    # TODO: first cast to float if bbox is int64 before convert_bounding_box_format
1577
    bounding_boxes = (
1578
        convert_bounding_box_format(bounding_boxes, old_format=format, new_format=tv_tensors.BoundingBoxFormat.XYXY)
1579
    ).reshape(-1, 4)
1580

1581
1582
    dtype = bounding_boxes.dtype if torch.is_floating_point(bounding_boxes) else torch.float32
    device = bounding_boxes.device
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613

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

1614
1615
    theta1 = torch.tensor(
        [[inv_coeffs[0], inv_coeffs[1], inv_coeffs[2]], [inv_coeffs[3], inv_coeffs[4], inv_coeffs[5]]],
1616
1617
1618
1619
        dtype=dtype,
        device=device,
    )

1620
1621
1622
1623
    theta2 = torch.tensor(
        [[inv_coeffs[6], inv_coeffs[7], 1.0], [inv_coeffs[6], inv_coeffs[7], 1.0]], dtype=dtype, device=device
    )

1624
1625
1626
1627
    # 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)]
1628
    points = bounding_boxes[:, [[0, 1], [2, 1], [2, 3], [0, 3]]].reshape(-1, 2)
1629
1630
1631
1632
1633
    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)

1634
1635
    numer_points = torch.matmul(points, theta1.T)
    denom_points = torch.matmul(points, theta2.T)
1636
    transformed_points = numer_points.div_(denom_points)
1637
1638
1639

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

1643
1644
    out_bboxes = clamp_bounding_boxes(
        torch.cat([out_bbox_mins, out_bbox_maxs], dim=1).to(bounding_boxes.dtype),
1645
        format=tv_tensors.BoundingBoxFormat.XYXY,
Philip Meier's avatar
Philip Meier committed
1646
        canvas_size=canvas_size,
1647
    )
1648
1649
1650

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

Nicolas Hug's avatar
Nicolas Hug committed
1651
    return convert_bounding_box_format(
1652
        out_bboxes, old_format=tv_tensors.BoundingBoxFormat.XYXY, new_format=format, inplace=True
1653
    ).reshape(original_shape)
1654
1655


1656
@_register_kernel_internal(perspective, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False)
1657
def _perspective_bounding_boxes_dispatch(
1658
    inpt: tv_tensors.BoundingBoxes,
1659
1660
1661
1662
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
    coefficients: Optional[List[float]] = None,
    **kwargs,
1663
) -> tv_tensors.BoundingBoxes:
1664
1665
1666
1667
1668
1669
1670
1671
    output = perspective_bounding_boxes(
        inpt.as_subclass(torch.Tensor),
        format=inpt.format,
        canvas_size=inpt.canvas_size,
        startpoints=startpoints,
        endpoints=endpoints,
        coefficients=coefficients,
    )
1672
    return tv_tensors.wrap(output, like=inpt)
1673
1674


1675
1676
def perspective_mask(
    mask: torch.Tensor,
1677
1678
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
1679
    fill: _FillTypeJIT = None,
1680
    coefficients: Optional[List[float]] = None,
1681
) -> torch.Tensor:
1682
1683
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
1684
1685
1686
1687
        needs_squeeze = True
    else:
        needs_squeeze = False

1688
    output = perspective_image(
1689
        mask, startpoints, endpoints, interpolation=InterpolationMode.NEAREST, fill=fill, coefficients=coefficients
1690
    )
1691

1692
1693
1694
1695
1696
    if needs_squeeze:
        output = output.squeeze(0)

    return output

1697

1698
@_register_kernel_internal(perspective, tv_tensors.Mask, tv_tensor_wrapper=False)
1699
def _perspective_mask_dispatch(
1700
    inpt: tv_tensors.Mask,
1701
1702
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
1703
    fill: _FillTypeJIT = None,
1704
1705
    coefficients: Optional[List[float]] = None,
    **kwargs,
1706
) -> tv_tensors.Mask:
1707
1708
1709
1710
1711
1712
1713
    output = perspective_mask(
        inpt.as_subclass(torch.Tensor),
        startpoints=startpoints,
        endpoints=endpoints,
        fill=fill,
        coefficients=coefficients,
    )
1714
    return tv_tensors.wrap(output, like=inpt)
1715
1716


1717
@_register_kernel_internal(perspective, tv_tensors.Video)
1718
1719
def perspective_video(
    video: torch.Tensor,
1720
1721
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
1722
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1723
    fill: _FillTypeJIT = None,
1724
    coefficients: Optional[List[float]] = None,
1725
) -> torch.Tensor:
1726
    return perspective_image(
1727
1728
        video, startpoints, endpoints, interpolation=interpolation, fill=fill, coefficients=coefficients
    )
1729
1730


1731
def elastic(
1732
    inpt: torch.Tensor,
1733
    displacement: torch.Tensor,
1734
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1735
1736
    fill: _FillTypeJIT = None,
) -> torch.Tensor:
1737
    """See :class:`~torchvision.transforms.v2.ElasticTransform` for details."""
1738
    if torch.jit.is_scripting():
1739
        return elastic_image(inpt, displacement=displacement, interpolation=interpolation, fill=fill)
1740
1741
1742
1743
1744

    _log_api_usage_once(elastic)

    kernel = _get_kernel(elastic, type(inpt))
    return kernel(inpt, displacement=displacement, interpolation=interpolation, fill=fill)
1745
1746


1747
1748
1749
elastic_transform = elastic


1750
@_register_kernel_internal(elastic, torch.Tensor)
1751
@_register_kernel_internal(elastic, tv_tensors.Image)
1752
def elastic_image(
1753
    image: torch.Tensor,
1754
    displacement: torch.Tensor,
1755
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1756
    fill: _FillTypeJIT = None,
1757
) -> torch.Tensor:
Philip Meier's avatar
Philip Meier committed
1758
1759
1760
    if not isinstance(displacement, torch.Tensor):
        raise TypeError("Argument displacement should be a Tensor")

1761
1762
    interpolation = _check_interpolation(interpolation)

1763
    height, width = image.shape[-2:]
1764
    device = image.device
1765
    dtype = image.dtype if torch.is_floating_point(image) else torch.float32
1766
1767
1768
1769
1770
1771
1772

    # 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

1773
    # We are aware that if input image dtype is uint8 and displacement is float64 then
1774
    # displacement will be cast to float32 and all computations will be done with float32
1775
    # We can fix this later if needed
1776

1777
    expected_shape = (1, height, width, 2)
1778
1779
1780
    if expected_shape != displacement.shape:
        raise ValueError(f"Argument displacement shape should be {expected_shape}, but given {displacement.shape}")

1781
1782
1783
    grid = _create_identity_grid((height, width), device=device, dtype=dtype).add_(
        displacement.to(dtype=dtype, device=device)
    )
1784
    output = _apply_grid_transform(image, grid, interpolation.value, fill=fill)
1785

1786
1787
1788
    if is_cpu_half:
        output = output.to(torch.float16)

1789
    return output
1790
1791


1792
@_register_kernel_internal(elastic, PIL.Image.Image)
1793
def _elastic_image_pil(
1794
    image: PIL.Image.Image,
1795
    displacement: torch.Tensor,
1796
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1797
    fill: _FillTypeJIT = None,
1798
) -> PIL.Image.Image:
1799
    t_img = pil_to_tensor(image)
1800
    output = elastic_image(t_img, displacement, interpolation=interpolation, fill=fill)
1801
    return to_pil_image(output, mode=image.mode)
1802
1803


1804
def _create_identity_grid(size: Tuple[int, int], device: torch.device, dtype: torch.dtype) -> torch.Tensor:
1805
    sy, sx = size
1806
1807
    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)
1808
1809
    base_grid[..., 0].copy_(x_grid)

1810
    y_grid = torch.linspace((-sy + 1) / sy, (sy - 1) / sy, sy, device=device, dtype=dtype).unsqueeze_(-1)
1811
1812
1813
1814
1815
    base_grid[..., 1].copy_(y_grid)

    return base_grid


1816
1817
def elastic_bounding_boxes(
    bounding_boxes: torch.Tensor,
1818
    format: tv_tensors.BoundingBoxFormat,
Philip Meier's avatar
Philip Meier committed
1819
    canvas_size: Tuple[int, int],
1820
1821
    displacement: torch.Tensor,
) -> torch.Tensor:
Philip Meier's avatar
Philip Meier committed
1822
1823
1824
1825
1826
1827
    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}")

1828
1829
    if bounding_boxes.numel() == 0:
        return bounding_boxes
1830

1831
    # TODO: add in docstring about approximation we are doing for grid inversion
1832
1833
    device = bounding_boxes.device
    dtype = bounding_boxes.dtype if torch.is_floating_point(bounding_boxes) else torch.float32
1834
1835
1836

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

1838
    original_shape = bounding_boxes.shape
Nicolas Hug's avatar
Nicolas Hug committed
1839
    # TODO: first cast to float if bbox is int64 before convert_bounding_box_format
1840
    bounding_boxes = (
1841
        convert_bounding_box_format(bounding_boxes, old_format=format, new_format=tv_tensors.BoundingBoxFormat.XYXY)
1842
    ).reshape(-1, 4)
1843

Philip Meier's avatar
Philip Meier committed
1844
    id_grid = _create_identity_grid(canvas_size, device=device, dtype=dtype)
1845
1846
    # We construct an approximation of inverse grid as inv_grid = id_grid - displacement
    # This is not an exact inverse of the grid
1847
    inv_grid = id_grid.sub_(displacement)
1848
1849

    # Get points from bboxes
1850
    points = bounding_boxes[:, [[0, 1], [2, 1], [2, 3], [0, 3]]].reshape(-1, 2)
1851
1852
1853
1854
1855
    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]

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

1860
    transformed_points = transformed_points.reshape(-1, 4, 2)
1861
    out_bbox_mins, out_bbox_maxs = torch.aminmax(transformed_points, dim=1)
1862
1863
    out_bboxes = clamp_bounding_boxes(
        torch.cat([out_bbox_mins, out_bbox_maxs], dim=1).to(bounding_boxes.dtype),
1864
        format=tv_tensors.BoundingBoxFormat.XYXY,
Philip Meier's avatar
Philip Meier committed
1865
        canvas_size=canvas_size,
1866
    )
1867

Nicolas Hug's avatar
Nicolas Hug committed
1868
    return convert_bounding_box_format(
1869
        out_bboxes, old_format=tv_tensors.BoundingBoxFormat.XYXY, new_format=format, inplace=True
1870
    ).reshape(original_shape)
1871
1872


1873
@_register_kernel_internal(elastic, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False)
1874
def _elastic_bounding_boxes_dispatch(
1875
1876
    inpt: tv_tensors.BoundingBoxes, displacement: torch.Tensor, **kwargs
) -> tv_tensors.BoundingBoxes:
1877
1878
1879
    output = elastic_bounding_boxes(
        inpt.as_subclass(torch.Tensor), format=inpt.format, canvas_size=inpt.canvas_size, displacement=displacement
    )
1880
    return tv_tensors.wrap(output, like=inpt)
1881
1882


1883
1884
1885
def elastic_mask(
    mask: torch.Tensor,
    displacement: torch.Tensor,
1886
    fill: _FillTypeJIT = None,
1887
) -> torch.Tensor:
1888
1889
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
1890
1891
1892
1893
        needs_squeeze = True
    else:
        needs_squeeze = False

1894
    output = elastic_image(mask, displacement=displacement, interpolation=InterpolationMode.NEAREST, fill=fill)
1895
1896
1897
1898
1899

    if needs_squeeze:
        output = output.squeeze(0)

    return output
1900
1901


1902
@_register_kernel_internal(elastic, tv_tensors.Mask, tv_tensor_wrapper=False)
1903
def _elastic_mask_dispatch(
1904
1905
    inpt: tv_tensors.Mask, displacement: torch.Tensor, fill: _FillTypeJIT = None, **kwargs
) -> tv_tensors.Mask:
1906
    output = elastic_mask(inpt.as_subclass(torch.Tensor), displacement=displacement, fill=fill)
1907
    return tv_tensors.wrap(output, like=inpt)
1908
1909


1910
@_register_kernel_internal(elastic, tv_tensors.Video)
1911
1912
1913
def elastic_video(
    video: torch.Tensor,
    displacement: torch.Tensor,
1914
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1915
    fill: _FillTypeJIT = None,
1916
) -> torch.Tensor:
1917
    return elastic_image(video, displacement, interpolation=interpolation, fill=fill)
1918
1919


1920
def center_crop(inpt: torch.Tensor, output_size: List[int]) -> torch.Tensor:
1921
    """See :class:`~torchvision.transforms.v2.RandomCrop` for details."""
1922
    if torch.jit.is_scripting():
1923
        return center_crop_image(inpt, output_size=output_size)
1924
1925
1926
1927
1928

    _log_api_usage_once(center_crop)

    kernel = _get_kernel(center_crop, type(inpt))
    return kernel(inpt, output_size=output_size)
1929
1930


1931
1932
def _center_crop_parse_output_size(output_size: List[int]) -> List[int]:
    if isinstance(output_size, numbers.Number):
1933
1934
        s = int(output_size)
        return [s, s]
1935
    elif isinstance(output_size, (tuple, list)) and len(output_size) == 1:
1936
        return [output_size[0], output_size[0]]
1937
1938
    else:
        return list(output_size)
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957


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


1958
@_register_kernel_internal(center_crop, torch.Tensor)
1959
@_register_kernel_internal(center_crop, tv_tensors.Image)
1960
def center_crop_image(image: torch.Tensor, output_size: List[int]) -> torch.Tensor:
1961
    crop_height, crop_width = _center_crop_parse_output_size(output_size)
1962
1963
1964
1965
    shape = image.shape
    if image.numel() == 0:
        return image.reshape(shape[:-2] + (crop_height, crop_width))
    image_height, image_width = shape[-2:]
1966
1967
1968

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

1971
        image_height, image_width = image.shape[-2:]
1972
        if crop_width == image_width and crop_height == image_height:
1973
            return image
1974
1975

    crop_top, crop_left = _center_crop_compute_crop_anchor(crop_height, crop_width, image_height, image_width)
1976
    return image[..., crop_top : (crop_top + crop_height), crop_left : (crop_left + crop_width)]
1977
1978


1979
@_register_kernel_internal(center_crop, PIL.Image.Image)
1980
def _center_crop_image_pil(image: PIL.Image.Image, output_size: List[int]) -> PIL.Image.Image:
1981
    crop_height, crop_width = _center_crop_parse_output_size(output_size)
1982
    image_height, image_width = _get_size_image_pil(image)
1983
1984
1985

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

1988
        image_height, image_width = _get_size_image_pil(image)
1989
        if crop_width == image_width and crop_height == image_height:
1990
            return image
1991
1992

    crop_top, crop_left = _center_crop_compute_crop_anchor(crop_height, crop_width, image_height, image_width)
1993
    return _crop_image_pil(image, crop_top, crop_left, crop_height, crop_width)
1994
1995


1996
1997
def center_crop_bounding_boxes(
    bounding_boxes: torch.Tensor,
1998
    format: tv_tensors.BoundingBoxFormat,
Philip Meier's avatar
Philip Meier committed
1999
    canvas_size: Tuple[int, int],
2000
    output_size: List[int],
2001
) -> Tuple[torch.Tensor, Tuple[int, int]]:
2002
    crop_height, crop_width = _center_crop_parse_output_size(output_size)
Philip Meier's avatar
Philip Meier committed
2003
    crop_top, crop_left = _center_crop_compute_crop_anchor(crop_height, crop_width, *canvas_size)
2004
2005
2006
    return crop_bounding_boxes(
        bounding_boxes, format, top=crop_top, left=crop_left, height=crop_height, width=crop_width
    )
2007
2008


2009
@_register_kernel_internal(center_crop, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False)
2010
def _center_crop_bounding_boxes_dispatch(
2011
2012
    inpt: tv_tensors.BoundingBoxes, output_size: List[int]
) -> tv_tensors.BoundingBoxes:
2013
2014
2015
    output, canvas_size = center_crop_bounding_boxes(
        inpt.as_subclass(torch.Tensor), format=inpt.format, canvas_size=inpt.canvas_size, output_size=output_size
    )
2016
    return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size)
2017
2018


2019
@_register_kernel_internal(center_crop, tv_tensors.Mask)
2020
2021
2022
def center_crop_mask(mask: torch.Tensor, output_size: List[int]) -> torch.Tensor:
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
2023
2024
2025
2026
        needs_squeeze = True
    else:
        needs_squeeze = False

2027
    output = center_crop_image(image=mask, output_size=output_size)
2028
2029
2030
2031
2032

    if needs_squeeze:
        output = output.squeeze(0)

    return output
2033
2034


2035
@_register_kernel_internal(center_crop, tv_tensors.Video)
2036
def center_crop_video(video: torch.Tensor, output_size: List[int]) -> torch.Tensor:
2037
    return center_crop_image(video, output_size)
2038
2039


2040
def resized_crop(
2041
    inpt: torch.Tensor,
2042
2043
2044
2045
2046
2047
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
2048
    antialias: Optional[bool] = True,
2049
) -> torch.Tensor:
2050
    """See :class:`~torchvision.transforms.v2.RandomResizedCrop` for details."""
2051
    if torch.jit.is_scripting():
2052
        return resized_crop_image(
2053
2054
2055
2056
2057
2058
2059
2060
            inpt,
            top=top,
            left=left,
            height=height,
            width=width,
            size=size,
            interpolation=interpolation,
            antialias=antialias,
2061
        )
2062

2063
2064
2065
2066
2067
2068
2069
2070
2071
2072
2073
2074
2075
    _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,
    )
2076

2077
2078

@_register_kernel_internal(resized_crop, torch.Tensor)
2079
@_register_kernel_internal(resized_crop, tv_tensors.Image)
2080
def resized_crop_image(
2081
    image: torch.Tensor,
2082
2083
2084
2085
2086
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
2087
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
2088
    antialias: Optional[bool] = True,
2089
) -> torch.Tensor:
2090
2091
    image = crop_image(image, top, left, height, width)
    return resize_image(image, size, interpolation=interpolation, antialias=antialias)
2092
2093


2094
def _resized_crop_image_pil(
2095
    image: PIL.Image.Image,
2096
2097
2098
2099
2100
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
2101
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
2102
) -> PIL.Image.Image:
2103
2104
    image = _crop_image_pil(image, top, left, height, width)
    return _resize_image_pil(image, size, interpolation=interpolation)
2105
2106


2107
@_register_kernel_internal(resized_crop, PIL.Image.Image)
2108
def _resized_crop_image_pil_dispatch(
2109
2110
2111
2112
2113
2114
2115
    image: PIL.Image.Image,
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
2116
    antialias: Optional[bool] = True,
2117
2118
2119
) -> PIL.Image.Image:
    if antialias is False:
        warnings.warn("Anti-alias option is always applied for PIL Image input. Argument antialias is ignored.")
2120
    return _resized_crop_image_pil(
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
        image,
        top=top,
        left=left,
        height=height,
        width=width,
        size=size,
        interpolation=interpolation,
    )


2131
2132
def resized_crop_bounding_boxes(
    bounding_boxes: torch.Tensor,
2133
    format: tv_tensors.BoundingBoxFormat,
2134
2135
2136
2137
2138
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
2139
) -> Tuple[torch.Tensor, Tuple[int, int]]:
2140
2141
2142
2143
    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)


2144
@_register_kernel_internal(resized_crop, tv_tensors.BoundingBoxes, tv_tensor_wrapper=False)
2145
def _resized_crop_bounding_boxes_dispatch(
2146
2147
    inpt: tv_tensors.BoundingBoxes, top: int, left: int, height: int, width: int, size: List[int], **kwargs
) -> tv_tensors.BoundingBoxes:
2148
2149
2150
    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
    )
2151
    return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size)
2152
2153


2154
def resized_crop_mask(
2155
2156
2157
2158
2159
2160
2161
    mask: torch.Tensor,
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
) -> torch.Tensor:
2162
2163
    mask = crop_mask(mask, top, left, height, width)
    return resize_mask(mask, size)
2164
2165


2166
@_register_kernel_internal(resized_crop, tv_tensors.Mask, tv_tensor_wrapper=False)
2167
def _resized_crop_mask_dispatch(
2168
2169
    inpt: tv_tensors.Mask, top: int, left: int, height: int, width: int, size: List[int], **kwargs
) -> tv_tensors.Mask:
2170
2171
2172
    output = resized_crop_mask(
        inpt.as_subclass(torch.Tensor), top=top, left=left, height=height, width=width, size=size
    )
2173
    return tv_tensors.wrap(output, like=inpt)
2174
2175


2176
@_register_kernel_internal(resized_crop, tv_tensors.Video)
2177
2178
2179
2180
2181
2182
2183
def resized_crop_video(
    video: torch.Tensor,
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
2184
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
2185
    antialias: Optional[bool] = True,
2186
) -> torch.Tensor:
2187
    return resized_crop_image(
2188
2189
2190
2191
        video, top, left, height, width, antialias=antialias, size=size, interpolation=interpolation
    )


2192
def five_crop(
2193
2194
    inpt: torch.Tensor, size: List[int]
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
2195
    """See :class:`~torchvision.transforms.v2.FiveCrop` for details."""
2196
    if torch.jit.is_scripting():
2197
        return five_crop_image(inpt, size=size)
2198
2199
2200
2201
2202

    _log_api_usage_once(five_crop)

    kernel = _get_kernel(five_crop, type(inpt))
    return kernel(inpt, size=size)
2203
2204


2205
2206
def _parse_five_crop_size(size: List[int]) -> List[int]:
    if isinstance(size, numbers.Number):
2207
2208
        s = int(size)
        size = [s, s]
2209
    elif isinstance(size, (tuple, list)) and len(size) == 1:
2210
2211
        s = size[0]
        size = [s, s]
2212
2213
2214
2215
2216
2217
2218

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

    return size


2219
@_register_five_ten_crop_kernel_internal(five_crop, torch.Tensor)
2220
@_register_five_ten_crop_kernel_internal(five_crop, tv_tensors.Image)
2221
def five_crop_image(
2222
    image: torch.Tensor, size: List[int]
2223
2224
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    crop_height, crop_width = _parse_five_crop_size(size)
2225
    image_height, image_width = image.shape[-2:]
2226
2227

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

2230
2231
2232
2233
2234
    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])
2235
2236
2237
2238

    return tl, tr, bl, br, center


2239
@_register_five_ten_crop_kernel_internal(five_crop, PIL.Image.Image)
2240
def _five_crop_image_pil(
2241
    image: PIL.Image.Image, size: List[int]
2242
2243
) -> 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)
2244
    image_height, image_width = _get_size_image_pil(image)
2245
2246

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

2249
2250
2251
2252
2253
    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])
2254
2255
2256
2257

    return tl, tr, bl, br, center


2258
@_register_five_ten_crop_kernel_internal(five_crop, tv_tensors.Video)
2259
2260
2261
def five_crop_video(
    video: torch.Tensor, size: List[int]
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
2262
    return five_crop_image(video, size)
2263
2264


2265
def ten_crop(
2266
    inpt: torch.Tensor, size: List[int], vertical_flip: bool = False
2267
) -> Tuple[
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
2278
]:
2279
    """See :class:`~torchvision.transforms.v2.TenCrop` for details."""
2280
    if torch.jit.is_scripting():
2281
        return ten_crop_image(inpt, size=size, vertical_flip=vertical_flip)
2282
2283
2284
2285
2286

    _log_api_usage_once(ten_crop)

    kernel = _get_kernel(ten_crop, type(inpt))
    return kernel(inpt, size=size, vertical_flip=vertical_flip)
2287
2288


2289
@_register_five_ten_crop_kernel_internal(ten_crop, torch.Tensor)
2290
@_register_five_ten_crop_kernel_internal(ten_crop, tv_tensors.Image)
2291
def ten_crop_image(
Philip Meier's avatar
Philip Meier committed
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
    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,
]:
2305
    non_flipped = five_crop_image(image, size)
2306
2307

    if vertical_flip:
2308
        image = vertical_flip_image(image)
2309
    else:
2310
        image = horizontal_flip_image(image)
2311

2312
    flipped = five_crop_image(image, size)
2313

Philip Meier's avatar
Philip Meier committed
2314
    return non_flipped + flipped
2315
2316


2317
@_register_five_ten_crop_kernel_internal(ten_crop, PIL.Image.Image)
2318
def _ten_crop_image_pil(
Philip Meier's avatar
Philip Meier committed
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
    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,
]:
2332
    non_flipped = _five_crop_image_pil(image, size)
2333
2334

    if vertical_flip:
2335
        image = _vertical_flip_image_pil(image)
2336
    else:
2337
        image = _horizontal_flip_image_pil(image)
2338

2339
    flipped = _five_crop_image_pil(image, size)
Philip Meier's avatar
Philip Meier committed
2340
2341
2342
2343

    return non_flipped + flipped


2344
@_register_five_ten_crop_kernel_internal(ten_crop, tv_tensors.Video)
Philip Meier's avatar
Philip Meier committed
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
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,
]:
2359
    return ten_crop_image(video, size, vertical_flip=vertical_flip)