"docs/GETTING_STARTED.md" did not exist on "d00d0be15a201a2f044d62229f73a65dd02c0a4b"
_geometry.py 75.4 KB
Newer Older
1
import math
2
import numbers
3
import warnings
4
from typing import List, Optional, Sequence, Tuple, Union
5
6
7

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

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

24
25
from torchvision.utils import _log_api_usage_once

26
from ._meta import clamp_bounding_box, convert_format_bounding_box, get_spatial_size_image_pil
27

28
29
from ._utils import is_simple_tensor

30

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


42
43
44
45
def horizontal_flip_image_tensor(image: torch.Tensor) -> torch.Tensor:
    return image.flip(-1)


46
47
48
horizontal_flip_image_pil = _FP.hflip


49
50
def horizontal_flip_mask(mask: torch.Tensor) -> torch.Tensor:
    return horizontal_flip_image_tensor(mask)
51
52


53
def horizontal_flip_bounding_box(
54
    bounding_box: torch.Tensor, format: datapoints.BoundingBoxFormat, spatial_size: Tuple[int, int]
55
56
57
) -> torch.Tensor:
    shape = bounding_box.shape

58
    bounding_box = bounding_box.clone().reshape(-1, 4)
59

60
    if format == datapoints.BoundingBoxFormat.XYXY:
61
        bounding_box[:, [2, 0]] = bounding_box[:, [0, 2]].sub_(spatial_size[1]).neg_()
62
    elif format == datapoints.BoundingBoxFormat.XYWH:
63
        bounding_box[:, 0].add_(bounding_box[:, 2]).sub_(spatial_size[1]).neg_()
64
    else:  # format == datapoints.BoundingBoxFormat.CXCYWH:
65
        bounding_box[:, 0].sub_(spatial_size[1]).neg_()
66

67
    return bounding_box.reshape(shape)
68
69


70
71
72
73
def horizontal_flip_video(video: torch.Tensor) -> torch.Tensor:
    return horizontal_flip_image_tensor(video)


Philip Meier's avatar
Philip Meier committed
74
def horizontal_flip(inpt: datapoints._InputTypeJIT) -> datapoints._InputTypeJIT:
75
76
77
    if not torch.jit.is_scripting():
        _log_api_usage_once(horizontal_flip)

78
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
79
        return horizontal_flip_image_tensor(inpt)
80
    elif isinstance(inpt, datapoints._datapoint.Datapoint):
81
        return inpt.horizontal_flip()
82
    elif isinstance(inpt, PIL.Image.Image):
83
        return horizontal_flip_image_pil(inpt)
84
85
    else:
        raise TypeError(
86
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
87
88
            f"but got {type(inpt)} instead."
        )
89
90


91
92
93
94
def vertical_flip_image_tensor(image: torch.Tensor) -> torch.Tensor:
    return image.flip(-2)


95
96
97
vertical_flip_image_pil = _FP.vflip


98
99
def vertical_flip_mask(mask: torch.Tensor) -> torch.Tensor:
    return vertical_flip_image_tensor(mask)
100
101
102


def vertical_flip_bounding_box(
103
    bounding_box: torch.Tensor, format: datapoints.BoundingBoxFormat, spatial_size: Tuple[int, int]
104
105
106
) -> torch.Tensor:
    shape = bounding_box.shape

107
    bounding_box = bounding_box.clone().reshape(-1, 4)
108

109
    if format == datapoints.BoundingBoxFormat.XYXY:
110
        bounding_box[:, [1, 3]] = bounding_box[:, [3, 1]].sub_(spatial_size[0]).neg_()
111
    elif format == datapoints.BoundingBoxFormat.XYWH:
112
        bounding_box[:, 1].add_(bounding_box[:, 3]).sub_(spatial_size[0]).neg_()
113
    else:  # format == datapoints.BoundingBoxFormat.CXCYWH:
114
        bounding_box[:, 1].sub_(spatial_size[0]).neg_()
115

116
    return bounding_box.reshape(shape)
117
118


119
120
121
122
def vertical_flip_video(video: torch.Tensor) -> torch.Tensor:
    return vertical_flip_image_tensor(video)


Philip Meier's avatar
Philip Meier committed
123
def vertical_flip(inpt: datapoints._InputTypeJIT) -> datapoints._InputTypeJIT:
124
125
126
    if not torch.jit.is_scripting():
        _log_api_usage_once(vertical_flip)

127
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
128
        return vertical_flip_image_tensor(inpt)
129
    elif isinstance(inpt, datapoints._datapoint.Datapoint):
130
        return inpt.vertical_flip()
131
    elif isinstance(inpt, PIL.Image.Image):
132
        return vertical_flip_image_pil(inpt)
133
134
    else:
        raise TypeError(
135
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
136
137
            f"but got {type(inpt)} instead."
        )
138
139


140
141
142
143
144
145
# 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


146
def _compute_resized_output_size(
147
    spatial_size: Tuple[int, int], size: List[int], max_size: Optional[int] = None
148
149
150
) -> List[int]:
    if isinstance(size, int):
        size = [size]
151
152
153
154
155
    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."
        )
156
    return __compute_resized_output_size(spatial_size, size=size, max_size=max_size)
157
158


159
160
161
def resize_image_tensor(
    image: torch.Tensor,
    size: List[int],
162
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
163
    max_size: Optional[int] = None,
164
    antialias: Optional[Union[str, bool]] = "warn",
165
) -> torch.Tensor:
166
    interpolation = _check_interpolation(interpolation)
167
168
    antialias = _check_antialias(img=image, antialias=antialias, interpolation=interpolation)
    assert not isinstance(antialias, str)
169
    antialias = False if antialias is None else antialias
170
171
172
    align_corners: Optional[bool] = None
    if interpolation == InterpolationMode.BILINEAR or interpolation == InterpolationMode.BICUBIC:
        align_corners = False
173
174
175
176
    else:
        # The default of antialias should be True from 0.17, so we don't warn or
        # error if other interpolation modes are used. This is documented.
        antialias = False
177

178
179
    shape = image.shape
    num_channels, old_height, old_width = shape[-3:]
vfdev's avatar
vfdev committed
180
    new_height, new_width = _compute_resized_output_size((old_height, old_width), size=size, max_size=max_size)
181
182

    if image.numel() > 0:
183
        image = image.reshape(-1, num_channels, old_height, old_width)
184

185
186
187
188
189
190
        dtype = image.dtype
        need_cast = dtype not in (torch.float32, torch.float64)
        if need_cast:
            image = image.to(dtype=torch.float32)

        image = interpolate(
191
192
            image,
            size=[new_height, new_width],
193
194
            mode=interpolation.value,
            align_corners=align_corners,
195
196
            antialias=antialias,
        )
197

198
199
200
201
202
        if need_cast:
            if interpolation == InterpolationMode.BICUBIC and dtype == torch.uint8:
                image = image.clamp_(min=0, max=255)
            image = image.round_().to(dtype=dtype)

203
    return image.reshape(shape[:-3] + (num_channels, new_height, new_width))
204
205


206
@torch.jit.unused
207
def resize_image_pil(
208
    image: PIL.Image.Image,
209
    size: Union[Sequence[int], int],
210
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
211
212
    max_size: Optional[int] = None,
) -> PIL.Image.Image:
213
    interpolation = _check_interpolation(interpolation)
214
    size = _compute_resized_output_size(image.size[::-1], size=size, max_size=max_size)  # type: ignore[arg-type]
215
    return _FP.resize(image, size, interpolation=pil_modes_mapping[interpolation])
216
217


218
219
220
def resize_mask(mask: torch.Tensor, size: List[int], max_size: Optional[int] = None) -> torch.Tensor:
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
221
222
223
224
        needs_squeeze = True
    else:
        needs_squeeze = False

225
    output = resize_image_tensor(mask, size=size, interpolation=InterpolationMode.NEAREST, max_size=max_size)
226
227
228
229
230

    if needs_squeeze:
        output = output.squeeze(0)

    return output
231
232


233
def resize_bounding_box(
234
    bounding_box: torch.Tensor, spatial_size: Tuple[int, int], size: List[int], max_size: Optional[int] = None
235
) -> Tuple[torch.Tensor, Tuple[int, int]]:
236
237
    old_height, old_width = spatial_size
    new_height, new_width = _compute_resized_output_size(spatial_size, size=size, max_size=max_size)
238
239
240
    w_ratio = new_width / old_width
    h_ratio = new_height / old_height
    ratios = torch.tensor([w_ratio, h_ratio, w_ratio, h_ratio], device=bounding_box.device)
241
    return (
242
        bounding_box.mul(ratios).to(bounding_box.dtype),
243
244
        (new_height, new_width),
    )
245
246


247
248
249
def resize_video(
    video: torch.Tensor,
    size: List[int],
250
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
251
    max_size: Optional[int] = None,
252
    antialias: Optional[Union[str, bool]] = "warn",
253
254
255
256
) -> torch.Tensor:
    return resize_image_tensor(video, size=size, interpolation=interpolation, max_size=max_size, antialias=antialias)


257
def resize(
Philip Meier's avatar
Philip Meier committed
258
    inpt: datapoints._InputTypeJIT,
259
    size: List[int],
260
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
261
    max_size: Optional[int] = None,
262
    antialias: Optional[Union[str, bool]] = "warn",
Philip Meier's avatar
Philip Meier committed
263
) -> datapoints._InputTypeJIT:
264
265
    if not torch.jit.is_scripting():
        _log_api_usage_once(resize)
266
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
267
        return resize_image_tensor(inpt, size, interpolation=interpolation, max_size=max_size, antialias=antialias)
268
    elif isinstance(inpt, datapoints._datapoint.Datapoint):
269
        return inpt.resize(size, interpolation=interpolation, max_size=max_size, antialias=antialias)
270
    elif isinstance(inpt, PIL.Image.Image):
271
        if antialias is False:
272
273
            warnings.warn("Anti-alias option is always applied for PIL Image input. Argument antialias is ignored.")
        return resize_image_pil(inpt, size, interpolation=interpolation, max_size=max_size)
274
275
    else:
        raise TypeError(
276
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
277
278
            f"but got {type(inpt)} instead."
        )
279
280


281
def _affine_parse_args(
282
    angle: Union[int, float],
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
    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}")

325
326
327
328
329
    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]
330
331
332
333

    return angle, translate, shear, center


334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
def _get_inverse_affine_matrix(
    center: List[float], angle: float, translate: List[float], scale: float, shear: List[float], inverted: bool = True
) -> List[float]:
    # Helper method to compute inverse matrix for affine transformation

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

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

    cx, cy = center
    tx, ty = translate

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

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

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

    return matrix


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

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

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

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


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

434
435
436
437
    # 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)

438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
    shape = float_img.shape
    if shape[0] > 1:
        # Apply same grid to a batch of images
        grid = grid.expand(shape[0], -1, -1, -1)

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

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

    # Fill with required color
    if fill is not None:
        float_img, mask = torch.tensor_split(float_img, indices=(-1,), dim=-3)
        mask = mask.expand_as(float_img)
454
        fill_list = fill if isinstance(fill, (tuple, list)) else [float(fill)]  # type: ignore[arg-type]
455
456
457
458
459
460
461
462
463
        fill_img = torch.tensor(fill_list, dtype=float_img.dtype, device=float_img.device).view(1, -1, 1, 1)
        if mode == "nearest":
            bool_mask = mask < 0.5
            float_img[bool_mask] = fill_img.expand_as(float_img)[bool_mask]
        else:  # 'bilinear'
            # The following is mathematically equivalent to:
            # img * mask + (1.0 - mask) * fill = img * mask - fill * mask + fill = mask * (img - fill) + fill
            float_img = float_img.sub_(fill_img).mul_(mask).add_(fill_img)

464
465
466
    img = float_img.round_().to(img.dtype) if not fp else float_img

    return img
467
468
469
470
471
472


def _assert_grid_transform_inputs(
    image: torch.Tensor,
    matrix: Optional[List[float]],
    interpolation: str,
Philip Meier's avatar
Philip Meier committed
473
    fill: datapoints._FillTypeJIT,
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
    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)


529
def affine_image_tensor(
530
    image: torch.Tensor,
531
    angle: Union[int, float],
532
533
534
    translate: List[float],
    scale: float,
    shear: List[float],
535
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
Philip Meier's avatar
Philip Meier committed
536
    fill: datapoints._FillTypeJIT = None,
537
538
    center: Optional[List[float]] = None,
) -> torch.Tensor:
539
540
    interpolation = _check_interpolation(interpolation)

541
542
    if image.numel() == 0:
        return image
543

544
    shape = image.shape
545
    ndim = image.ndim
546

547
548
549
550
551
552
553
554
555
556
    if ndim > 4:
        image = image.reshape((-1,) + shape[-3:])
        needs_unsquash = True
    elif ndim == 3:
        image = image.unsqueeze(0)
        needs_unsquash = True
    else:
        needs_unsquash = False

    height, width = shape[-2:]
557
558
559
560
561
    angle, translate, shear, center = _affine_parse_args(angle, translate, scale, shear, interpolation, center)

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

564
    translate_f = [float(t) for t in translate]
565
566
    matrix = _get_inverse_affine_matrix(center_f, angle, translate_f, scale, shear)

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

569
    dtype = image.dtype if torch.is_floating_point(image) else torch.float32
570
571
    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)
572
    output = _apply_grid_transform(image, grid, interpolation.value, fill=fill)
573
574
575
576
577

    if needs_unsquash:
        output = output.reshape(shape)

    return output
578
579


580
@torch.jit.unused
581
def affine_image_pil(
582
    image: PIL.Image.Image,
583
    angle: Union[int, float],
584
585
586
    translate: List[float],
    scale: float,
    shear: List[float],
587
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
Philip Meier's avatar
Philip Meier committed
588
    fill: datapoints._FillTypeJIT = None,
589
590
    center: Optional[List[float]] = None,
) -> PIL.Image.Image:
591
    interpolation = _check_interpolation(interpolation)
592
593
594
595
596
597
    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:
598
        height, width = get_spatial_size_image_pil(image)
599
600
601
        center = [width * 0.5, height * 0.5]
    matrix = _get_inverse_affine_matrix(center, angle, translate, scale, shear)

602
    return _FP.affine(image, matrix, interpolation=pil_modes_mapping[interpolation], fill=fill)
603
604


605
def _affine_bounding_box_with_expand(
606
    bounding_box: torch.Tensor,
607
    format: datapoints.BoundingBoxFormat,
608
    spatial_size: Tuple[int, int],
609
610
611
612
    angle: Union[int, float],
    translate: List[float],
    scale: float,
    shear: List[float],
613
    center: Optional[List[float]] = None,
614
    expand: bool = False,
615
) -> Tuple[torch.Tensor, Tuple[int, int]]:
616
617
618
    if bounding_box.numel() == 0:
        return bounding_box, spatial_size

619
620
621
622
623
624
625
626
627
628
629
    original_shape = bounding_box.shape
    original_dtype = bounding_box.dtype
    bounding_box = bounding_box.clone() if bounding_box.is_floating_point() else bounding_box.float()
    dtype = bounding_box.dtype
    device = bounding_box.device
    bounding_box = (
        convert_format_bounding_box(
            bounding_box, old_format=format, new_format=datapoints.BoundingBoxFormat.XYXY, inplace=True
        )
    ).reshape(-1, 4)

630
631
632
    angle, translate, shear, center = _affine_parse_args(
        angle, translate, scale, shear, InterpolationMode.NEAREST, center
    )
633

634
    if center is None:
635
        height, width = spatial_size
636
637
        center = [width * 0.5, height * 0.5]

638
639
640
641
642
643
644
    affine_vector = _get_inverse_affine_matrix(center, angle, translate, scale, shear, inverted=False)
    transposed_affine_matrix = (
        torch.tensor(
            affine_vector,
            dtype=dtype,
            device=device,
        )
645
        .reshape(2, 3)
646
647
        .T
    )
648
649
650
651
    # 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)]
652
    points = bounding_box[:, [[0, 1], [2, 1], [2, 3], [0, 3]]].reshape(-1, 2)
653
    points = torch.cat([points, torch.ones(points.shape[0], 1, device=device, dtype=dtype)], dim=-1)
654
    # 2) Now let's transform the points using affine matrix
655
    transformed_points = torch.matmul(points, transposed_affine_matrix)
656
657
    # 3) Reshape transformed points to [N boxes, 4 points, x/y coords]
    # and compute bounding box from 4 transformed points:
658
    transformed_points = transformed_points.reshape(-1, 4, 2)
659
    out_bbox_mins, out_bbox_maxs = torch.aminmax(transformed_points, dim=1)
660
    out_bboxes = torch.cat([out_bbox_mins, out_bbox_maxs], dim=1)
661
662
663
664

    if expand:
        # Compute minimum point for transformed image frame:
        # Points are Top-Left, Top-Right, Bottom-Left, Bottom-Right points.
665
        height, width = spatial_size
666
667
668
        points = torch.tensor(
            [
                [0.0, 0.0, 1.0],
669
670
671
                [0.0, float(height), 1.0],
                [float(width), float(height), 1.0],
                [float(width), 0.0, 1.0],
672
673
674
675
            ],
            dtype=dtype,
            device=device,
        )
676
        new_points = torch.matmul(points, transposed_affine_matrix)
677
        tr = torch.amin(new_points, dim=0, keepdim=True)
678
        # Translate bounding boxes
679
        out_bboxes.sub_(tr.repeat((1, 2)))
680
681
        # Estimate meta-data for image with inverted=True and with center=[0,0]
        affine_vector = _get_inverse_affine_matrix([0.0, 0.0], angle, translate, scale, shear)
682
        new_width, new_height = _compute_affine_output_size(affine_vector, width, height)
683
        spatial_size = (new_height, new_width)
684

685
686
687
688
689
690
691
    out_bboxes = clamp_bounding_box(out_bboxes, format=datapoints.BoundingBoxFormat.XYXY, spatial_size=spatial_size)
    out_bboxes = convert_format_bounding_box(
        out_bboxes, old_format=datapoints.BoundingBoxFormat.XYXY, new_format=format, inplace=True
    ).reshape(original_shape)

    out_bboxes = out_bboxes.to(original_dtype)
    return out_bboxes, spatial_size
692
693
694
695


def affine_bounding_box(
    bounding_box: torch.Tensor,
696
    format: datapoints.BoundingBoxFormat,
697
    spatial_size: Tuple[int, int],
698
    angle: Union[int, float],
699
700
701
702
703
    translate: List[float],
    scale: float,
    shear: List[float],
    center: Optional[List[float]] = None,
) -> torch.Tensor:
704
705
706
707
708
709
710
711
712
713
714
715
    out_box, _ = _affine_bounding_box_with_expand(
        bounding_box,
        format=format,
        spatial_size=spatial_size,
        angle=angle,
        translate=translate,
        scale=scale,
        shear=shear,
        center=center,
        expand=False,
    )
    return out_box
716
717


718
719
def affine_mask(
    mask: torch.Tensor,
720
    angle: Union[int, float],
721
722
723
    translate: List[float],
    scale: float,
    shear: List[float],
Philip Meier's avatar
Philip Meier committed
724
    fill: datapoints._FillTypeJIT = None,
725
726
    center: Optional[List[float]] = None,
) -> torch.Tensor:
727
728
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
729
730
731
732
733
        needs_squeeze = True
    else:
        needs_squeeze = False

    output = affine_image_tensor(
734
        mask,
735
736
737
738
739
        angle=angle,
        translate=translate,
        scale=scale,
        shear=shear,
        interpolation=InterpolationMode.NEAREST,
740
        fill=fill,
741
742
743
        center=center,
    )

744
745
746
747
748
    if needs_squeeze:
        output = output.squeeze(0)

    return output

749

750
751
752
753
754
755
def affine_video(
    video: torch.Tensor,
    angle: Union[int, float],
    translate: List[float],
    scale: float,
    shear: List[float],
756
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
Philip Meier's avatar
Philip Meier committed
757
    fill: datapoints._FillTypeJIT = None,
758
759
760
761
762
763
764
765
766
767
768
769
770
771
    center: Optional[List[float]] = None,
) -> torch.Tensor:
    return affine_image_tensor(
        video,
        angle=angle,
        translate=translate,
        scale=scale,
        shear=shear,
        interpolation=interpolation,
        fill=fill,
        center=center,
    )


772
def affine(
Philip Meier's avatar
Philip Meier committed
773
    inpt: datapoints._InputTypeJIT,
774
    angle: Union[int, float],
775
776
777
    translate: List[float],
    scale: float,
    shear: List[float],
778
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
Philip Meier's avatar
Philip Meier committed
779
    fill: datapoints._FillTypeJIT = None,
780
    center: Optional[List[float]] = None,
Philip Meier's avatar
Philip Meier committed
781
) -> datapoints._InputTypeJIT:
782
783
784
    if not torch.jit.is_scripting():
        _log_api_usage_once(affine)

785
    # TODO: consider deprecating integers from angle and shear on the future
786
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
787
        return affine_image_tensor(
788
789
790
791
792
793
794
795
796
            inpt,
            angle,
            translate=translate,
            scale=scale,
            shear=shear,
            interpolation=interpolation,
            fill=fill,
            center=center,
        )
797
    elif isinstance(inpt, datapoints._datapoint.Datapoint):
798
799
800
        return inpt.affine(
            angle, translate=translate, scale=scale, shear=shear, interpolation=interpolation, fill=fill, center=center
        )
801
    elif isinstance(inpt, PIL.Image.Image):
802
        return affine_image_pil(
803
804
805
806
807
808
809
810
811
            inpt,
            angle,
            translate=translate,
            scale=scale,
            shear=shear,
            interpolation=interpolation,
            fill=fill,
            center=center,
        )
812
813
    else:
        raise TypeError(
814
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
815
816
            f"but got {type(inpt)} instead."
        )
817
818


819
def rotate_image_tensor(
820
    image: torch.Tensor,
821
    angle: float,
822
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
823
824
    expand: bool = False,
    center: Optional[List[float]] = None,
Philip Meier's avatar
Philip Meier committed
825
    fill: datapoints._FillTypeJIT = None,
826
) -> torch.Tensor:
827
828
    interpolation = _check_interpolation(interpolation)

829
830
    shape = image.shape
    num_channels, height, width = shape[-3:]
831

832
833
    center_f = [0.0, 0.0]
    if center is not None:
834
        if expand:
835
            # TODO: Do we actually want to warn, or just document this?
836
            warnings.warn("The provided center argument has no effect on the result if expand is True")
837
838
        # Center values should be in pixel coordinates but translated such that (0, 0) corresponds to image center.
        center_f = [(c - s * 0.5) for c, s in zip(center, [width, height])]
839
840
841
842

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

844
    if image.numel() > 0:
845
846
847
848
849
        image = image.reshape(-1, num_channels, height, width)

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

        ow, oh = _compute_affine_output_size(matrix, width, height) if expand else (width, height)
850
        dtype = image.dtype if torch.is_floating_point(image) else torch.float32
851
852
        theta = torch.tensor(matrix, dtype=dtype, device=image.device).reshape(1, 2, 3)
        grid = _affine_grid(theta, w=width, h=height, ow=ow, oh=oh)
853
        output = _apply_grid_transform(image, grid, interpolation.value, fill=fill)
854
855

        new_height, new_width = output.shape[-2:]
856
    else:
857
858
        output = image
        new_width, new_height = _compute_affine_output_size(matrix, width, height) if expand else (width, height)
859

860
    return output.reshape(shape[:-3] + (num_channels, new_height, new_width))
861
862


863
@torch.jit.unused
864
def rotate_image_pil(
865
    image: PIL.Image.Image,
866
    angle: float,
867
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
868
869
    expand: bool = False,
    center: Optional[List[float]] = None,
Philip Meier's avatar
Philip Meier committed
870
    fill: datapoints._FillTypeJIT = None,
871
) -> PIL.Image.Image:
872
873
    interpolation = _check_interpolation(interpolation)

874
    if center is not None and expand:
875
        warnings.warn("The provided center argument has no effect on the result if expand is True")
876
877
        center = None

878
    return _FP.rotate(
879
        image, angle, interpolation=pil_modes_mapping[interpolation], expand=expand, fill=fill, center=center
880
881
882
    )


883
884
def rotate_bounding_box(
    bounding_box: torch.Tensor,
885
    format: datapoints.BoundingBoxFormat,
886
    spatial_size: Tuple[int, int],
887
888
889
    angle: float,
    expand: bool = False,
    center: Optional[List[float]] = None,
890
) -> Tuple[torch.Tensor, Tuple[int, int]]:
891
892
893
894
    if center is not None and expand:
        warnings.warn("The provided center argument has no effect on the result if expand is True")
        center = None

895
    return _affine_bounding_box_with_expand(
896
        bounding_box,
897
898
        format=format,
        spatial_size=spatial_size,
899
900
901
902
903
904
905
        angle=-angle,
        translate=[0.0, 0.0],
        scale=1.0,
        shear=[0.0, 0.0],
        center=center,
        expand=expand,
    )
906
907


908
909
def rotate_mask(
    mask: torch.Tensor,
910
911
912
    angle: float,
    expand: bool = False,
    center: Optional[List[float]] = None,
Philip Meier's avatar
Philip Meier committed
913
    fill: datapoints._FillTypeJIT = None,
914
) -> torch.Tensor:
915
916
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
917
918
919
920
921
        needs_squeeze = True
    else:
        needs_squeeze = False

    output = rotate_image_tensor(
922
        mask,
923
924
925
        angle=angle,
        expand=expand,
        interpolation=InterpolationMode.NEAREST,
926
        fill=fill,
927
928
929
        center=center,
    )

930
931
932
933
934
    if needs_squeeze:
        output = output.squeeze(0)

    return output

935

936
937
938
def rotate_video(
    video: torch.Tensor,
    angle: float,
939
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
940
941
    expand: bool = False,
    center: Optional[List[float]] = None,
Philip Meier's avatar
Philip Meier committed
942
    fill: datapoints._FillTypeJIT = None,
943
944
945
946
) -> torch.Tensor:
    return rotate_image_tensor(video, angle, interpolation=interpolation, expand=expand, fill=fill, center=center)


947
def rotate(
Philip Meier's avatar
Philip Meier committed
948
    inpt: datapoints._InputTypeJIT,
949
    angle: float,
950
    interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
951
952
    expand: bool = False,
    center: Optional[List[float]] = None,
Philip Meier's avatar
Philip Meier committed
953
954
    fill: datapoints._FillTypeJIT = None,
) -> datapoints._InputTypeJIT:
955
956
957
    if not torch.jit.is_scripting():
        _log_api_usage_once(rotate)

958
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
959
        return rotate_image_tensor(inpt, angle, interpolation=interpolation, expand=expand, fill=fill, center=center)
960
    elif isinstance(inpt, datapoints._datapoint.Datapoint):
961
        return inpt.rotate(angle, interpolation=interpolation, expand=expand, fill=fill, center=center)
962
    elif isinstance(inpt, PIL.Image.Image):
963
        return rotate_image_pil(inpt, angle, interpolation=interpolation, expand=expand, fill=fill, center=center)
964
965
    else:
        raise TypeError(
966
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
967
968
            f"but got {type(inpt)} instead."
        )
969
970


971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
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]
993

994

995
def pad_image_tensor(
996
    image: torch.Tensor,
997
998
    padding: List[int],
    fill: Optional[Union[int, float, List[float]]] = None,
999
1000
    padding_mode: str = "constant",
) -> torch.Tensor:
1001
1002
1003
1004
1005
    # Be aware that while `padding` has order `[left, top, right, bottom]` has order, `torch_padding` uses
    # `[left, right, top, bottom]`. This stems from the fact that we align our API with PIL, but need to use `torch_pad`
    # internally.
    torch_padding = _parse_pad_padding(padding)

1006
    if padding_mode not in ("constant", "edge", "reflect", "symmetric"):
1007
1008
1009
1010
1011
        raise ValueError(
            f"`padding_mode` should be either `'constant'`, `'edge'`, `'reflect'` or `'symmetric'`, "
            f"but got `'{padding_mode}'`."
        )

1012
    if fill is None:
1013
1014
1015
1016
1017
1018
        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)
1019
    else:
1020
        return _pad_with_vector_fill(image, torch_padding, fill=fill, padding_mode=padding_mode)
1021
1022
1023


def _pad_with_scalar_fill(
1024
    image: torch.Tensor,
1025
1026
1027
    torch_padding: List[int],
    fill: Union[int, float],
    padding_mode: str,
1028
) -> torch.Tensor:
1029
1030
    shape = image.shape
    num_channels, height, width = shape[-3:]
1031

1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
    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
1055

1056
1057
1058
1059
1060
        image = torch_pad(image, torch_padding, mode=padding_mode)

        if needs_cast:
            image = image.to(dtype)
    else:  # padding_mode == "symmetric"
1061
        image = _pad_symmetric(image, torch_padding)
1062
1063

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

1065
    return image.reshape(shape[:-3] + (num_channels, new_height, new_width))
1066
1067


1068
# TODO: This should be removed once torch_pad supports non-scalar padding values
1069
def _pad_with_vector_fill(
1070
    image: torch.Tensor,
1071
    torch_padding: List[int],
1072
    fill: List[float],
1073
    padding_mode: str,
1074
1075
1076
1077
) -> torch.Tensor:
    if padding_mode != "constant":
        raise ValueError(f"Padding mode '{padding_mode}' is not supported if fill is not scalar")

1078
1079
    output = _pad_with_scalar_fill(image, torch_padding, fill=0, padding_mode="constant")
    left, right, top, bottom = torch_padding
1080
    fill = torch.tensor(fill, dtype=image.dtype, device=image.device).reshape(-1, 1, 1)
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092

    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


1093
1094
1095
pad_image_pil = _FP.pad


1096
1097
def pad_mask(
    mask: torch.Tensor,
1098
1099
    padding: List[int],
    fill: Optional[Union[int, float, List[float]]] = None,
1100
1101
    padding_mode: str = "constant",
) -> torch.Tensor:
1102
1103
1104
    if fill is None:
        fill = 0

1105
    if isinstance(fill, (tuple, list)):
1106
1107
        raise ValueError("Non-scalar fill value is not supported")

1108
1109
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
1110
1111
1112
1113
        needs_squeeze = True
    else:
        needs_squeeze = False

1114
    output = pad_image_tensor(mask, padding=padding, fill=fill, padding_mode=padding_mode)
1115
1116
1117
1118
1119

    if needs_squeeze:
        output = output.squeeze(0)

    return output
1120
1121


1122
def pad_bounding_box(
vfdev's avatar
vfdev committed
1123
    bounding_box: torch.Tensor,
1124
    format: datapoints.BoundingBoxFormat,
1125
    spatial_size: Tuple[int, int],
1126
    padding: List[int],
vfdev's avatar
vfdev committed
1127
    padding_mode: str = "constant",
1128
) -> Tuple[torch.Tensor, Tuple[int, int]]:
vfdev's avatar
vfdev committed
1129
1130
1131
1132
    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")

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

1135
    if format == datapoints.BoundingBoxFormat.XYXY:
1136
1137
1138
1139
        pad = [left, top, left, top]
    else:
        pad = [left, top, 0, 0]
    bounding_box = bounding_box + torch.tensor(pad, dtype=bounding_box.dtype, device=bounding_box.device)
1140

1141
    height, width = spatial_size
1142
1143
    height += top + bottom
    width += left + right
1144
    spatial_size = (height, width)
1145

1146
    return clamp_bounding_box(bounding_box, format=format, spatial_size=spatial_size), spatial_size
1147
1148


1149
1150
def pad_video(
    video: torch.Tensor,
1151
1152
    padding: List[int],
    fill: Optional[Union[int, float, List[float]]] = None,
1153
1154
1155
1156
1157
    padding_mode: str = "constant",
) -> torch.Tensor:
    return pad_image_tensor(video, padding, fill=fill, padding_mode=padding_mode)


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

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

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


1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
def crop_image_tensor(image: torch.Tensor, top: int, left: int, height: int, width: int) -> torch.Tensor:
    h, w = image.shape[-2:]

    right = left + width
    bottom = top + height

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


1199
1200
1201
crop_image_pil = _FP.crop


1202
1203
def crop_bounding_box(
    bounding_box: torch.Tensor,
1204
    format: datapoints.BoundingBoxFormat,
1205
1206
    top: int,
    left: int,
1207
1208
1209
    height: int,
    width: int,
) -> Tuple[torch.Tensor, Tuple[int, int]]:
1210

1211
    # Crop or implicit pad if left and/or top have negative values:
1212
    if format == datapoints.BoundingBoxFormat.XYXY:
1213
        sub = [left, top, left, top]
1214
    else:
1215
1216
1217
        sub = [left, top, 0, 0]

    bounding_box = bounding_box - torch.tensor(sub, dtype=bounding_box.dtype, device=bounding_box.device)
1218
    spatial_size = (height, width)
1219

1220
    return clamp_bounding_box(bounding_box, format=format, spatial_size=spatial_size), spatial_size
1221
1222


1223
def crop_mask(mask: torch.Tensor, top: int, left: int, height: int, width: int) -> torch.Tensor:
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
        needs_squeeze = True
    else:
        needs_squeeze = False

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

    if needs_squeeze:
        output = output.squeeze(0)

    return output
1236
1237


1238
1239
1240
1241
def crop_video(video: torch.Tensor, top: int, left: int, height: int, width: int) -> torch.Tensor:
    return crop_image_tensor(video, top, left, height, width)


Philip Meier's avatar
Philip Meier committed
1242
def crop(inpt: datapoints._InputTypeJIT, top: int, left: int, height: int, width: int) -> datapoints._InputTypeJIT:
1243
1244
1245
    if not torch.jit.is_scripting():
        _log_api_usage_once(crop)

1246
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
1247
        return crop_image_tensor(inpt, top, left, height, width)
1248
    elif isinstance(inpt, datapoints._datapoint.Datapoint):
1249
        return inpt.crop(top, left, height, width)
1250
    elif isinstance(inpt, PIL.Image.Image):
1251
        return crop_image_pil(inpt, top, left, height, width)
1252
1253
    else:
        raise TypeError(
1254
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
1255
1256
            f"but got {type(inpt)} instead."
        )
1257
1258


1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
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)
1274
    x_grid = torch.linspace(d, ow + d - 1.0, steps=ow, device=device, dtype=dtype)
1275
    base_grid[..., 0].copy_(x_grid)
1276
    y_grid = torch.linspace(d, oh + d - 1.0, steps=oh, device=device, dtype=dtype).unsqueeze_(-1)
1277
1278
1279
1280
    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))
1281
1282
1283
    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))
1284
1285
1286
1287
1288

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


1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
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.")


1306
def perspective_image_tensor(
1307
    image: torch.Tensor,
1308
1309
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
1310
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
Philip Meier's avatar
Philip Meier committed
1311
    fill: datapoints._FillTypeJIT = None,
1312
    coefficients: Optional[List[float]] = None,
1313
) -> torch.Tensor:
1314
    perspective_coeffs = _perspective_coefficients(startpoints, endpoints, coefficients)
1315
1316
    interpolation = _check_interpolation(interpolation)

1317
1318
1319
1320
    if image.numel() == 0:
        return image

    shape = image.shape
1321
    ndim = image.ndim
1322

1323
    if ndim > 4:
1324
        image = image.reshape((-1,) + shape[-3:])
1325
        needs_unsquash = True
1326
1327
1328
    elif ndim == 3:
        image = image.unsqueeze(0)
        needs_unsquash = True
1329
1330
1331
    else:
        needs_unsquash = False

1332
    _assert_grid_transform_inputs(
1333
1334
1335
1336
1337
1338
1339
1340
        image,
        matrix=None,
        interpolation=interpolation.value,
        fill=fill,
        supported_interpolation_modes=["nearest", "bilinear"],
        coeffs=perspective_coeffs,
    )

1341
    oh, ow = shape[-2:]
1342
    dtype = image.dtype if torch.is_floating_point(image) else torch.float32
1343
    grid = _perspective_grid(perspective_coeffs, ow=ow, oh=oh, dtype=dtype, device=image.device)
1344
    output = _apply_grid_transform(image, grid, interpolation.value, fill=fill)
1345
1346

    if needs_unsquash:
1347
        output = output.reshape(shape)
1348
1349

    return output
1350
1351


1352
@torch.jit.unused
1353
def perspective_image_pil(
1354
    image: PIL.Image.Image,
1355
1356
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
1357
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BICUBIC,
Philip Meier's avatar
Philip Meier committed
1358
    fill: datapoints._FillTypeJIT = None,
1359
    coefficients: Optional[List[float]] = None,
1360
) -> PIL.Image.Image:
1361
    perspective_coeffs = _perspective_coefficients(startpoints, endpoints, coefficients)
1362
    interpolation = _check_interpolation(interpolation)
1363
    return _FP.perspective(image, perspective_coeffs, interpolation=pil_modes_mapping[interpolation], fill=fill)
1364
1365


1366
1367
def perspective_bounding_box(
    bounding_box: torch.Tensor,
1368
    format: datapoints.BoundingBoxFormat,
1369
    spatial_size: Tuple[int, int],
1370
1371
1372
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
    coefficients: Optional[List[float]] = None,
1373
) -> torch.Tensor:
1374
1375
1376
    if bounding_box.numel() == 0:
        return bounding_box

1377
    perspective_coeffs = _perspective_coefficients(startpoints, endpoints, coefficients)
1378
1379

    original_shape = bounding_box.shape
1380
    # TODO: first cast to float if bbox is int64 before convert_format_bounding_box
1381
    bounding_box = (
1382
        convert_format_bounding_box(bounding_box, old_format=format, new_format=datapoints.BoundingBoxFormat.XYXY)
1383
    ).reshape(-1, 4)
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417

    dtype = bounding_box.dtype if torch.is_floating_point(bounding_box) else torch.float32
    device = bounding_box.device

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

1418
1419
    theta1 = torch.tensor(
        [[inv_coeffs[0], inv_coeffs[1], inv_coeffs[2]], [inv_coeffs[3], inv_coeffs[4], inv_coeffs[5]]],
1420
1421
1422
1423
        dtype=dtype,
        device=device,
    )

1424
1425
1426
1427
    theta2 = torch.tensor(
        [[inv_coeffs[6], inv_coeffs[7], 1.0], [inv_coeffs[6], inv_coeffs[7], 1.0]], dtype=dtype, device=device
    )

1428
1429
1430
1431
    # 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)]
1432
    points = bounding_box[:, [[0, 1], [2, 1], [2, 3], [0, 3]]].reshape(-1, 2)
1433
1434
1435
1436
1437
    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)

1438
1439
    numer_points = torch.matmul(points, theta1.T)
    denom_points = torch.matmul(points, theta2.T)
1440
    transformed_points = numer_points.div_(denom_points)
1441
1442
1443

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

1447
1448
1449
1450
1451
    out_bboxes = clamp_bounding_box(
        torch.cat([out_bbox_mins, out_bbox_maxs], dim=1).to(bounding_box.dtype),
        format=datapoints.BoundingBoxFormat.XYXY,
        spatial_size=spatial_size,
    )
1452
1453
1454

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

1455
    return convert_format_bounding_box(
1456
        out_bboxes, old_format=datapoints.BoundingBoxFormat.XYXY, new_format=format, inplace=True
1457
    ).reshape(original_shape)
1458
1459


1460
1461
def perspective_mask(
    mask: torch.Tensor,
1462
1463
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
Philip Meier's avatar
Philip Meier committed
1464
    fill: datapoints._FillTypeJIT = None,
1465
    coefficients: Optional[List[float]] = None,
1466
) -> torch.Tensor:
1467
1468
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
1469
1470
1471
1472
1473
        needs_squeeze = True
    else:
        needs_squeeze = False

    output = perspective_image_tensor(
1474
        mask, startpoints, endpoints, interpolation=InterpolationMode.NEAREST, fill=fill, coefficients=coefficients
1475
    )
1476

1477
1478
1479
1480
1481
    if needs_squeeze:
        output = output.squeeze(0)

    return output

1482

1483
1484
def perspective_video(
    video: torch.Tensor,
1485
1486
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
1487
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
Philip Meier's avatar
Philip Meier committed
1488
    fill: datapoints._FillTypeJIT = None,
1489
    coefficients: Optional[List[float]] = None,
1490
) -> torch.Tensor:
1491
1492
1493
    return perspective_image_tensor(
        video, startpoints, endpoints, interpolation=interpolation, fill=fill, coefficients=coefficients
    )
1494
1495


1496
def perspective(
Philip Meier's avatar
Philip Meier committed
1497
    inpt: datapoints._InputTypeJIT,
1498
1499
    startpoints: Optional[List[List[int]]],
    endpoints: Optional[List[List[int]]],
1500
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
Philip Meier's avatar
Philip Meier committed
1501
    fill: datapoints._FillTypeJIT = None,
1502
    coefficients: Optional[List[float]] = None,
Philip Meier's avatar
Philip Meier committed
1503
) -> datapoints._InputTypeJIT:
1504
1505
    if not torch.jit.is_scripting():
        _log_api_usage_once(perspective)
1506
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
1507
1508
1509
        return perspective_image_tensor(
            inpt, startpoints, endpoints, interpolation=interpolation, fill=fill, coefficients=coefficients
        )
1510
    elif isinstance(inpt, datapoints._datapoint.Datapoint):
1511
1512
1513
        return inpt.perspective(
            startpoints, endpoints, interpolation=interpolation, fill=fill, coefficients=coefficients
        )
1514
    elif isinstance(inpt, PIL.Image.Image):
1515
1516
1517
        return perspective_image_pil(
            inpt, startpoints, endpoints, interpolation=interpolation, fill=fill, coefficients=coefficients
        )
1518
1519
    else:
        raise TypeError(
1520
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
1521
1522
            f"but got {type(inpt)} instead."
        )
1523
1524


1525
def elastic_image_tensor(
1526
    image: torch.Tensor,
1527
    displacement: torch.Tensor,
1528
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
Philip Meier's avatar
Philip Meier committed
1529
    fill: datapoints._FillTypeJIT = None,
1530
) -> torch.Tensor:
1531
1532
    interpolation = _check_interpolation(interpolation)

1533
1534
1535
1536
    if image.numel() == 0:
        return image

    shape = image.shape
1537
    ndim = image.ndim
1538

1539
    device = image.device
1540
    dtype = image.dtype if torch.is_floating_point(image) else torch.float32
1541
1542
1543
1544
1545
1546
1547

    # 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

1548
1549
1550
    # We are aware that if input image dtype is uint8 and displacement is float64 then
    # displacement will be casted to float32 and all computations will be done with float32
    # We can fix this later if needed
1551

1552
1553
1554
1555
    expected_shape = (1,) + shape[-2:] + (2,)
    if expected_shape != displacement.shape:
        raise ValueError(f"Argument displacement shape should be {expected_shape}, but given {displacement.shape}")

1556
    if ndim > 4:
1557
        image = image.reshape((-1,) + shape[-3:])
1558
        needs_unsquash = True
1559
1560
1561
    elif ndim == 3:
        image = image.unsqueeze(0)
        needs_unsquash = True
1562
1563
1564
    else:
        needs_unsquash = False

1565
1566
    if displacement.dtype != dtype or displacement.device != device:
        displacement = displacement.to(dtype=dtype, device=device)
1567

1568
1569
1570
    image_height, image_width = shape[-2:]
    grid = _create_identity_grid((image_height, image_width), device=device, dtype=dtype).add_(displacement)
    output = _apply_grid_transform(image, grid, interpolation.value, fill=fill)
1571
1572

    if needs_unsquash:
1573
        output = output.reshape(shape)
1574

1575
1576
1577
    if is_cpu_half:
        output = output.to(torch.float16)

1578
    return output
1579
1580


1581
@torch.jit.unused
1582
def elastic_image_pil(
1583
    image: PIL.Image.Image,
1584
    displacement: torch.Tensor,
1585
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
Philip Meier's avatar
Philip Meier committed
1586
    fill: datapoints._FillTypeJIT = None,
1587
) -> PIL.Image.Image:
1588
    t_img = pil_to_tensor(image)
1589
    output = elastic_image_tensor(t_img, displacement, interpolation=interpolation, fill=fill)
1590
    return to_pil_image(output, mode=image.mode)
1591
1592


1593
def _create_identity_grid(size: Tuple[int, int], device: torch.device, dtype: torch.dtype) -> torch.Tensor:
1594
    sy, sx = size
1595
1596
    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)
1597
1598
    base_grid[..., 0].copy_(x_grid)

1599
    y_grid = torch.linspace((-sy + 1) / sy, (sy - 1) / sy, sy, device=device, dtype=dtype).unsqueeze_(-1)
1600
1601
1602
1603
1604
    base_grid[..., 1].copy_(y_grid)

    return base_grid


1605
1606
def elastic_bounding_box(
    bounding_box: torch.Tensor,
1607
    format: datapoints.BoundingBoxFormat,
1608
    spatial_size: Tuple[int, int],
1609
1610
    displacement: torch.Tensor,
) -> torch.Tensor:
1611
1612
1613
    if bounding_box.numel() == 0:
        return bounding_box

1614
    # TODO: add in docstring about approximation we are doing for grid inversion
1615
1616
1617
1618
1619
    device = bounding_box.device
    dtype = bounding_box.dtype if torch.is_floating_point(bounding_box) else torch.float32

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

    original_shape = bounding_box.shape
1622
    # TODO: first cast to float if bbox is int64 before convert_format_bounding_box
1623
    bounding_box = (
1624
        convert_format_bounding_box(bounding_box, old_format=format, new_format=datapoints.BoundingBoxFormat.XYXY)
1625
    ).reshape(-1, 4)
1626

1627
    id_grid = _create_identity_grid(spatial_size, device=device, dtype=dtype)
1628
1629
    # We construct an approximation of inverse grid as inv_grid = id_grid - displacement
    # This is not an exact inverse of the grid
1630
    inv_grid = id_grid.sub_(displacement)
1631
1632

    # Get points from bboxes
1633
    points = bounding_box[:, [[0, 1], [2, 1], [2, 3], [0, 3]]].reshape(-1, 2)
1634
1635
1636
1637
1638
    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]

1639
    # Transform points:
1640
    t_size = torch.tensor(spatial_size[::-1], device=displacement.device, dtype=displacement.dtype)
1641
    transformed_points = inv_grid[0, index_y, index_x, :].add_(1).mul_(0.5 * t_size).sub_(0.5)
1642

1643
    transformed_points = transformed_points.reshape(-1, 4, 2)
1644
    out_bbox_mins, out_bbox_maxs = torch.aminmax(transformed_points, dim=1)
1645
1646
1647
1648
1649
    out_bboxes = clamp_bounding_box(
        torch.cat([out_bbox_mins, out_bbox_maxs], dim=1).to(bounding_box.dtype),
        format=datapoints.BoundingBoxFormat.XYXY,
        spatial_size=spatial_size,
    )
1650

1651
    return convert_format_bounding_box(
1652
        out_bboxes, old_format=datapoints.BoundingBoxFormat.XYXY, new_format=format, inplace=True
1653
    ).reshape(original_shape)
1654
1655


1656
1657
1658
def elastic_mask(
    mask: torch.Tensor,
    displacement: torch.Tensor,
Philip Meier's avatar
Philip Meier committed
1659
    fill: datapoints._FillTypeJIT = None,
1660
) -> torch.Tensor:
1661
1662
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
1663
1664
1665
1666
        needs_squeeze = True
    else:
        needs_squeeze = False

1667
    output = elastic_image_tensor(mask, displacement=displacement, interpolation=InterpolationMode.NEAREST, fill=fill)
1668
1669
1670
1671
1672

    if needs_squeeze:
        output = output.squeeze(0)

    return output
1673
1674


1675
1676
1677
def elastic_video(
    video: torch.Tensor,
    displacement: torch.Tensor,
1678
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
Philip Meier's avatar
Philip Meier committed
1679
    fill: datapoints._FillTypeJIT = None,
1680
) -> torch.Tensor:
1681
    return elastic_image_tensor(video, displacement, interpolation=interpolation, fill=fill)
1682
1683


1684
def elastic(
Philip Meier's avatar
Philip Meier committed
1685
    inpt: datapoints._InputTypeJIT,
1686
    displacement: torch.Tensor,
1687
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
Philip Meier's avatar
Philip Meier committed
1688
1689
    fill: datapoints._FillTypeJIT = None,
) -> datapoints._InputTypeJIT:
1690
1691
1692
    if not torch.jit.is_scripting():
        _log_api_usage_once(elastic)

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

1696
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
1697
        return elastic_image_tensor(inpt, displacement, interpolation=interpolation, fill=fill)
1698
    elif isinstance(inpt, datapoints._datapoint.Datapoint):
1699
        return inpt.elastic(displacement, interpolation=interpolation, fill=fill)
1700
    elif isinstance(inpt, PIL.Image.Image):
1701
        return elastic_image_pil(inpt, displacement, interpolation=interpolation, fill=fill)
1702
1703
    else:
        raise TypeError(
1704
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
1705
1706
            f"but got {type(inpt)} instead."
        )
1707
1708
1709
1710
1711


elastic_transform = elastic


1712
1713
def _center_crop_parse_output_size(output_size: List[int]) -> List[int]:
    if isinstance(output_size, numbers.Number):
1714
1715
        s = int(output_size)
        return [s, s]
1716
    elif isinstance(output_size, (tuple, list)) and len(output_size) == 1:
1717
        return [output_size[0], output_size[0]]
1718
1719
    else:
        return list(output_size)
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738


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


1739
def center_crop_image_tensor(image: torch.Tensor, output_size: List[int]) -> torch.Tensor:
1740
    crop_height, crop_width = _center_crop_parse_output_size(output_size)
1741
1742
1743
1744
    shape = image.shape
    if image.numel() == 0:
        return image.reshape(shape[:-2] + (crop_height, crop_width))
    image_height, image_width = shape[-2:]
1745
1746
1747

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

1750
        image_height, image_width = image.shape[-2:]
1751
        if crop_width == image_width and crop_height == image_height:
1752
            return image
1753
1754

    crop_top, crop_left = _center_crop_compute_crop_anchor(crop_height, crop_width, image_height, image_width)
1755
    return image[..., crop_top : (crop_top + crop_height), crop_left : (crop_left + crop_width)]
1756
1757


1758
@torch.jit.unused
1759
def center_crop_image_pil(image: PIL.Image.Image, output_size: List[int]) -> PIL.Image.Image:
1760
    crop_height, crop_width = _center_crop_parse_output_size(output_size)
1761
    image_height, image_width = get_spatial_size_image_pil(image)
1762
1763
1764

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

1767
        image_height, image_width = get_spatial_size_image_pil(image)
1768
        if crop_width == image_width and crop_height == image_height:
1769
            return image
1770
1771

    crop_top, crop_left = _center_crop_compute_crop_anchor(crop_height, crop_width, image_height, image_width)
1772
    return crop_image_pil(image, crop_top, crop_left, crop_height, crop_width)
1773
1774


1775
1776
def center_crop_bounding_box(
    bounding_box: torch.Tensor,
1777
    format: datapoints.BoundingBoxFormat,
1778
    spatial_size: Tuple[int, int],
1779
    output_size: List[int],
1780
) -> Tuple[torch.Tensor, Tuple[int, int]]:
1781
    crop_height, crop_width = _center_crop_parse_output_size(output_size)
1782
    crop_top, crop_left = _center_crop_compute_crop_anchor(crop_height, crop_width, *spatial_size)
1783
    return crop_bounding_box(bounding_box, format, top=crop_top, left=crop_left, height=crop_height, width=crop_width)
1784
1785


1786
1787
1788
def center_crop_mask(mask: torch.Tensor, output_size: List[int]) -> torch.Tensor:
    if mask.ndim < 3:
        mask = mask.unsqueeze(0)
1789
1790
1791
1792
        needs_squeeze = True
    else:
        needs_squeeze = False

1793
    output = center_crop_image_tensor(image=mask, output_size=output_size)
1794
1795
1796
1797
1798

    if needs_squeeze:
        output = output.squeeze(0)

    return output
1799
1800


1801
1802
1803
1804
def center_crop_video(video: torch.Tensor, output_size: List[int]) -> torch.Tensor:
    return center_crop_image_tensor(video, output_size)


Philip Meier's avatar
Philip Meier committed
1805
def center_crop(inpt: datapoints._InputTypeJIT, output_size: List[int]) -> datapoints._InputTypeJIT:
1806
1807
1808
    if not torch.jit.is_scripting():
        _log_api_usage_once(center_crop)

1809
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
1810
        return center_crop_image_tensor(inpt, output_size)
1811
    elif isinstance(inpt, datapoints._datapoint.Datapoint):
1812
        return inpt.center_crop(output_size)
1813
    elif isinstance(inpt, PIL.Image.Image):
1814
        return center_crop_image_pil(inpt, output_size)
1815
1816
    else:
        raise TypeError(
1817
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
1818
1819
            f"but got {type(inpt)} instead."
        )
1820
1821


1822
def resized_crop_image_tensor(
1823
    image: torch.Tensor,
1824
1825
1826
1827
1828
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
1829
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1830
    antialias: Optional[Union[str, bool]] = "warn",
1831
) -> torch.Tensor:
1832
1833
    image = crop_image_tensor(image, top, left, height, width)
    return resize_image_tensor(image, size, interpolation=interpolation, antialias=antialias)
1834
1835


1836
@torch.jit.unused
1837
def resized_crop_image_pil(
1838
    image: PIL.Image.Image,
1839
1840
1841
1842
1843
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
1844
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1845
) -> PIL.Image.Image:
1846
1847
    image = crop_image_pil(image, top, left, height, width)
    return resize_image_pil(image, size, interpolation=interpolation)
1848
1849


1850
1851
def resized_crop_bounding_box(
    bounding_box: torch.Tensor,
1852
    format: datapoints.BoundingBoxFormat,
1853
1854
1855
1856
1857
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
1858
1859
) -> Tuple[torch.Tensor, Tuple[int, int]]:
    bounding_box, _ = crop_bounding_box(bounding_box, format, top, left, height, width)
1860
    return resize_bounding_box(bounding_box, spatial_size=(height, width), size=size)
1861
1862


1863
def resized_crop_mask(
1864
1865
1866
1867
1868
1869
1870
    mask: torch.Tensor,
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
) -> torch.Tensor:
1871
1872
    mask = crop_mask(mask, top, left, height, width)
    return resize_mask(mask, size)
1873
1874


1875
1876
1877
1878
1879
1880
1881
def resized_crop_video(
    video: torch.Tensor,
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
1882
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1883
    antialias: Optional[Union[str, bool]] = "warn",
1884
1885
1886
1887
1888
1889
) -> torch.Tensor:
    return resized_crop_image_tensor(
        video, top, left, height, width, antialias=antialias, size=size, interpolation=interpolation
    )


1890
def resized_crop(
Philip Meier's avatar
Philip Meier committed
1891
    inpt: datapoints._InputTypeJIT,
1892
1893
1894
1895
1896
    top: int,
    left: int,
    height: int,
    width: int,
    size: List[int],
1897
    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
1898
    antialias: Optional[Union[str, bool]] = "warn",
Philip Meier's avatar
Philip Meier committed
1899
) -> datapoints._InputTypeJIT:
1900
1901
1902
    if not torch.jit.is_scripting():
        _log_api_usage_once(resized_crop)

1903
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
1904
1905
1906
        return resized_crop_image_tensor(
            inpt, top, left, height, width, antialias=antialias, size=size, interpolation=interpolation
        )
1907
    elif isinstance(inpt, datapoints._datapoint.Datapoint):
1908
        return inpt.resized_crop(top, left, height, width, antialias=antialias, size=size, interpolation=interpolation)
1909
    elif isinstance(inpt, PIL.Image.Image):
1910
        return resized_crop_image_pil(inpt, top, left, height, width, size=size, interpolation=interpolation)
1911
1912
    else:
        raise TypeError(
1913
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
1914
1915
            f"but got {type(inpt)} instead."
        )
1916
1917


1918
1919
def _parse_five_crop_size(size: List[int]) -> List[int]:
    if isinstance(size, numbers.Number):
1920
1921
        s = int(size)
        size = [s, s]
1922
    elif isinstance(size, (tuple, list)) and len(size) == 1:
1923
1924
        s = size[0]
        size = [s, s]
1925
1926
1927
1928
1929
1930
1931
1932

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

    return size


def five_crop_image_tensor(
1933
    image: torch.Tensor, size: List[int]
1934
1935
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    crop_height, crop_width = _parse_five_crop_size(size)
1936
    image_height, image_width = image.shape[-2:]
1937
1938

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

1941
1942
1943
1944
1945
    tl = crop_image_tensor(image, 0, 0, crop_height, crop_width)
    tr = crop_image_tensor(image, 0, image_width - crop_width, crop_height, crop_width)
    bl = crop_image_tensor(image, image_height - crop_height, 0, crop_height, crop_width)
    br = crop_image_tensor(image, image_height - crop_height, image_width - crop_width, crop_height, crop_width)
    center = center_crop_image_tensor(image, [crop_height, crop_width])
1946
1947
1948
1949

    return tl, tr, bl, br, center


1950
@torch.jit.unused
1951
def five_crop_image_pil(
1952
    image: PIL.Image.Image, size: List[int]
1953
1954
) -> 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)
1955
    image_height, image_width = get_spatial_size_image_pil(image)
1956
1957

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

1960
1961
1962
1963
1964
    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])
1965
1966
1967
1968

    return tl, tr, bl, br, center


1969
1970
1971
1972
1973
1974
def five_crop_video(
    video: torch.Tensor, size: List[int]
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    return five_crop_image_tensor(video, size)


Philip Meier's avatar
Philip Meier committed
1975
ImageOrVideoTypeJIT = Union[datapoints._ImageTypeJIT, datapoints._VideoTypeJIT]
1976
1977


1978
def five_crop(
1979
1980
    inpt: ImageOrVideoTypeJIT, size: List[int]
) -> Tuple[ImageOrVideoTypeJIT, ImageOrVideoTypeJIT, ImageOrVideoTypeJIT, ImageOrVideoTypeJIT, ImageOrVideoTypeJIT]:
1981
1982
1983
    if not torch.jit.is_scripting():
        _log_api_usage_once(five_crop)

1984
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
1985
        return five_crop_image_tensor(inpt, size)
1986
    elif isinstance(inpt, datapoints.Image):
1987
        output = five_crop_image_tensor(inpt.as_subclass(torch.Tensor), size)
1988
1989
        return tuple(datapoints.Image.wrap_like(inpt, item) for item in output)  # type: ignore[return-value]
    elif isinstance(inpt, datapoints.Video):
1990
        output = five_crop_video(inpt.as_subclass(torch.Tensor), size)
1991
        return tuple(datapoints.Video.wrap_like(inpt, item) for item in output)  # type: ignore[return-value]
1992
    elif isinstance(inpt, PIL.Image.Image):
1993
        return five_crop_image_pil(inpt, size)
1994
1995
    else:
        raise TypeError(
1996
            f"Input can either be a plain tensor, an `Image` or `Video` datapoint, or a PIL image, "
1997
1998
            f"but got {type(inpt)} instead."
        )
1999
2000


Philip Meier's avatar
Philip Meier committed
2001
2002
2003
2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
def ten_crop_image_tensor(
    image: torch.Tensor, size: List[int], vertical_flip: bool = False
) -> Tuple[
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
]:
    non_flipped = five_crop_image_tensor(image, size)
2016
2017

    if vertical_flip:
2018
        image = vertical_flip_image_tensor(image)
2019
    else:
2020
        image = horizontal_flip_image_tensor(image)
2021

Philip Meier's avatar
Philip Meier committed
2022
    flipped = five_crop_image_tensor(image, size)
2023

Philip Meier's avatar
Philip Meier committed
2024
    return non_flipped + flipped
2025
2026


2027
@torch.jit.unused
Philip Meier's avatar
Philip Meier committed
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
2038
2039
2040
2041
2042
def ten_crop_image_pil(
    image: PIL.Image.Image, size: List[int], vertical_flip: bool = False
) -> Tuple[
    PIL.Image.Image,
    PIL.Image.Image,
    PIL.Image.Image,
    PIL.Image.Image,
    PIL.Image.Image,
    PIL.Image.Image,
    PIL.Image.Image,
    PIL.Image.Image,
    PIL.Image.Image,
    PIL.Image.Image,
]:
    non_flipped = five_crop_image_pil(image, size)
2043
2044

    if vertical_flip:
2045
        image = vertical_flip_image_pil(image)
2046
    else:
2047
        image = horizontal_flip_image_pil(image)
2048

Philip Meier's avatar
Philip Meier committed
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
2067
    flipped = five_crop_image_pil(image, size)

    return non_flipped + flipped


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,
]:
2068
2069
2070
2071
    return ten_crop_image_tensor(video, size, vertical_flip=vertical_flip)


def ten_crop(
Philip Meier's avatar
Philip Meier committed
2072
    inpt: Union[datapoints._ImageTypeJIT, datapoints._VideoTypeJIT], size: List[int], vertical_flip: bool = False
Philip Meier's avatar
Philip Meier committed
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
) -> Tuple[
    ImageOrVideoTypeJIT,
    ImageOrVideoTypeJIT,
    ImageOrVideoTypeJIT,
    ImageOrVideoTypeJIT,
    ImageOrVideoTypeJIT,
    ImageOrVideoTypeJIT,
    ImageOrVideoTypeJIT,
    ImageOrVideoTypeJIT,
    ImageOrVideoTypeJIT,
    ImageOrVideoTypeJIT,
]:
2085
2086
2087
    if not torch.jit.is_scripting():
        _log_api_usage_once(ten_crop)

2088
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
2089
        return ten_crop_image_tensor(inpt, size, vertical_flip=vertical_flip)
2090
    elif isinstance(inpt, datapoints.Image):
2091
        output = ten_crop_image_tensor(inpt.as_subclass(torch.Tensor), size, vertical_flip=vertical_flip)
2092
        return tuple(datapoints.Image.wrap_like(inpt, item) for item in output)  # type: ignore[return-value]
2093
    elif isinstance(inpt, datapoints.Video):
2094
        output = ten_crop_video(inpt.as_subclass(torch.Tensor), size, vertical_flip=vertical_flip)
2095
        return tuple(datapoints.Video.wrap_like(inpt, item) for item in output)  # type: ignore[return-value]
2096
    elif isinstance(inpt, PIL.Image.Image):
2097
        return ten_crop_image_pil(inpt, size, vertical_flip=vertical_flip)
2098
2099
    else:
        raise TypeError(
2100
            f"Input can either be a plain tensor, an `Image` or `Video` datapoint, or a PIL image, "
2101
2102
            f"but got {type(inpt)} instead."
        )