functional_tensor.py 37.5 KB
Newer Older
vfdev's avatar
vfdev committed
1
import warnings
vfdev's avatar
vfdev committed
2
from typing import Optional, Dict, Tuple
vfdev's avatar
vfdev committed
3

4
import torch
5
from torch import Tensor
6
from torch.nn.functional import grid_sample
7
from torch.jit.annotations import List, BroadcastingList2
8
9


vfdev's avatar
vfdev committed
10
11
def _is_tensor_a_torch_image(x: Tensor) -> bool:
    return x.ndim >= 2
12
13


vfdev's avatar
vfdev committed
14
def _get_image_size(img: Tensor) -> List[int]:
vfdev's avatar
vfdev committed
15
    """Returns (w, h) of tensor image"""
vfdev's avatar
vfdev committed
16
17
18
19
20
    if _is_tensor_a_torch_image(img):
        return [img.shape[-1], img.shape[-2]]
    raise TypeError("Unexpected type {}".format(type(img)))


21
22
23
24
25
26
27
28
29
def _get_image_num_channels(img: Tensor) -> int:
    if img.ndim == 2:
        return 1
    elif img.ndim > 2:
        return img.shape[-3]

    raise TypeError("Unexpected type {}".format(type(img)))


vfdev's avatar
vfdev committed
30
def vflip(img: Tensor) -> Tensor:
31
32
33
34
35
36
    """PRIVATE METHOD. Vertically flip the given the Image Tensor.

    .. warning::

        Module ``transforms.functional_tensor`` is private and should not be used in user application.
        Please, consider instead using methods from `transforms.functional` module.
37
38

    Args:
39
        img (Tensor): Image Tensor to be flipped in the form [..., C, H, W].
40
41
42
43

    Returns:
        Tensor:  Vertically flipped image Tensor.
    """
44
    if not _is_tensor_a_torch_image(img):
45
46
        raise TypeError('tensor is not a torch image.')

47
    return img.flip(-2)
48
49


vfdev's avatar
vfdev committed
50
def hflip(img: Tensor) -> Tensor:
51
52
53
54
55
56
    """PRIVATE METHOD. Horizontally flip the given the Image Tensor.

    .. warning::

        Module ``transforms.functional_tensor`` is private and should not be used in user application.
        Please, consider instead using methods from `transforms.functional` module.
57
58

    Args:
59
        img (Tensor): Image Tensor to be flipped in the form [..., C, H, W].
60
61
62
63

    Returns:
        Tensor:  Horizontally flipped image Tensor.
    """
64
    if not _is_tensor_a_torch_image(img):
65
66
        raise TypeError('tensor is not a torch image.')

67
    return img.flip(-1)
ekka's avatar
ekka committed
68
69


vfdev's avatar
vfdev committed
70
def crop(img: Tensor, top: int, left: int, height: int, width: int) -> Tensor:
71
72
73
74
75
76
    """PRIVATE METHOD. Crop the given Image Tensor.

    .. warning::

        Module ``transforms.functional_tensor`` is private and should not be used in user application.
        Please, consider instead using methods from `transforms.functional` module.
77

ekka's avatar
ekka committed
78
    Args:
vfdev's avatar
vfdev committed
79
        img (Tensor): Image to be cropped in the form [..., H, W]. (0,0) denotes the top left corner of the image.
ekka's avatar
ekka committed
80
81
82
83
        top (int): Vertical component of the top left corner of the crop box.
        left (int): Horizontal component of the top left corner of the crop box.
        height (int): Height of the crop box.
        width (int): Width of the crop box.
84

ekka's avatar
ekka committed
85
86
87
    Returns:
        Tensor: Cropped image.
    """
88
    if not _is_tensor_a_torch_image(img):
vfdev's avatar
vfdev committed
89
        raise TypeError("tensor is not a torch image.")
ekka's avatar
ekka committed
90
91

    return img[..., top:top + height, left:left + width]
92
93


94
def rgb_to_grayscale(img: Tensor, num_output_channels: int = 1) -> Tensor:
95
96
97
98
99
100
101
    """PRIVATE METHOD. Convert the given RGB Image Tensor to Grayscale.

    .. warning::

        Module ``transforms.functional_tensor`` is private and should not be used in user application.
        Please, consider instead using methods from `transforms.functional` module.

102
103
104
105
106
    For RGB to Grayscale conversion, ITU-R 601-2 luma transform is performed which
    is L = R * 0.2989 + G * 0.5870 + B * 0.1140

    Args:
        img (Tensor): Image to be converted to Grayscale in the form [C, H, W].
107
        num_output_channels (int): number of channels of the output image. Value can be 1 or 3. Default, 1.
108
109

    Returns:
110
111
112
113
        Tensor: Grayscale version of the image.
            if num_output_channels = 1 : returned image is single channel

            if num_output_channels = 3 : returned image is 3 channel with r = g = b
114
115

    """
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
    if img.ndim < 3:
        raise TypeError("Input image tensor should have at least 3 dimensions, but found {}".format(img.ndim))
    c = img.shape[-3]
    if c != 3:
        raise TypeError("Input image tensor should 3 channels, but found {}".format(c))

    if num_output_channels not in (1, 3):
        raise ValueError('num_output_channels should be either 1 or 3')

    r, g, b = img.unbind(dim=-3)
    # This implementation closely follows the TF one:
    # https://github.com/tensorflow/tensorflow/blob/v2.3.0/tensorflow/python/ops/image_ops_impl.py#L2105-L2138
    l_img = (0.2989 * r + 0.587 * g + 0.114 * b).to(img.dtype)
    l_img = l_img.unsqueeze(dim=-3)

    if num_output_channels == 3:
        return l_img.expand(img.shape)
133

134
    return l_img
135
136


vfdev's avatar
vfdev committed
137
def adjust_brightness(img: Tensor, brightness_factor: float) -> Tensor:
138
139
140
141
142
143
    """PRIVATE METHOD. Adjust brightness of an RGB image.

    .. warning::

        Module ``transforms.functional_tensor`` is private and should not be used in user application.
        Please, consider instead using methods from `transforms.functional` module.
144
145
146
147
148
149
150
151
152
153

    Args:
        img (Tensor): Image to be adjusted.
        brightness_factor (float):  How much to adjust the brightness. Can be
            any non negative number. 0 gives a black image, 1 gives the
            original image while 2 increases the brightness by a factor of 2.

    Returns:
        Tensor: Brightness adjusted image.
    """
154
155
156
    if brightness_factor < 0:
        raise ValueError('brightness_factor ({}) is not non-negative.'.format(brightness_factor))

157
    if not _is_tensor_a_torch_image(img):
158
159
        raise TypeError('tensor is not a torch image.')

160
    return _blend(img, torch.zeros_like(img), brightness_factor)
161
162


vfdev's avatar
vfdev committed
163
def adjust_contrast(img: Tensor, contrast_factor: float) -> Tensor:
164
165
166
167
168
169
    """PRIVATE METHOD. Adjust contrast of an RGB image.

    .. warning::

        Module ``transforms.functional_tensor`` is private and should not be used in user application.
        Please, consider instead using methods from `transforms.functional` module.
170
171
172
173
174
175
176
177
178
179

    Args:
        img (Tensor): Image to be adjusted.
        contrast_factor (float): How much to adjust the contrast. Can be any
            non negative number. 0 gives a solid gray image, 1 gives the
            original image while 2 increases the contrast by a factor of 2.

    Returns:
        Tensor: Contrast adjusted image.
    """
180
181
182
    if contrast_factor < 0:
        raise ValueError('contrast_factor ({}) is not non-negative.'.format(contrast_factor))

183
    if not _is_tensor_a_torch_image(img):
184
185
        raise TypeError('tensor is not a torch image.')

186
187
    dtype = img.dtype if torch.is_floating_point(img) else torch.float32
    mean = torch.mean(rgb_to_grayscale(img).to(dtype), dim=(-3, -2, -1), keepdim=True)
188
189
190
191

    return _blend(img, mean, contrast_factor)


192
def adjust_hue(img: Tensor, hue_factor: float) -> Tensor:
193
194
195
196
197
198
    """PRIVATE METHOD. Adjust hue of an image.

    .. warning::

        Module ``transforms.functional_tensor`` is private and should not be used in user application.
        Please, consider instead using methods from `transforms.functional` module.
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221

    The image hue is adjusted by converting the image to HSV and
    cyclically shifting the intensities in the hue channel (H).
    The image is then converted back to original image mode.

    `hue_factor` is the amount of shift in H channel and must be in the
    interval `[-0.5, 0.5]`.

    See `Hue`_ for more details.

    .. _Hue: https://en.wikipedia.org/wiki/Hue

    Args:
        img (Tensor): Image to be adjusted. Image type is either uint8 or float.
        hue_factor (float):  How much to shift the hue channel. Should be in
            [-0.5, 0.5]. 0.5 and -0.5 give complete reversal of hue channel in
            HSV space in positive and negative direction respectively.
            0 means no shift. Therefore, both -0.5 and 0.5 will give an image
            with complementary colors while 0 gives the original image.

    Returns:
         Tensor: Hue adjusted image.
    """
222
    if not (-0.5 <= hue_factor <= 0.5):
223
224
        raise ValueError('hue_factor ({}) is not in [-0.5, 0.5].'.format(hue_factor))

225
226
    if not (isinstance(img, torch.Tensor) and _is_tensor_a_torch_image(img)):
        raise TypeError('img should be Tensor image. Got {}'.format(type(img)))
227
228
229
230
231
232

    orig_dtype = img.dtype
    if img.dtype == torch.uint8:
        img = img.to(dtype=torch.float32) / 255.0

    img = _rgb2hsv(img)
233
    h, s, v = img.unbind(dim=-3)
234
    h = (h + hue_factor) % 1.0
235
    img = torch.stack((h, s, v), dim=-3)
236
237
238
239
240
241
242
243
    img_hue_adj = _hsv2rgb(img)

    if orig_dtype == torch.uint8:
        img_hue_adj = (img_hue_adj * 255.0).to(dtype=orig_dtype)

    return img_hue_adj


vfdev's avatar
vfdev committed
244
def adjust_saturation(img: Tensor, saturation_factor: float) -> Tensor:
245
246
247
248
249
250
    """PRIVATE METHOD. Adjust color saturation of an RGB image.

    .. warning::

        Module ``transforms.functional_tensor`` is private and should not be used in user application.
        Please, consider instead using methods from `transforms.functional` module.
251
252
253

    Args:
        img (Tensor): Image to be adjusted.
254
255
256
        saturation_factor (float):  How much to adjust the saturation. Can be any
            non negative number. 0 gives a black and white image, 1 gives the
            original image while 2 enhances the saturation by a factor of 2.
257
258
259
260

    Returns:
        Tensor: Saturation adjusted image.
    """
261
262
263
    if saturation_factor < 0:
        raise ValueError('saturation_factor ({}) is not non-negative.'.format(saturation_factor))

264
    if not _is_tensor_a_torch_image(img):
265
266
        raise TypeError('tensor is not a torch image.')

267
    return _blend(img, rgb_to_grayscale(img), saturation_factor)
268
269


270
def adjust_gamma(img: Tensor, gamma: float, gain: float = 1) -> Tensor:
271
272
273
274
275
276
    r"""PRIVATE METHOD. Adjust gamma of an RGB image.

    .. warning::

        Module ``transforms.functional_tensor`` is private and should not be used in user application.
        Please, consider instead using methods from `transforms.functional` module.
277
278
279
280
281
282
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

    Also known as Power Law Transform. Intensities in RGB mode are adjusted
    based on the following equation:

    .. math::
        `I_{\text{out}} = 255 \times \text{gain} \times \left(\frac{I_{\text{in}}}{255}\right)^{\gamma}`

    See `Gamma Correction`_ for more details.

    .. _Gamma Correction: https://en.wikipedia.org/wiki/Gamma_correction

    Args:
        img (Tensor): Tensor of RBG values to be adjusted.
        gamma (float): Non negative real number, same as :math:`\gamma` in the equation.
            gamma larger than 1 make the shadows darker,
            while gamma smaller than 1 make dark regions lighter.
        gain (float): The constant multiplier.
    """

    if not isinstance(img, torch.Tensor):
        raise TypeError('img should be a Tensor. Got {}'.format(type(img)))

    if gamma < 0:
        raise ValueError('Gamma should be a non-negative real number')

    result = img
    dtype = img.dtype
    if not torch.is_floating_point(img):
        result = result / 255.0

    result = (gain * result ** gamma).clamp(0, 1)

    if result.dtype != dtype:
        eps = 1e-3
        result = (255 + 1.0 - eps) * result
    result = result.to(dtype)
    return result


vfdev's avatar
vfdev committed
316
def center_crop(img: Tensor, output_size: BroadcastingList2[int]) -> Tensor:
317
318
    """DEPRECATED. Crop the Image Tensor and resize it to desired size.

319
320
321
322
323
    .. warning::

        Module ``transforms.functional_tensor`` is private and should not be used in user application.
        Please, consider instead using methods from `transforms.functional` module.

324
325
326
327
    .. warning::

        This method is deprecated and will be removed in future releases.
        Please, use ``F.center_crop`` instead.
328
329

    Args:
vfdev's avatar
vfdev committed
330
        img (Tensor): Image to be cropped.
331
332
333
334
335
336
        output_size (sequence or int): (height, width) of the crop box. If int,
                it is used for both directions

    Returns:
            Tensor: Cropped image.
    """
337
338
339
340
341
    warnings.warn(
        "This method is deprecated and will be removed in future releases. "
        "Please, use ``F.center_crop`` instead."
    )

342
    if not _is_tensor_a_torch_image(img):
343
344
345
346
        raise TypeError('tensor is not a torch image.')

    _, image_width, image_height = img.size()
    crop_height, crop_width = output_size
vfdev's avatar
vfdev committed
347
348
349
350
351
352
353
354
    # crop_top = int(round((image_height - crop_height) / 2.))
    # Result can be different between python func and scripted func
    # Temporary workaround:
    crop_top = int((image_height - crop_height + 1) * 0.5)
    # crop_left = int(round((image_width - crop_width) / 2.))
    # Result can be different between python func and scripted func
    # Temporary workaround:
    crop_left = int((image_width - crop_width + 1) * 0.5)
355
356
357
358

    return crop(img, crop_top, crop_left, crop_height, crop_width)


vfdev's avatar
vfdev committed
359
def five_crop(img: Tensor, size: BroadcastingList2[int]) -> List[Tensor]:
360
361
    """DEPRECATED. Crop the given Image Tensor into four corners and the central crop.

362
363
364
365
366
    .. warning::

        Module ``transforms.functional_tensor`` is private and should not be used in user application.
        Please, consider instead using methods from `transforms.functional` module.

367
368
369
370
371
    .. warning::

        This method is deprecated and will be removed in future releases.
        Please, use ``F.five_crop`` instead.

372
    .. Note::
373

374
        This transform returns a List of Tensors and there may be a
375
376
377
        mismatch in the number of inputs and targets your ``Dataset`` returns.

    Args:
vfdev's avatar
vfdev committed
378
379
380
381
        img (Tensor): Image to be cropped.
        size (sequence or int): Desired output size of the crop. If size is an
            int instead of sequence like (h, w), a square crop (size, size) is
            made.
382
383

    Returns:
384
       List: List (tl, tr, bl, br, center)
385
386
                Corresponding top left, top right, bottom left, bottom right and center crop.
    """
387
388
389
390
391
    warnings.warn(
        "This method is deprecated and will be removed in future releases. "
        "Please, use ``F.five_crop`` instead."
    )

392
    if not _is_tensor_a_torch_image(img):
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
        raise TypeError('tensor is not a torch image.')

    assert len(size) == 2, "Please provide only two dimensions (h, w) for size."

    _, image_width, image_height = img.size()
    crop_height, crop_width = size
    if crop_width > image_width or crop_height > image_height:
        msg = "Requested crop size {} is bigger than input size {}"
        raise ValueError(msg.format(size, (image_height, image_width)))

    tl = crop(img, 0, 0, crop_width, crop_height)
    tr = crop(img, image_width - crop_width, 0, image_width, crop_height)
    bl = crop(img, 0, image_height - crop_height, crop_width, image_height)
    br = crop(img, image_width - crop_width, image_height - crop_height, image_width, image_height)
    center = center_crop(img, (crop_height, crop_width))

409
    return [tl, tr, bl, br, center]
410
411


vfdev's avatar
vfdev committed
412
def ten_crop(img: Tensor, size: BroadcastingList2[int], vertical_flip: bool = False) -> List[Tensor]:
413
    """DEPRECATED. Crop the given Image Tensor into four corners and the central crop plus the
414
        flipped version of these (horizontal flipping is used by default).
vfdev's avatar
vfdev committed
415

416
417
418
419
420
    .. warning::

        Module ``transforms.functional_tensor`` is private and should not be used in user application.
        Please, consider instead using methods from `transforms.functional` module.

421
422
423
424
425
    .. warning::

        This method is deprecated and will be removed in future releases.
        Please, use ``F.ten_crop`` instead.

426
    .. Note::
427

428
        This transform returns a List of images and there may be a
429
430
431
        mismatch in the number of inputs and targets your ``Dataset`` returns.

    Args:
vfdev's avatar
vfdev committed
432
433
        img (Tensor): Image to be cropped.
        size (sequence or int): Desired output size of the crop. If size is an
434
435
            int instead of sequence like (h, w), a square crop (size, size) is
            made.
vfdev's avatar
vfdev committed
436
        vertical_flip (bool): Use vertical flipping instead of horizontal
437
438

    Returns:
439
       List: List (tl, tr, bl, br, center, tl_flip, tr_flip, bl_flip, br_flip, center_flip)
440
441
442
                Corresponding top left, top right, bottom left, bottom right and center crop
                and same for the flipped image's tensor.
    """
443
444
445
446
447
    warnings.warn(
        "This method is deprecated and will be removed in future releases. "
        "Please, use ``F.ten_crop`` instead."
    )

448
    if not _is_tensor_a_torch_image(img):
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
        raise TypeError('tensor is not a torch image.')

    assert len(size) == 2, "Please provide only two dimensions (h, w) for size."
    first_five = five_crop(img, size)

    if vertical_flip:
        img = vflip(img)
    else:
        img = hflip(img)

    second_five = five_crop(img, size)

    return first_five + second_five


vfdev's avatar
vfdev committed
464
def _blend(img1: Tensor, img2: Tensor, ratio: float) -> Tensor:
465
466
    bound = 1.0 if img1.is_floating_point() else 255.0
    return (ratio * img1 + (1.0 - ratio) * img2).clamp(0, bound).to(img1.dtype)
467
468
469


def _rgb2hsv(img):
470
    r, g, b = img.unbind(dim=-3)
471

472
473
    # Implementation is based on https://github.com/python-pillow/Pillow/blob/4174d4267616897df3746d315d5a2d0f82c656ee/
    # src/libImaging/Convert.c#L330
474
475
    maxc = torch.max(img, dim=-3).values
    minc = torch.min(img, dim=-3).values
476
477
478
479
480
481
482
483
484
485

    # The algorithm erases S and H channel where `maxc = minc`. This avoids NaN
    # from happening in the results, because
    #   + S channel has division by `maxc`, which is zero only if `maxc = minc`
    #   + H channel has division by `(maxc - minc)`.
    #
    # Instead of overwriting NaN afterwards, we just prevent it from occuring so
    # we don't need to deal with it in case we save the NaN in a buffer in
    # backprop, if it is ever supported, but it doesn't hurt to do so.
    eqc = maxc == minc
486
487

    cr = maxc - minc
488
    # Since `eqc => cr = 0`, replacing denominator with 1 when `eqc` is fine.
489
490
    ones = torch.ones_like(maxc)
    s = cr / torch.where(eqc, ones, maxc)
491
492
493
494
    # Note that `eqc => maxc = minc = r = g = b`. So the following calculation
    # of `h` would reduce to `bc - gc + 2 + rc - bc + 4 + rc - bc = 6` so it
    # would not matter what values `rc`, `gc`, and `bc` have here, and thus
    # replacing denominator with 1 when `eqc` is fine.
495
    cr_divisor = torch.where(eqc, ones, cr)
496
497
498
    rc = (maxc - r) / cr_divisor
    gc = (maxc - g) / cr_divisor
    bc = (maxc - b) / cr_divisor
499
500
501
502
503
504

    hr = (maxc == r) * (bc - gc)
    hg = ((maxc == g) & (maxc != r)) * (2.0 + rc - bc)
    hb = ((maxc != g) & (maxc != r)) * (4.0 + gc - rc)
    h = (hr + hg + hb)
    h = torch.fmod((h / 6.0 + 1.0), 1.0)
505
    return torch.stack((h, s, maxc), dim=-3)
506
507
508


def _hsv2rgb(img):
509
    h, s, v = img.unbind(dim=-3)
510
511
512
513
514
515
516
517
518
    i = torch.floor(h * 6.0)
    f = (h * 6.0) - i
    i = i.to(dtype=torch.int32)

    p = torch.clamp((v * (1.0 - s)), 0.0, 1.0)
    q = torch.clamp((v * (1.0 - s * f)), 0.0, 1.0)
    t = torch.clamp((v * (1.0 - s * (1.0 - f))), 0.0, 1.0)
    i = i % 6

519
    mask = i.unsqueeze(dim=-3) == torch.arange(6, device=i.device).view(-1, 1, 1)
520

521
522
523
524
    a1 = torch.stack((v, q, p, p, t, v), dim=-3)
    a2 = torch.stack((t, v, v, q, p, p), dim=-3)
    a3 = torch.stack((p, p, t, v, v, q), dim=-3)
    a4 = torch.stack((a1, a2, a3), dim=-4)
525

526
    return torch.einsum("...ijk, ...xijk -> ...xjk", mask.to(dtype=img.dtype), a4)
527
528


529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
def _pad_symmetric(img: Tensor, padding: List[int]) -> Tensor:
    # padding is left, right, top, bottom
    in_sizes = img.size()

    x_indices = [i for i in range(in_sizes[-1])]  # [0, 1, 2, 3, ...]
    left_indices = [i for i in range(padding[0] - 1, -1, -1)]  # e.g. [3, 2, 1, 0]
    right_indices = [-(i + 1) for i in range(padding[1])]  # e.g. [-1, -2, -3]
    x_indices = torch.tensor(left_indices + x_indices + right_indices)

    y_indices = [i for i in range(in_sizes[-2])]
    top_indices = [i for i in range(padding[2] - 1, -1, -1)]
    bottom_indices = [-(i + 1) for i in range(padding[3])]
    y_indices = torch.tensor(top_indices + y_indices + bottom_indices)

    ndim = img.ndim
    if ndim == 3:
        return img[:, y_indices[:, None], x_indices[None, :]]
    elif ndim == 4:
        return img[:, :, y_indices[:, None], x_indices[None, :]]
    else:
        raise RuntimeError("Symmetric padding of N-D tensors are not supported yet")


552
def pad(img: Tensor, padding: List[int], fill: int = 0, padding_mode: str = "constant") -> Tensor:
553
554
555
556
557
558
    r"""PRIVATE METHOD. Pad the given Tensor Image on all sides with specified padding mode and fill value.

    .. warning::

        Module ``transforms.functional_tensor`` is private and should not be used in user application.
        Please, consider instead using methods from `transforms.functional` module.
559
560
561
562
563
564
565
566
567
568
569

    Args:
        img (Tensor): Image to be padded.
        padding (int or tuple or list): Padding on each border. If a single int is provided this
            is used to pad all borders. If a tuple or list of length 2 is provided this is the padding
            on left/right and top/bottom respectively. If a tuple or list of length 4 is provided
            this is the padding for the left, top, right and bottom borders
            respectively. In torchscript mode padding as single int is not supported, use a tuple or
            list of length 1: ``[padding, ]``.
        fill (int): Pixel fill value for constant fill. Default is 0.
            This value is only used when the padding_mode is constant
vfdev's avatar
vfdev committed
570
571
        padding_mode (str): Type of padding. Should be: constant, edge or reflect. Default is constant.
            Mode symmetric is not yet supported for Tensor inputs.
572
573
574

            - constant: pads with a constant value, this value is specified with fill

575
576
577
578
579
580
581
            - edge: pads with the last value on the edge of the image

            - reflect: pads with reflection of image (without repeating the last value on the edge)

                       padding [1, 2, 3, 4] with 2 elements on both sides in reflect mode
                       will result in [3, 2, 1, 2, 3, 4, 3, 2]

582
583
584
585
586
            - symmetric: pads with reflection of image (repeating the last value on the edge)

                         padding [1, 2, 3, 4] with 2 elements on both sides in symmetric mode
                         will result in [2, 1, 1, 2, 3, 4, 4, 3]

587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
    Returns:
        Tensor: Padded image.
    """
    if not _is_tensor_a_torch_image(img):
        raise TypeError("tensor is not a torch image.")

    if not isinstance(padding, (int, tuple, list)):
        raise TypeError("Got inappropriate padding arg")
    if not isinstance(fill, (int, float)):
        raise TypeError("Got inappropriate fill arg")
    if not isinstance(padding_mode, str):
        raise TypeError("Got inappropriate padding_mode arg")

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

    if isinstance(padding, list) and len(padding) not in [1, 2, 4]:
        raise ValueError("Padding must be an int or a 1, 2, or 4 element tuple, not a " +
                         "{} element tuple".format(len(padding)))

607
608
    if padding_mode not in ["constant", "edge", "reflect", "symmetric"]:
        raise ValueError("Padding mode should be either constant, edge, reflect or symmetric")
609
610
611

    if isinstance(padding, int):
        if torch.jit.is_scripting():
vfdev's avatar
vfdev committed
612
            # This maybe unreachable
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
            raise ValueError("padding can't be an int while torchscripting, set it as a list [value, ]")
        pad_left = pad_right = pad_top = pad_bottom = padding
    elif 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]
    else:
        pad_left = padding[0]
        pad_top = padding[1]
        pad_right = padding[2]
        pad_bottom = padding[3]

    p = [pad_left, pad_right, pad_top, pad_bottom]

628
629
630
    if padding_mode == "edge":
        # remap padding_mode str
        padding_mode = "replicate"
631
632
633
634
635
    elif padding_mode == "symmetric":
        # route to another implementation
        if p[0] < 0 or p[1] < 0 or p[2] < 0 or p[3] < 0:  # no any support for torch script
            raise ValueError("Padding can not be negative for symmetric padding_mode")
        return _pad_symmetric(img, p)
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650

    need_squeeze = False
    if img.ndim < 4:
        img = img.unsqueeze(dim=0)
        need_squeeze = True

    out_dtype = img.dtype
    need_cast = False
    if (padding_mode != "constant") and img.dtype not in (torch.float32, torch.float64):
        # Here we temporary cast input tensor to float
        # until pytorch issue is resolved :
        # https://github.com/pytorch/pytorch/issues/40763
        need_cast = True
        img = img.to(torch.float32)

651
    img = torch.nn.functional.pad(img, p, mode=padding_mode, value=float(fill))
652
653
654
655
656
657
658

    if need_squeeze:
        img = img.squeeze(dim=0)

    if need_cast:
        img = img.to(out_dtype)

659
    return img
vfdev's avatar
vfdev committed
660
661
662


def resize(img: Tensor, size: List[int], interpolation: int = 2) -> Tensor:
663
664
665
666
667
668
    r"""PRIVATE METHOD. Resize the input Tensor to the given size.

    .. warning::

        Module ``transforms.functional_tensor`` is private and should not be used in user application.
        Please, consider instead using methods from `transforms.functional` module.
vfdev's avatar
vfdev committed
669
670
671
672
673
674
675
676
677
678

    Args:
        img (Tensor): Image to be resized.
        size (int or tuple or list): Desired output size. If size is a sequence like
            (h, w), the output size will be matched to this. If size is an int,
            the smaller edge of the image will be matched to this number maintaining
            the aspect ratio. i.e, if height > width, then image will be rescaled to
            :math:`\left(\text{size} \times \frac{\text{height}}{\text{width}}, \text{size}\right)`.
            In torchscript mode padding as a single int is not supported, use a tuple or
            list of length 1: ``[size, ]``.
vfdev's avatar
vfdev committed
679
680
        interpolation (int, optional): Desired interpolation. Default is bilinear (=2). Other supported values:
            nearest(=0) and bicubic(=3).
vfdev's avatar
vfdev committed
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715

    Returns:
        Tensor: Resized image.
    """
    if not _is_tensor_a_torch_image(img):
        raise TypeError("tensor is not a torch image.")

    if not isinstance(size, (int, tuple, list)):
        raise TypeError("Got inappropriate size arg")
    if not isinstance(interpolation, int):
        raise TypeError("Got inappropriate interpolation arg")

    _interpolation_modes = {
        0: "nearest",
        2: "bilinear",
        3: "bicubic",
    }

    if interpolation not in _interpolation_modes:
        raise ValueError("This interpolation mode is unsupported with Tensor input")

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

    if isinstance(size, list) and len(size) not in [1, 2]:
        raise ValueError("Size must be an int or a 1 or 2 element tuple/list, not a "
                         "{} element tuple/list".format(len(size)))

    w, h = _get_image_size(img)

    if isinstance(size, int):
        size_w, size_h = size, size
    elif len(size) < 2:
        size_w, size_h = size[0], size[0]
    else:
716
        size_w, size_h = size[1], size[0]  # Convention (h, w)
vfdev's avatar
vfdev committed
717
718
719
720
721
722
723

    if isinstance(size, int) or len(size) < 2:
        if w < h:
            size_h = int(size_w * h / w)
        else:
            size_w = int(size_h * w / h)

724
725
        if (w <= h and w == size_w) or (h <= w and h == size_h):
            return img
vfdev's avatar
vfdev committed
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754

    # make image NCHW
    need_squeeze = False
    if img.ndim < 4:
        img = img.unsqueeze(dim=0)
        need_squeeze = True

    mode = _interpolation_modes[interpolation]

    out_dtype = img.dtype
    need_cast = False
    if img.dtype not in (torch.float32, torch.float64):
        need_cast = True
        img = img.to(torch.float32)

    # Define align_corners to avoid warnings
    align_corners = False if mode in ["bilinear", "bicubic"] else None

    img = torch.nn.functional.interpolate(img, size=(size_h, size_w), mode=mode, align_corners=align_corners)

    if need_squeeze:
        img = img.squeeze(dim=0)

    if need_cast:
        if mode == "bicubic":
            img = img.clamp(min=0, max=255)
        img = img.to(out_dtype)

    return img
vfdev's avatar
vfdev committed
755
756


vfdev's avatar
vfdev committed
757
def _assert_grid_transform_inputs(
758
759
760
761
762
763
        img: Tensor,
        matrix: Optional[List[float]],
        resample: int,
        fillcolor: Optional[int],
        _interpolation_modes: Dict[int, str],
        coeffs: Optional[List[float]] = None,
vfdev's avatar
vfdev committed
764
765
766
):
    if not (isinstance(img, torch.Tensor) and _is_tensor_a_torch_image(img)):
        raise TypeError("img should be Tensor Image. Got {}".format(type(img)))
vfdev's avatar
vfdev committed
767

768
    if matrix is not None and not isinstance(matrix, list):
vfdev's avatar
vfdev committed
769
        raise TypeError("Argument matrix should be a list. Got {}".format(type(matrix)))
vfdev's avatar
vfdev committed
770

771
    if matrix is not None and len(matrix) != 6:
vfdev's avatar
vfdev committed
772
        raise ValueError("Argument matrix should have 6 float values")
vfdev's avatar
vfdev committed
773

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

vfdev's avatar
vfdev committed
777
    if fillcolor is not None:
vfdev's avatar
vfdev committed
778
        warnings.warn("Argument fill/fillcolor is not supported for Tensor input. Fill value is zero")
vfdev's avatar
vfdev committed
779
780

    if resample not in _interpolation_modes:
781
        raise ValueError("Resampling mode '{}' is unsupported with Tensor input".format(resample))
vfdev's avatar
vfdev committed
782
783


vfdev's avatar
vfdev committed
784
def _apply_grid_transform(img: Tensor, grid: Tensor, mode: str) -> Tensor:
vfdev's avatar
vfdev committed
785
786
787
788
789
790
791
792
    # make image NCHW
    need_squeeze = False
    if img.ndim < 4:
        img = img.unsqueeze(dim=0)
        need_squeeze = True

    out_dtype = img.dtype
    need_cast = False
793
    if out_dtype != grid.dtype:
vfdev's avatar
vfdev committed
794
        need_cast = True
795
        img = img.to(grid)
vfdev's avatar
vfdev committed
796

797
798
799
    if img.shape[0] > 1:
        # Apply same grid to a batch of images
        grid = grid.expand(img.shape[0], grid.shape[1], grid.shape[2], grid.shape[3])
vfdev's avatar
vfdev committed
800
801
802
803
804
805
806
807
808
809
    img = grid_sample(img, grid, mode=mode, padding_mode="zeros", align_corners=False)

    if need_squeeze:
        img = img.squeeze(dim=0)

    if need_cast:
        # it is better to round before cast
        img = torch.round(img).to(out_dtype)

    return img
vfdev's avatar
vfdev committed
810
811


812
813
814
815
816
817
818
819
820
821
def _gen_affine_grid(
        theta: Tensor, w: int, h: int, ow: int, oh: int,
) -> 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

    d = 0.5
822
    base_grid = torch.empty(1, oh, ow, 3, dtype=theta.dtype, device=theta.device)
823
824
825
826
    base_grid[..., 0].copy_(torch.linspace(-ow * 0.5 + d, ow * 0.5 + d - 1, steps=ow))
    base_grid[..., 1].copy_(torch.linspace(-oh * 0.5 + d, oh * 0.5 + d - 1, steps=oh).unsqueeze_(-1))
    base_grid[..., 2].fill_(1)

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


vfdev's avatar
vfdev committed
832
833
834
def affine(
        img: Tensor, matrix: List[float], resample: int = 0, fillcolor: Optional[int] = None
) -> Tensor:
835
836
837
838
839
840
    """PRIVATE METHOD. Apply affine transformation on the Tensor image keeping image center invariant.

    .. warning::

        Module ``transforms.functional_tensor`` is private and should not be used in user application.
        Please, consider instead using methods from `transforms.functional` module.
vfdev's avatar
vfdev committed
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859

    Args:
        img (Tensor): image to be rotated.
        matrix (list of floats): list of 6 float values representing inverse matrix for affine transformation.
        resample (int, optional): An optional resampling filter. Default is nearest (=0). Other supported values:
            bilinear(=2).
        fillcolor (int, optional): this option is not supported for Tensor input. Fill value for the area outside the
            transform in the output image is always 0.

    Returns:
        Tensor: Transformed image.
    """
    _interpolation_modes = {
        0: "nearest",
        2: "bilinear",
    }

    _assert_grid_transform_inputs(img, matrix, resample, fillcolor, _interpolation_modes)

860
861
    dtype = img.dtype if torch.is_floating_point(img) else torch.float32
    theta = torch.tensor(matrix, dtype=dtype, device=img.device).reshape(1, 2, 3)
vfdev's avatar
vfdev committed
862
    shape = img.shape
863
    # grid will be generated on the same device as theta and img
864
    grid = _gen_affine_grid(theta, w=shape[-1], h=shape[-2], ow=shape[-1], oh=shape[-2])
vfdev's avatar
vfdev committed
865
866
867
868
    mode = _interpolation_modes[resample]
    return _apply_grid_transform(img, grid, mode)


869
def _compute_output_size(matrix: List[float], w: int, h: int) -> Tuple[int, int]:
vfdev's avatar
vfdev committed
870

871
872
873
    # Inspired of PIL implementation:
    # https://github.com/python-pillow/Pillow/blob/11de3318867e4398057373ee9f12dcb33db7335c/src/PIL/Image.py#L2054

vfdev's avatar
vfdev committed
874
875
    # pts are Top-Left, Top-Right, Bottom-Left, Bottom-Right points.
    pts = torch.tensor([
876
877
878
879
        [-0.5 * w, -0.5 * h, 1.0],
        [-0.5 * w, 0.5 * h, 1.0],
        [0.5 * w, 0.5 * h, 1.0],
        [0.5 * w, -0.5 * h, 1.0],
vfdev's avatar
vfdev committed
880
    ])
881
    theta = torch.tensor(matrix, dtype=torch.float).reshape(1, 2, 3)
882
    new_pts = pts.view(1, 4, 3).bmm(theta.transpose(1, 2)).view(4, 2)
vfdev's avatar
vfdev committed
883
884
885
    min_vals, _ = new_pts.min(dim=0)
    max_vals, _ = new_pts.max(dim=0)

886
887
888
889
890
891
    # Truncate precision to 1e-4 to avoid ceil of Xe-15 to 1.0
    tol = 1e-4
    cmax = torch.ceil((max_vals / tol).trunc_() * tol)
    cmin = torch.floor((min_vals / tol).trunc_() * tol)
    size = cmax - cmin
    return int(size[0]), int(size[1])
vfdev's avatar
vfdev committed
892
893
894
895
896


def rotate(
        img: Tensor, matrix: List[float], resample: int = 0, expand: bool = False, fill: Optional[int] = None
) -> Tensor:
897
898
899
900
901
902
    """PRIVATE METHOD. Rotate the Tensor image by angle.

    .. warning::

        Module ``transforms.functional_tensor`` is private and should not be used in user application.
        Please, consider instead using methods from `transforms.functional` module.
vfdev's avatar
vfdev committed
903
904
905
906

    Args:
        img (Tensor): image to be rotated.
        matrix (list of floats): list of 6 float values representing inverse matrix for rotation transformation.
907
            Translation part (``matrix[2]`` and ``matrix[5]``) should be in pixel coordinates.
vfdev's avatar
vfdev committed
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
        resample (int, optional): An optional resampling filter. Default is nearest (=0). Other supported values:
            bilinear(=2).
        expand (bool, optional): Optional expansion flag.
            If true, expands the output image to make it large enough to hold the entire rotated image.
            If false or omitted, make the output image the same size as the input image.
            Note that the expand flag assumes rotation around the center and no translation.
        fill (n-tuple or int or float): this option is not supported for Tensor input.
            Fill value for the area outside the transform in the output image is always 0.

    Returns:
        Tensor: Rotated image.

    .. _filters: https://pillow.readthedocs.io/en/latest/handbook/concepts.html#filters

    """
    _interpolation_modes = {
        0: "nearest",
        2: "bilinear",
    }

    _assert_grid_transform_inputs(img, matrix, resample, fill, _interpolation_modes)
929
    w, h = img.shape[-1], img.shape[-2]
930
    ow, oh = _compute_output_size(matrix, w, h) if expand else (w, h)
931
932
    dtype = img.dtype if torch.is_floating_point(img) else torch.float32
    theta = torch.tensor(matrix, dtype=dtype, device=img.device).reshape(1, 2, 3)
933
    # grid will be generated on the same device as theta and img
934
    grid = _gen_affine_grid(theta, w=w, h=h, ow=ow, oh=oh)
vfdev's avatar
vfdev committed
935
936
937
    mode = _interpolation_modes[resample]

    return _apply_grid_transform(img, grid, mode)
938
939


940
def _perspective_grid(coeffs: List[float], ow: int, oh: int, dtype: torch.dtype, device: torch.device):
941
942
943
944
945
946
947
948
949
950
    # 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]]
951
    ]], dtype=dtype, device=device)
952
953
954
    theta2 = torch.tensor([[
        [coeffs[6], coeffs[7], 1.0],
        [coeffs[6], coeffs[7], 1.0]
955
    ]], dtype=dtype, device=device)
956
957

    d = 0.5
958
    base_grid = torch.empty(1, oh, ow, 3, dtype=dtype, device=device)
959
960
961
962
    base_grid[..., 0].copy_(torch.linspace(d, ow * 1.0 + d - 1.0, steps=ow))
    base_grid[..., 1].copy_(torch.linspace(d, oh * 1.0 + d - 1.0, steps=oh).unsqueeze_(-1))
    base_grid[..., 2].fill_(1)

963
    rescaled_theta1 = theta1.transpose(1, 2) / torch.tensor([0.5 * ow, 0.5 * oh], dtype=dtype, device=device)
964
    output_grid1 = base_grid.view(1, oh * ow, 3).bmm(rescaled_theta1)
965
966
967
968
969
970
971
972
973
    output_grid2 = base_grid.view(1, oh * ow, 3).bmm(theta2.transpose(1, 2))

    output_grid = output_grid1 / output_grid2 - 1.0
    return output_grid.view(1, oh, ow, 2)


def perspective(
        img: Tensor, perspective_coeffs: List[float], interpolation: int = 2, fill: Optional[int] = None
) -> Tensor:
974
975
976
977
978
979
    """PRIVATE METHOD. Perform perspective transform of the given Tensor image.

    .. warning::

        Module ``transforms.functional_tensor`` is private and should not be used in user application.
        Please, consider instead using methods from `transforms.functional` module.
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008

    Args:
        img (Tensor): Image to be transformed.
        perspective_coeffs (list of float): perspective transformation coefficients.
        interpolation (int): Interpolation type. Default, ``PIL.Image.BILINEAR``.
        fill (n-tuple or int or float): this option is not supported for Tensor input. Fill value for the area
            outside the transform in the output image is always 0.

    Returns:
        Tensor: transformed image.
    """
    if not (isinstance(img, torch.Tensor) and _is_tensor_a_torch_image(img)):
        raise TypeError('img should be Tensor Image. Got {}'.format(type(img)))

    _interpolation_modes = {
        0: "nearest",
        2: "bilinear",
    }

    _assert_grid_transform_inputs(
        img,
        matrix=None,
        resample=interpolation,
        fillcolor=fill,
        _interpolation_modes=_interpolation_modes,
        coeffs=perspective_coeffs
    )

    ow, oh = img.shape[-1], img.shape[-2]
1009
1010
    dtype = img.dtype if torch.is_floating_point(img) else torch.float32
    grid = _perspective_grid(perspective_coeffs, ow=ow, oh=oh, dtype=dtype, device=img.device)
1011
1012
1013
    mode = _interpolation_modes[interpolation]

    return _apply_grid_transform(img, grid, mode)