transforms.py 77.7 KB
Newer Older
1
import math
vfdev's avatar
vfdev committed
2
import numbers
3
import random
vfdev's avatar
vfdev committed
4
import warnings
vfdev's avatar
vfdev committed
5
from collections.abc import Sequence
6
from typing import Tuple, List, Optional
vfdev's avatar
vfdev committed
7
8
9
10

import torch
from torch import Tensor

11
12
13
14
15
16
try:
    import accimage
except ImportError:
    accimage = None

from . import functional as F
17
from .functional import InterpolationMode, _interpolation_modes_from_int
18

19

20
21
22
23
__all__ = ["Compose", "ToTensor", "PILToTensor", "ConvertImageDtype", "ToPILImage", "Normalize", "Resize", "Scale",
           "CenterCrop", "Pad", "Lambda", "RandomApply", "RandomChoice", "RandomOrder", "RandomCrop",
           "RandomHorizontalFlip", "RandomVerticalFlip", "RandomResizedCrop", "RandomSizedCrop", "FiveCrop", "TenCrop",
           "LinearTransformation", "ColorJitter", "RandomRotation", "RandomAffine", "Grayscale", "RandomGrayscale",
24
25
           "RandomPerspective", "RandomErasing", "GaussianBlur", "InterpolationMode", "RandomInvert", "RandomPosterize",
           "RandomSolarize", "RandomAdjustSharpness", "RandomAutocontrast", "RandomEqualize"]
26

27

28
class Compose:
29
30
    """Composes several transforms together. This transform does not support torchscript.
    Please, see the note below.
31
32
33
34
35
36
37
38
39

    Args:
        transforms (list of ``Transform`` objects): list of transforms to compose.

    Example:
        >>> transforms.Compose([
        >>>     transforms.CenterCrop(10),
        >>>     transforms.ToTensor(),
        >>> ])
40
41
42
43
44
45
46
47
48
49
50
51
52

    .. note::
        In order to script the transformations, please use ``torch.nn.Sequential`` as below.

        >>> transforms = torch.nn.Sequential(
        >>>     transforms.CenterCrop(10),
        >>>     transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        >>> )
        >>> scripted_transforms = torch.jit.script(transforms)

        Make sure to use only scriptable transformations, i.e. that work with ``torch.Tensor``, does not require
        `lambda` functions or ``PIL.Image``.

53
54
55
56
57
58
59
60
61
62
    """

    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, img):
        for t in self.transforms:
            img = t(img)
        return img

63
64
65
66
67
68
69
70
    def __repr__(self):
        format_string = self.__class__.__name__ + '('
        for t in self.transforms:
            format_string += '\n'
            format_string += '    {0}'.format(t)
        format_string += '\n)'
        return format_string

71

72
class ToTensor:
73
    """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor. This transform does not support torchscript.
74
75

    Converts a PIL Image or numpy.ndarray (H x W x C) in the range
surgan12's avatar
surgan12 committed
76
77
78
79
80
    [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0]
    if the PIL Image belongs to one of the modes (L, LA, P, I, F, RGB, YCbCr, RGBA, CMYK, 1)
    or if the numpy.ndarray has dtype = np.uint8

    In the other cases, tensors are returned without scaling.
81
82
83
84
85
86

    .. note::
        Because the input image is scaled to [0.0, 1.0], this transformation should not be used when
        transforming target image masks. See the `references`_ for implementing the transforms for image masks.

    .. _references: https://github.com/pytorch/vision/tree/master/references/segmentation
87
88
89
90
91
92
93
94
95
96
97
98
    """

    def __call__(self, pic):
        """
        Args:
            pic (PIL Image or numpy.ndarray): Image to be converted to tensor.

        Returns:
            Tensor: Converted image.
        """
        return F.to_tensor(pic)

99
100
101
    def __repr__(self):
        return self.__class__.__name__ + '()'

102

103
class PILToTensor:
104
    """Convert a ``PIL Image`` to a tensor of the same type. This transform does not support torchscript.
105

vfdev's avatar
vfdev committed
106
    Converts a PIL Image (H x W x C) to a Tensor of shape (C x H x W).
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
    """

    def __call__(self, pic):
        """
        Args:
            pic (PIL Image): Image to be converted to tensor.

        Returns:
            Tensor: Converted image.
        """
        return F.pil_to_tensor(pic)

    def __repr__(self):
        return self.__class__.__name__ + '()'


123
class ConvertImageDtype(torch.nn.Module):
124
    """Convert a tensor image to the given ``dtype`` and scale the values accordingly
125
    This function does not support PIL Image.
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142

    Args:
        dtype (torch.dtype): Desired data type of the output

    .. note::

        When converting from a smaller to a larger integer ``dtype`` the maximum values are **not** mapped exactly.
        If converted back and forth, this mismatch has no effect.

    Raises:
        RuntimeError: When trying to cast :class:`torch.float32` to :class:`torch.int32` or :class:`torch.int64` as
            well as for trying to cast :class:`torch.float64` to :class:`torch.int64`. These conversions might lead to
            overflow errors since the floating point ``dtype`` cannot store consecutive integers over the whole range
            of the integer ``dtype``.
    """

    def __init__(self, dtype: torch.dtype) -> None:
143
        super().__init__()
144
145
        self.dtype = dtype

vfdev's avatar
vfdev committed
146
    def forward(self, image):
147
148
149
        return F.convert_image_dtype(image, self.dtype)


150
class ToPILImage:
151
    """Convert a tensor or an ndarray to PIL Image. This transform does not support torchscript.
152
153
154
155
156
157
158

    Converts a torch.*Tensor of shape C x H x W or a numpy ndarray of shape
    H x W x C to a PIL Image while preserving the value range.

    Args:
        mode (`PIL.Image mode`_): color space and pixel depth of input data (optional).
            If ``mode`` is ``None`` (default) there are some assumptions made about the input data:
vfdev's avatar
vfdev committed
159
160
161
162
163
            - If the input has 4 channels, the ``mode`` is assumed to be ``RGBA``.
            - If the input has 3 channels, the ``mode`` is assumed to be ``RGB``.
            - If the input has 2 channels, the ``mode`` is assumed to be ``LA``.
            - If the input has 1 channel, the ``mode`` is determined by the data type (i.e ``int``, ``float``,
            ``short``).
164

csukuangfj's avatar
csukuangfj committed
165
    .. _PIL.Image mode: https://pillow.readthedocs.io/en/latest/handbook/concepts.html#concept-modes
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
    """
    def __init__(self, mode=None):
        self.mode = mode

    def __call__(self, pic):
        """
        Args:
            pic (Tensor or numpy.ndarray): Image to be converted to PIL Image.

        Returns:
            PIL Image: Image converted to PIL Image.

        """
        return F.to_pil_image(pic, self.mode)

181
    def __repr__(self):
182
183
184
185
186
        format_string = self.__class__.__name__ + '('
        if self.mode is not None:
            format_string += 'mode={0}'.format(self.mode)
        format_string += ')'
        return format_string
187

188

189
class Normalize(torch.nn.Module):
Fang Gao's avatar
Fang Gao committed
190
    """Normalize a tensor image with mean and standard deviation.
191
    This transform does not support PIL Image.
192
193
194
    Given mean: ``(mean[1],...,mean[n])`` and std: ``(std[1],..,std[n])`` for ``n``
    channels, this transform will normalize each channel of the input
    ``torch.*Tensor`` i.e.,
abdjava's avatar
abdjava committed
195
    ``output[channel] = (input[channel] - mean[channel]) / std[channel]``
196

197
    .. note::
198
        This transform acts out of place, i.e., it does not mutate the input tensor.
199

200
201
202
    Args:
        mean (sequence): Sequence of means for each channel.
        std (sequence): Sequence of standard deviations for each channel.
203
204
        inplace(bool,optional): Bool to make this operation in-place.

205
206
    """

surgan12's avatar
surgan12 committed
207
    def __init__(self, mean, std, inplace=False):
208
        super().__init__()
209
210
        self.mean = mean
        self.std = std
surgan12's avatar
surgan12 committed
211
        self.inplace = inplace
212

213
    def forward(self, tensor: Tensor) -> Tensor:
214
215
        """
        Args:
vfdev's avatar
vfdev committed
216
            tensor (Tensor): Tensor image to be normalized.
217
218
219
220

        Returns:
            Tensor: Normalized Tensor image.
        """
surgan12's avatar
surgan12 committed
221
        return F.normalize(tensor, self.mean, self.std, self.inplace)
222

223
224
225
    def __repr__(self):
        return self.__class__.__name__ + '(mean={0}, std={1})'.format(self.mean, self.std)

226

vfdev's avatar
vfdev committed
227
228
class Resize(torch.nn.Module):
    """Resize the input image to the given size.
229
    If the image is torch Tensor, it is expected
vfdev's avatar
vfdev committed
230
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions
231

232
233
234
235
236
237
    .. warning::
        The output image might be different depending on its type: when downsampling, the interpolation of PIL images
        and tensors is slightly different, because PIL applies antialiasing. This may lead to significant differences
        in the performance of a network. Therefore, it is preferable to train and serve a model with the same input
        types.

238
239
240
241
242
    Args:
        size (sequence or int): Desired output size. If size is a sequence like
            (h, w), output size will be matched to this. If size is an int,
            smaller edge of the image will be matched to this number.
            i.e, if height > width, then image will be rescaled to
vfdev's avatar
vfdev committed
243
            (size * height / width, size).
244
245
246

            .. note::
                In torchscript mode size as single int is not supported, use a sequence of length 1: ``[size, ]``.
247
248
249
250
        interpolation (InterpolationMode): Desired interpolation enum defined by
            :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``.
            If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` and
            ``InterpolationMode.BICUBIC`` are supported.
251
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
252
253
254
255
        max_size (int, optional): The maximum allowed for the longer edge of
            the resized image: if the longer edge of the image is greater
            than ``max_size`` after being resized according to ``size``, then
            the image is resized again so that the longer edge is equal to
256
            ``max_size``. As a result, ``size`` might be overruled, i.e the
257
258
259
            smaller edge may be shorter than ``size``. This is only supported
            if ``size`` is an int (or a sequence of length 1 in torchscript
            mode).
260
261
262
263
264
265
        antialias (bool, optional): antialias flag. If ``img`` is PIL Image, the flag is ignored and anti-alias
            is always used. If ``img`` is Tensor, the flag is False by default and can be set True for
            ``InterpolationMode.BILINEAR`` only mode.

            .. warning::
                There is no autodiff support for ``antialias=True`` option with input ``img`` as Tensor.
266

267
268
    """

269
    def __init__(self, size, interpolation=InterpolationMode.BILINEAR, max_size=None, antialias=None):
vfdev's avatar
vfdev committed
270
        super().__init__()
271
272
273
274
275
        if not isinstance(size, (int, Sequence)):
            raise TypeError("Size should be int or sequence. Got {}".format(type(size)))
        if isinstance(size, Sequence) and len(size) not in (1, 2):
            raise ValueError("If size is a sequence, it should have 1 or 2 values")
        self.size = size
276
        self.max_size = max_size
277
278
279
280

        # Backward compatibility with integer value
        if isinstance(interpolation, int):
            warnings.warn(
281
282
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
283
284
285
            )
            interpolation = _interpolation_modes_from_int(interpolation)

286
        self.interpolation = interpolation
287
        self.antialias = antialias
288

vfdev's avatar
vfdev committed
289
    def forward(self, img):
290
291
        """
        Args:
vfdev's avatar
vfdev committed
292
            img (PIL Image or Tensor): Image to be scaled.
293
294

        Returns:
vfdev's avatar
vfdev committed
295
            PIL Image or Tensor: Rescaled image.
296
        """
297
        return F.resize(img, self.size, self.interpolation, self.max_size, self.antialias)
298

299
    def __repr__(self):
300
        interpolate_str = self.interpolation.value
301
302
        return self.__class__.__name__ + '(size={0}, interpolation={1}, max_size={2}, antialias={3})'.format(
            self.size, interpolate_str, self.max_size, self.antialias)
303

304
305
306
307
308
309
310
311
312
313
314

class Scale(Resize):
    """
    Note: This transform is deprecated in favor of Resize.
    """
    def __init__(self, *args, **kwargs):
        warnings.warn("The use of the transforms.Scale transform is deprecated, " +
                      "please use transforms.Resize instead.")
        super(Scale, self).__init__(*args, **kwargs)


vfdev's avatar
vfdev committed
315
316
class CenterCrop(torch.nn.Module):
    """Crops the given image at the center.
317
    If the image is torch Tensor, it is expected
318
319
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
    If image size is smaller than output size along any edge, image is padded with 0 and then center cropped.
320
321
322
323

    Args:
        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
324
            made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
325
326
327
    """

    def __init__(self, size):
vfdev's avatar
vfdev committed
328
        super().__init__()
329
        self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
330

vfdev's avatar
vfdev committed
331
    def forward(self, img):
332
333
        """
        Args:
vfdev's avatar
vfdev committed
334
            img (PIL Image or Tensor): Image to be cropped.
335
336

        Returns:
vfdev's avatar
vfdev committed
337
            PIL Image or Tensor: Cropped image.
338
339
340
        """
        return F.center_crop(img, self.size)

341
342
343
    def __repr__(self):
        return self.__class__.__name__ + '(size={0})'.format(self.size)

344

345
346
class Pad(torch.nn.Module):
    """Pad the given image on all sides with the given "pad" value.
347
    If the image is torch Tensor, it is expected
348
349
350
    to have [..., H, W] shape, where ... means at most 2 leading dimensions for mode reflect and symmetric,
    at most 3 leading dimensions for mode edge,
    and an arbitrary number of leading dimensions for mode constant
351
352

    Args:
353
354
355
        padding (int or sequence): Padding on each border. If a single int is provided this
            is used to pad all borders. If sequence of length 2 is provided this is the padding
            on left/right and top/bottom respectively. If a sequence of length 4 is provided
356
            this is the padding for the left, top, right and bottom borders respectively.
357
358
359
360

            .. note::
                In torchscript mode padding as single int is not supported, use a sequence of
                length 1: ``[padding, ]``.
361
        fill (number or str or tuple): Pixel fill value for constant fill. Default is 0. If a tuple of
362
            length 3, it is used to fill R, G, B channels respectively.
363
364
365
            This value is only used when the padding_mode is constant.
            Only number is supported for torch Tensor.
            Only int or str or tuple value is supported for PIL Image.
366
        padding_mode (str): Type of padding. Should be: constant, edge, reflect or symmetric.
367
            Default is constant.
368
369
370

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

371
372
            - edge: pads with the last value at the edge of the image.
              If input a 5D torch Tensor, the last 3 dimensions will be padded instead of the last 2
373

374
375
376
            - reflect: pads with reflection of image without repeating the last value on the edge.
              For example, 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]
377

378
379
380
            - symmetric: pads with reflection of image repeating the last value on the edge.
              For example, 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]
381
382
    """

383
384
385
386
387
388
389
390
391
392
393
394
395
    def __init__(self, padding, fill=0, padding_mode="constant"):
        super().__init__()
        if not isinstance(padding, (numbers.Number, tuple, list)):
            raise TypeError("Got inappropriate padding arg")

        if not isinstance(fill, (numbers.Number, str, tuple)):
            raise TypeError("Got inappropriate fill arg")

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

        if isinstance(padding, Sequence) 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 " +
396
397
398
399
                             "{} element tuple".format(len(padding)))

        self.padding = padding
        self.fill = fill
400
        self.padding_mode = padding_mode
401

402
    def forward(self, img):
403
404
        """
        Args:
405
            img (PIL Image or Tensor): Image to be padded.
406
407

        Returns:
408
            PIL Image or Tensor: Padded image.
409
        """
410
        return F.pad(img, self.padding, self.fill, self.padding_mode)
411

412
    def __repr__(self):
413
414
        return self.__class__.__name__ + '(padding={0}, fill={1}, padding_mode={2})'.\
            format(self.padding, self.fill, self.padding_mode)
415

416

417
class Lambda:
418
    """Apply a user-defined lambda as a transform. This transform does not support torchscript.
419
420
421
422
423
424

    Args:
        lambd (function): Lambda/function to be used for transform.
    """

    def __init__(self, lambd):
425
426
        if not callable(lambd):
            raise TypeError("Argument lambd should be callable, got {}".format(repr(type(lambd).__name__)))
427
428
429
430
431
        self.lambd = lambd

    def __call__(self, img):
        return self.lambd(img)

432
433
434
    def __repr__(self):
        return self.__class__.__name__ + '()'

435

436
class RandomTransforms:
437
438
439
    """Base class for a list of transformations with randomness

    Args:
440
        transforms (sequence): list of transformations
441
442
443
    """

    def __init__(self, transforms):
444
445
        if not isinstance(transforms, Sequence):
            raise TypeError("Argument transforms should be a sequence")
446
447
448
449
450
451
452
453
454
455
456
457
458
459
        self.transforms = transforms

    def __call__(self, *args, **kwargs):
        raise NotImplementedError()

    def __repr__(self):
        format_string = self.__class__.__name__ + '('
        for t in self.transforms:
            format_string += '\n'
            format_string += '    {0}'.format(t)
        format_string += '\n)'
        return format_string


460
class RandomApply(torch.nn.Module):
461
    """Apply randomly a list of transformations with a given probability.
462
463
464
465
466
467
468
469
470
471
472
473

    .. note::
        In order to script the transformation, please use ``torch.nn.ModuleList`` as input instead of list/tuple of
        transforms as shown below:

        >>> transforms = transforms.RandomApply(torch.nn.ModuleList([
        >>>     transforms.ColorJitter(),
        >>> ]), p=0.3)
        >>> scripted_transforms = torch.jit.script(transforms)

        Make sure to use only scriptable transformations, i.e. that work with ``torch.Tensor``, does not require
        `lambda` functions or ``PIL.Image``.
474
475

    Args:
476
        transforms (sequence or torch.nn.Module): list of transformations
477
478
479
480
        p (float): probability
    """

    def __init__(self, transforms, p=0.5):
481
482
        super().__init__()
        self.transforms = transforms
483
484
        self.p = p

485
486
    def forward(self, img):
        if self.p < torch.rand(1):
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
            return img
        for t in self.transforms:
            img = t(img)
        return img

    def __repr__(self):
        format_string = self.__class__.__name__ + '('
        format_string += '\n    p={}'.format(self.p)
        for t in self.transforms:
            format_string += '\n'
            format_string += '    {0}'.format(t)
        format_string += '\n)'
        return format_string


class RandomOrder(RandomTransforms):
503
    """Apply a list of transformations in a random order. This transform does not support torchscript.
504
505
506
507
508
509
510
511
512
513
    """
    def __call__(self, img):
        order = list(range(len(self.transforms)))
        random.shuffle(order)
        for i in order:
            img = self.transforms[i](img)
        return img


class RandomChoice(RandomTransforms):
514
    """Apply single transformation randomly picked from a list. This transform does not support torchscript.
515
516
517
518
519
520
    """
    def __call__(self, img):
        t = random.choice(self.transforms)
        return t(img)


vfdev's avatar
vfdev committed
521
522
class RandomCrop(torch.nn.Module):
    """Crop the given image at a random location.
523
    If the image is torch Tensor, it is expected
524
525
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions,
    but if non-constant padding is used, the input is expected to have at most 2 leading dimensions
526
527
528
529

    Args:
        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
530
            made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
531
        padding (int or sequence, optional): Optional padding on each border
vfdev's avatar
vfdev committed
532
            of the image. Default is None. If a single int is provided this
533
534
            is used to pad all borders. If sequence of length 2 is provided this is the padding
            on left/right and top/bottom respectively. If a sequence of length 4 is provided
vfdev's avatar
vfdev committed
535
            this is the padding for the left, top, right and bottom borders respectively.
536
537
538
539

            .. note::
                In torchscript mode padding as single int is not supported, use a sequence of
                length 1: ``[padding, ]``.
540
        pad_if_needed (boolean): It will pad the image if smaller than the
ekka's avatar
ekka committed
541
            desired size to avoid raising an exception. Since cropping is done
542
            after padding, the padding seems to be done at a random offset.
543
        fill (number or str or tuple): Pixel fill value for constant fill. Default is 0. If a tuple of
544
            length 3, it is used to fill R, G, B channels respectively.
545
546
547
            This value is only used when the padding_mode is constant.
            Only number is supported for torch Tensor.
            Only int or str or tuple value is supported for PIL Image.
548
549
        padding_mode (str): Type of padding. Should be: constant, edge, reflect or symmetric.
            Default is constant.
550

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

553
554
            - edge: pads with the last value at the edge of the image.
              If input a 5D torch Tensor, the last 3 dimensions will be padded instead of the last 2
555

556
557
558
            - reflect: pads with reflection of image without repeating the last value on the edge.
              For example, 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]
559

560
561
562
            - symmetric: pads with reflection of image repeating the last value on the edge.
              For example, 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]
563
564
565
    """

    @staticmethod
vfdev's avatar
vfdev committed
566
    def get_params(img: Tensor, output_size: Tuple[int, int]) -> Tuple[int, int, int, int]:
567
568
569
        """Get parameters for ``crop`` for a random crop.

        Args:
vfdev's avatar
vfdev committed
570
            img (PIL Image or Tensor): Image to be cropped.
571
572
573
574
575
            output_size (tuple): Expected output size of the crop.

        Returns:
            tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
        """
vfdev's avatar
vfdev committed
576
        w, h = F._get_image_size(img)
577
        th, tw = output_size
vfdev's avatar
vfdev committed
578
579
580
581
582
583

        if h + 1 < th or w + 1 < tw:
            raise ValueError(
                "Required crop size {} is larger then input image size {}".format((th, tw), (h, w))
            )

584
585
586
        if w == tw and h == th:
            return 0, 0, h, w

587
588
        i = torch.randint(0, h - th + 1, size=(1, )).item()
        j = torch.randint(0, w - tw + 1, size=(1, )).item()
589
590
        return i, j, th, tw

vfdev's avatar
vfdev committed
591
592
593
    def __init__(self, size, padding=None, pad_if_needed=False, fill=0, padding_mode="constant"):
        super().__init__()

594
595
596
597
        self.size = tuple(_setup_size(
            size, error_msg="Please provide only two dimensions (h, w) for size."
        ))

vfdev's avatar
vfdev committed
598
599
600
601
602
603
        self.padding = padding
        self.pad_if_needed = pad_if_needed
        self.fill = fill
        self.padding_mode = padding_mode

    def forward(self, img):
604
605
        """
        Args:
vfdev's avatar
vfdev committed
606
            img (PIL Image or Tensor): Image to be cropped.
607
608

        Returns:
vfdev's avatar
vfdev committed
609
            PIL Image or Tensor: Cropped image.
610
        """
611
612
        if self.padding is not None:
            img = F.pad(img, self.padding, self.fill, self.padding_mode)
613

vfdev's avatar
vfdev committed
614
        width, height = F._get_image_size(img)
615
        # pad the width if needed
vfdev's avatar
vfdev committed
616
617
618
        if self.pad_if_needed and width < self.size[1]:
            padding = [self.size[1] - width, 0]
            img = F.pad(img, padding, self.fill, self.padding_mode)
619
        # pad the height if needed
vfdev's avatar
vfdev committed
620
621
622
        if self.pad_if_needed and height < self.size[0]:
            padding = [0, self.size[0] - height]
            img = F.pad(img, padding, self.fill, self.padding_mode)
623

624
625
626
627
        i, j, h, w = self.get_params(img, self.size)

        return F.crop(img, i, j, h, w)

628
    def __repr__(self):
vfdev's avatar
vfdev committed
629
        return self.__class__.__name__ + "(size={0}, padding={1})".format(self.size, self.padding)
630

631

632
633
class RandomHorizontalFlip(torch.nn.Module):
    """Horizontally flip the given image randomly with a given probability.
634
    If the image is torch Tensor, it is expected
635
636
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
637
638
639
640
641
642

    Args:
        p (float): probability of the image being flipped. Default value is 0.5
    """

    def __init__(self, p=0.5):
643
        super().__init__()
644
        self.p = p
645

646
    def forward(self, img):
647
648
        """
        Args:
649
            img (PIL Image or Tensor): Image to be flipped.
650
651

        Returns:
652
            PIL Image or Tensor: Randomly flipped image.
653
        """
654
        if torch.rand(1) < self.p:
655
656
657
            return F.hflip(img)
        return img

658
    def __repr__(self):
659
        return self.__class__.__name__ + '(p={})'.format(self.p)
660

661

662
class RandomVerticalFlip(torch.nn.Module):
vfdev's avatar
vfdev committed
663
    """Vertically flip the given image randomly with a given probability.
664
    If the image is torch Tensor, it is expected
665
666
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
667
668
669
670
671
672

    Args:
        p (float): probability of the image being flipped. Default value is 0.5
    """

    def __init__(self, p=0.5):
673
        super().__init__()
674
        self.p = p
675

676
    def forward(self, img):
677
678
        """
        Args:
679
            img (PIL Image or Tensor): Image to be flipped.
680
681

        Returns:
682
            PIL Image or Tensor: Randomly flipped image.
683
        """
684
        if torch.rand(1) < self.p:
685
686
687
            return F.vflip(img)
        return img

688
    def __repr__(self):
689
        return self.__class__.__name__ + '(p={})'.format(self.p)
690

691

692
693
class RandomPerspective(torch.nn.Module):
    """Performs a random perspective transformation of the given image with a given probability.
694
    If the image is torch Tensor, it is expected
695
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
696
697

    Args:
698
699
700
        distortion_scale (float): argument to control the degree of distortion and ranges from 0 to 1.
            Default is 0.5.
        p (float): probability of the image being transformed. Default is 0.5.
701
702
703
        interpolation (InterpolationMode): Desired interpolation enum defined by
            :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``.
            If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported.
704
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
705
706
        fill (sequence or number): Pixel fill value for the area outside the transformed
            image. Default is ``0``. If given a number, the value is used for all bands respectively.
707
            If input is PIL Image, the options is only available for ``Pillow>=5.0.0``.
708
709
    """

710
    def __init__(self, distortion_scale=0.5, p=0.5, interpolation=InterpolationMode.BILINEAR, fill=0):
711
        super().__init__()
712
        self.p = p
713
714
715
716

        # Backward compatibility with integer value
        if isinstance(interpolation, int):
            warnings.warn(
717
718
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
719
720
721
            )
            interpolation = _interpolation_modes_from_int(interpolation)

722
723
        self.interpolation = interpolation
        self.distortion_scale = distortion_scale
724
725
726
727
728
729

        if fill is None:
            fill = 0
        elif not isinstance(fill, (Sequence, numbers.Number)):
            raise TypeError("Fill should be either a sequence or a number.")

730
        self.fill = fill
731

732
    def forward(self, img):
733
734
        """
        Args:
735
            img (PIL Image or Tensor): Image to be Perspectively transformed.
736
737

        Returns:
738
            PIL Image or Tensor: Randomly transformed image.
739
        """
740
741
742
743
744
745
746
747

        fill = self.fill
        if isinstance(img, Tensor):
            if isinstance(fill, (int, float)):
                fill = [float(fill)] * F._get_image_num_channels(img)
            else:
                fill = [float(f) for f in fill]

748
749
        if torch.rand(1) < self.p:
            width, height = F._get_image_size(img)
750
            startpoints, endpoints = self.get_params(width, height, self.distortion_scale)
751
            return F.perspective(img, startpoints, endpoints, self.interpolation, fill)
752
753
754
        return img

    @staticmethod
755
    def get_params(width: int, height: int, distortion_scale: float) -> Tuple[List[List[int]], List[List[int]]]:
756
757
758
        """Get parameters for ``perspective`` for a random perspective transform.

        Args:
759
760
761
            width (int): width of the image.
            height (int): height of the image.
            distortion_scale (float): argument to control the degree of distortion and ranges from 0 to 1.
762
763

        Returns:
764
            List containing [top-left, top-right, bottom-right, bottom-left] of the original image,
765
766
            List containing [top-left, top-right, bottom-right, bottom-left] of the transformed image.
        """
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
        half_height = height // 2
        half_width = width // 2
        topleft = [
            int(torch.randint(0, int(distortion_scale * half_width) + 1, size=(1, )).item()),
            int(torch.randint(0, int(distortion_scale * half_height) + 1, size=(1, )).item())
        ]
        topright = [
            int(torch.randint(width - int(distortion_scale * half_width) - 1, width, size=(1, )).item()),
            int(torch.randint(0, int(distortion_scale * half_height) + 1, size=(1, )).item())
        ]
        botright = [
            int(torch.randint(width - int(distortion_scale * half_width) - 1, width, size=(1, )).item()),
            int(torch.randint(height - int(distortion_scale * half_height) - 1, height, size=(1, )).item())
        ]
        botleft = [
            int(torch.randint(0, int(distortion_scale * half_width) + 1, size=(1, )).item()),
            int(torch.randint(height - int(distortion_scale * half_height) - 1, height, size=(1, )).item())
        ]
        startpoints = [[0, 0], [width - 1, 0], [width - 1, height - 1], [0, height - 1]]
786
787
788
789
790
791
792
        endpoints = [topleft, topright, botright, botleft]
        return startpoints, endpoints

    def __repr__(self):
        return self.__class__.__name__ + '(p={})'.format(self.p)


793
class RandomResizedCrop(torch.nn.Module):
794
795
    """Crop a random portion of image and resize it to a given size.

796
    If the image is torch Tensor, it is expected
797
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions
798

799
800
801
    A crop of the original image is made: the crop has a random area (H * W)
    and a random aspect ratio. This crop is finally resized to the given
    size. This is popularly used to train the Inception networks.
802
803

    Args:
804
        size (int or sequence): expected output size of the crop, for each edge. If size is an
805
            int instead of sequence like (h, w), a square output size ``(size, size)`` is
806
            made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
807
808
809

            .. note::
                In torchscript mode size as single int is not supported, use a sequence of length 1: ``[size, ]``.
Nicolas Hug's avatar
Nicolas Hug committed
810
811
        scale (tuple of float): Specifies the lower and upper bounds for the random area of the crop,
            before resizing. The scale is defined with respect to the area of the original image.
812
813
        ratio (tuple of float): lower and upper bounds for the random aspect ratio of the crop, before
            resizing.
814
815
816
817
        interpolation (InterpolationMode): Desired interpolation enum defined by
            :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``.
            If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` and
            ``InterpolationMode.BICUBIC`` are supported.
818
819
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.

820
821
    """

822
    def __init__(self, size, scale=(0.08, 1.0), ratio=(3. / 4., 4. / 3.), interpolation=InterpolationMode.BILINEAR):
823
        super().__init__()
824
        self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
825

826
        if not isinstance(scale, Sequence):
827
            raise TypeError("Scale should be a sequence")
828
        if not isinstance(ratio, Sequence):
829
            raise TypeError("Ratio should be a sequence")
830
        if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
831
            warnings.warn("Scale and ratio should be of kind (min, max)")
832

833
834
835
        # Backward compatibility with integer value
        if isinstance(interpolation, int):
            warnings.warn(
836
837
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
838
839
840
            )
            interpolation = _interpolation_modes_from_int(interpolation)

841
        self.interpolation = interpolation
842
843
        self.scale = scale
        self.ratio = ratio
844
845

    @staticmethod
846
    def get_params(
847
            img: Tensor, scale: List[float], ratio: List[float]
848
    ) -> Tuple[int, int, int, int]:
849
850
851
        """Get parameters for ``crop`` for a random sized crop.

        Args:
852
            img (PIL Image or Tensor): Input image.
853
854
            scale (list): range of scale of the origin size cropped
            ratio (list): range of aspect ratio of the origin aspect ratio cropped
855
856
857

        Returns:
            tuple: params (i, j, h, w) to be passed to ``crop`` for a random
858
            sized crop.
859
        """
vfdev's avatar
vfdev committed
860
        width, height = F._get_image_size(img)
Zhicheng Yan's avatar
Zhicheng Yan committed
861
        area = height * width
862

863
        log_ratio = torch.log(torch.tensor(ratio))
864
        for _ in range(10):
865
            target_area = area * torch.empty(1).uniform_(scale[0], scale[1]).item()
866
867
868
            aspect_ratio = torch.exp(
                torch.empty(1).uniform_(log_ratio[0], log_ratio[1])
            ).item()
869
870
871
872

            w = int(round(math.sqrt(target_area * aspect_ratio)))
            h = int(round(math.sqrt(target_area / aspect_ratio)))

Zhicheng Yan's avatar
Zhicheng Yan committed
873
            if 0 < w <= width and 0 < h <= height:
874
875
                i = torch.randint(0, height - h + 1, size=(1,)).item()
                j = torch.randint(0, width - w + 1, size=(1,)).item()
876
877
                return i, j, h, w

878
        # Fallback to central crop
Zhicheng Yan's avatar
Zhicheng Yan committed
879
        in_ratio = float(width) / float(height)
880
        if in_ratio < min(ratio):
Zhicheng Yan's avatar
Zhicheng Yan committed
881
            w = width
882
            h = int(round(w / min(ratio)))
883
        elif in_ratio > max(ratio):
Zhicheng Yan's avatar
Zhicheng Yan committed
884
            h = height
885
            w = int(round(h * max(ratio)))
886
        else:  # whole image
Zhicheng Yan's avatar
Zhicheng Yan committed
887
888
889
890
            w = width
            h = height
        i = (height - h) // 2
        j = (width - w) // 2
891
        return i, j, h, w
892

893
    def forward(self, img):
894
895
        """
        Args:
896
            img (PIL Image or Tensor): Image to be cropped and resized.
897
898

        Returns:
899
            PIL Image or Tensor: Randomly cropped and resized image.
900
        """
901
        i, j, h, w = self.get_params(img, self.scale, self.ratio)
902
903
        return F.resized_crop(img, i, j, h, w, self.size, self.interpolation)

904
    def __repr__(self):
905
        interpolate_str = self.interpolation.value
906
        format_string = self.__class__.__name__ + '(size={0}'.format(self.size)
907
908
        format_string += ', scale={0}'.format(tuple(round(s, 4) for s in self.scale))
        format_string += ', ratio={0}'.format(tuple(round(r, 4) for r in self.ratio))
909
910
        format_string += ', interpolation={0})'.format(interpolate_str)
        return format_string
911

912
913
914
915
916
917
918
919
920
921
922

class RandomSizedCrop(RandomResizedCrop):
    """
    Note: This transform is deprecated in favor of RandomResizedCrop.
    """
    def __init__(self, *args, **kwargs):
        warnings.warn("The use of the transforms.RandomSizedCrop transform is deprecated, " +
                      "please use transforms.RandomResizedCrop instead.")
        super(RandomSizedCrop, self).__init__(*args, **kwargs)


vfdev's avatar
vfdev committed
923
924
class FiveCrop(torch.nn.Module):
    """Crop the given image into four corners and the central crop.
925
    If the image is torch Tensor, it is expected
vfdev's avatar
vfdev committed
926
927
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
928
929
930
931
932
933
934
935
936

    .. Note::
         This transform returns a tuple of images and there may be a mismatch in the number of
         inputs and targets your Dataset returns. See below for an example of how to deal with
         this.

    Args:
         size (sequence or int): Desired output size of the crop. If size is an ``int``
            instead of sequence like (h, w), a square crop of size (size, size) is made.
937
            If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
938
939
940
941
942
943
944
945
946
947
948
949
950
951

    Example:
         >>> transform = Compose([
         >>>    FiveCrop(size), # this is a list of PIL Images
         >>>    Lambda(lambda crops: torch.stack([ToTensor()(crop) for crop in crops])) # returns a 4D tensor
         >>> ])
         >>> #In your test loop you can do the following:
         >>> input, target = batch # input is a 5d tensor, target is 2d
         >>> bs, ncrops, c, h, w = input.size()
         >>> result = model(input.view(-1, c, h, w)) # fuse batch size and ncrops
         >>> result_avg = result.view(bs, ncrops, -1).mean(1) # avg over crops
    """

    def __init__(self, size):
vfdev's avatar
vfdev committed
952
        super().__init__()
953
        self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
954

vfdev's avatar
vfdev committed
955
956
957
958
959
960
961
962
    def forward(self, img):
        """
        Args:
            img (PIL Image or Tensor): Image to be cropped.

        Returns:
            tuple of 5 images. Image can be PIL Image or Tensor
        """
963
964
        return F.five_crop(img, self.size)

965
966
967
    def __repr__(self):
        return self.__class__.__name__ + '(size={0})'.format(self.size)

968

vfdev's avatar
vfdev committed
969
970
971
class TenCrop(torch.nn.Module):
    """Crop the given image into four corners and the central crop plus the flipped version of
    these (horizontal flipping is used by default).
972
    If the image is torch Tensor, it is expected
vfdev's avatar
vfdev committed
973
974
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
975
976
977
978
979
980
981
982
983

    .. Note::
         This transform returns a tuple of images and there may be a mismatch in the number of
         inputs and targets your Dataset returns. See below for an example of how to deal with
         this.

    Args:
        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
984
            made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
985
        vertical_flip (bool): Use vertical flipping instead of horizontal
986
987
988
989
990
991
992
993
994
995
996
997
998
999

    Example:
         >>> transform = Compose([
         >>>    TenCrop(size), # this is a list of PIL Images
         >>>    Lambda(lambda crops: torch.stack([ToTensor()(crop) for crop in crops])) # returns a 4D tensor
         >>> ])
         >>> #In your test loop you can do the following:
         >>> input, target = batch # input is a 5d tensor, target is 2d
         >>> bs, ncrops, c, h, w = input.size()
         >>> result = model(input.view(-1, c, h, w)) # fuse batch size and ncrops
         >>> result_avg = result.view(bs, ncrops, -1).mean(1) # avg over crops
    """

    def __init__(self, size, vertical_flip=False):
vfdev's avatar
vfdev committed
1000
        super().__init__()
1001
        self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
1002
1003
        self.vertical_flip = vertical_flip

vfdev's avatar
vfdev committed
1004
1005
1006
1007
1008
1009
1010
1011
    def forward(self, img):
        """
        Args:
            img (PIL Image or Tensor): Image to be cropped.

        Returns:
            tuple of 10 images. Image can be PIL Image or Tensor
        """
1012
1013
        return F.ten_crop(img, self.size, self.vertical_flip)

1014
    def __repr__(self):
1015
        return self.__class__.__name__ + '(size={0}, vertical_flip={1})'.format(self.size, self.vertical_flip)
1016

1017

1018
class LinearTransformation(torch.nn.Module):
ekka's avatar
ekka committed
1019
    """Transform a tensor image with a square transformation matrix and a mean_vector computed
1020
    offline.
1021
    This transform does not support PIL Image.
ekka's avatar
ekka committed
1022
1023
1024
    Given transformation_matrix and mean_vector, will flatten the torch.*Tensor and
    subtract mean_vector from it which is then followed by computing the dot
    product with the transformation matrix and then reshaping the tensor to its
1025
    original shape.
1026

1027
    Applications:
1028
        whitening transformation: Suppose X is a column vector zero-centered data.
1029
1030
1031
        Then compute the data covariance matrix [D x D] with torch.mm(X.t(), X),
        perform SVD on this matrix and pass it as transformation_matrix.

1032
1033
    Args:
        transformation_matrix (Tensor): tensor [D x D], D = C x H x W
ekka's avatar
ekka committed
1034
        mean_vector (Tensor): tensor [D], D = C x H x W
1035
1036
    """

ekka's avatar
ekka committed
1037
    def __init__(self, transformation_matrix, mean_vector):
1038
        super().__init__()
1039
1040
1041
        if transformation_matrix.size(0) != transformation_matrix.size(1):
            raise ValueError("transformation_matrix should be square. Got " +
                             "[{} x {}] rectangular matrix.".format(*transformation_matrix.size()))
ekka's avatar
ekka committed
1042
1043
1044

        if mean_vector.size(0) != transformation_matrix.size(0):
            raise ValueError("mean_vector should have the same length {}".format(mean_vector.size(0)) +
Francisco Massa's avatar
Francisco Massa committed
1045
1046
                             " as any one of the dimensions of the transformation_matrix [{}]"
                             .format(tuple(transformation_matrix.size())))
ekka's avatar
ekka committed
1047

1048
1049
1050
1051
        if transformation_matrix.device != mean_vector.device:
            raise ValueError("Input tensors should be on the same device. Got {} and {}"
                             .format(transformation_matrix.device, mean_vector.device))

1052
        self.transformation_matrix = transformation_matrix
ekka's avatar
ekka committed
1053
        self.mean_vector = mean_vector
1054

1055
    def forward(self, tensor: Tensor) -> Tensor:
1056
1057
        """
        Args:
vfdev's avatar
vfdev committed
1058
            tensor (Tensor): Tensor image to be whitened.
1059
1060
1061
1062

        Returns:
            Tensor: Transformed image.
        """
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
        shape = tensor.shape
        n = shape[-3] * shape[-2] * shape[-1]
        if n != self.transformation_matrix.shape[0]:
            raise ValueError("Input tensor and transformation matrix have incompatible shape." +
                             "[{} x {} x {}] != ".format(shape[-3], shape[-2], shape[-1]) +
                             "{}".format(self.transformation_matrix.shape[0]))

        if tensor.device.type != self.mean_vector.device.type:
            raise ValueError("Input tensor should be on the same device as transformation matrix and mean vector. "
                             "Got {} vs {}".format(tensor.device, self.mean_vector.device))

        flat_tensor = tensor.view(-1, n) - self.mean_vector
1075
        transformed_tensor = torch.mm(flat_tensor, self.transformation_matrix)
1076
        tensor = transformed_tensor.view(shape)
1077
1078
        return tensor

1079
    def __repr__(self):
ekka's avatar
ekka committed
1080
1081
1082
        format_string = self.__class__.__name__ + '(transformation_matrix='
        format_string += (str(self.transformation_matrix.tolist()) + ')')
        format_string += (", (mean_vector=" + str(self.mean_vector.tolist()) + ')')
1083
1084
        return format_string

1085

1086
class ColorJitter(torch.nn.Module):
1087
    """Randomly change the brightness, contrast, saturation and hue of an image.
1088
    If the image is torch Tensor, it is expected
1089
1090
    to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
    If img is PIL Image, mode "1", "L", "I", "F" and modes with transparency (alpha channel) are not supported.
1091
1092

    Args:
yaox12's avatar
yaox12 committed
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
        brightness (float or tuple of float (min, max)): How much to jitter brightness.
            brightness_factor is chosen uniformly from [max(0, 1 - brightness), 1 + brightness]
            or the given [min, max]. Should be non negative numbers.
        contrast (float or tuple of float (min, max)): How much to jitter contrast.
            contrast_factor is chosen uniformly from [max(0, 1 - contrast), 1 + contrast]
            or the given [min, max]. Should be non negative numbers.
        saturation (float or tuple of float (min, max)): How much to jitter saturation.
            saturation_factor is chosen uniformly from [max(0, 1 - saturation), 1 + saturation]
            or the given [min, max]. Should be non negative numbers.
        hue (float or tuple of float (min, max)): How much to jitter hue.
            hue_factor is chosen uniformly from [-hue, hue] or the given [min, max].
            Should have 0<= hue <= 0.5 or -0.5 <= min <= max <= 0.5.
1105
    """
1106

1107
    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):
1108
        super().__init__()
yaox12's avatar
yaox12 committed
1109
1110
1111
1112
1113
1114
        self.brightness = self._check_input(brightness, 'brightness')
        self.contrast = self._check_input(contrast, 'contrast')
        self.saturation = self._check_input(saturation, 'saturation')
        self.hue = self._check_input(hue, 'hue', center=0, bound=(-0.5, 0.5),
                                     clip_first_on_zero=False)

1115
    @torch.jit.unused
yaox12's avatar
yaox12 committed
1116
1117
1118
1119
    def _check_input(self, value, name, center=1, bound=(0, float('inf')), clip_first_on_zero=True):
        if isinstance(value, numbers.Number):
            if value < 0:
                raise ValueError("If {} is a single number, it must be non negative.".format(name))
1120
            value = [center - float(value), center + float(value)]
yaox12's avatar
yaox12 committed
1121
            if clip_first_on_zero:
1122
                value[0] = max(value[0], 0.0)
yaox12's avatar
yaox12 committed
1123
1124
1125
1126
        elif isinstance(value, (tuple, list)) and len(value) == 2:
            if not bound[0] <= value[0] <= value[1] <= bound[1]:
                raise ValueError("{} values should be between {}".format(name, bound))
        else:
1127
            raise TypeError("{} should be a single number or a list/tuple with length 2.".format(name))
yaox12's avatar
yaox12 committed
1128
1129
1130
1131
1132
1133

        # if value is 0 or (1., 1.) for brightness/contrast/saturation
        # or (0., 0.) for hue, do nothing
        if value[0] == value[1] == center:
            value = None
        return value
1134
1135

    @staticmethod
1136
1137
1138
1139
1140
1141
    def get_params(brightness: Optional[List[float]],
                   contrast: Optional[List[float]],
                   saturation: Optional[List[float]],
                   hue: Optional[List[float]]
                   ) -> Tuple[Tensor, Optional[float], Optional[float], Optional[float], Optional[float]]:
        """Get the parameters for the randomized transform to be applied on image.
1142

1143
1144
1145
1146
1147
1148
1149
1150
1151
        Args:
            brightness (tuple of float (min, max), optional): The range from which the brightness_factor is chosen
                uniformly. Pass None to turn off the transformation.
            contrast (tuple of float (min, max), optional): The range from which the contrast_factor is chosen
                uniformly. Pass None to turn off the transformation.
            saturation (tuple of float (min, max), optional): The range from which the saturation_factor is chosen
                uniformly. Pass None to turn off the transformation.
            hue (tuple of float (min, max), optional): The range from which the hue_factor is chosen uniformly.
                Pass None to turn off the transformation.
1152
1153

        Returns:
1154
1155
            tuple: The parameters used to apply the randomized transform
            along with their random order.
1156
        """
1157
        fn_idx = torch.randperm(4)
1158

1159
1160
1161
1162
        b = None if brightness is None else float(torch.empty(1).uniform_(brightness[0], brightness[1]))
        c = None if contrast is None else float(torch.empty(1).uniform_(contrast[0], contrast[1]))
        s = None if saturation is None else float(torch.empty(1).uniform_(saturation[0], saturation[1]))
        h = None if hue is None else float(torch.empty(1).uniform_(hue[0], hue[1]))
1163

1164
        return fn_idx, b, c, s, h
1165

1166
    def forward(self, img):
1167
1168
        """
        Args:
1169
            img (PIL Image or Tensor): Input image.
1170
1171

        Returns:
1172
1173
            PIL Image or Tensor: Color jittered image.
        """
1174
1175
1176
        fn_idx, brightness_factor, contrast_factor, saturation_factor, hue_factor = \
            self.get_params(self.brightness, self.contrast, self.saturation, self.hue)

1177
        for fn_id in fn_idx:
1178
            if fn_id == 0 and brightness_factor is not None:
1179
                img = F.adjust_brightness(img, brightness_factor)
1180
            elif fn_id == 1 and contrast_factor is not None:
1181
                img = F.adjust_contrast(img, contrast_factor)
1182
            elif fn_id == 2 and saturation_factor is not None:
1183
                img = F.adjust_saturation(img, saturation_factor)
1184
            elif fn_id == 3 and hue_factor is not None:
1185
1186
1187
                img = F.adjust_hue(img, hue_factor)

        return img
1188

1189
    def __repr__(self):
1190
1191
1192
1193
1194
1195
        format_string = self.__class__.__name__ + '('
        format_string += 'brightness={0}'.format(self.brightness)
        format_string += ', contrast={0}'.format(self.contrast)
        format_string += ', saturation={0}'.format(self.saturation)
        format_string += ', hue={0})'.format(self.hue)
        return format_string
1196

1197

1198
class RandomRotation(torch.nn.Module):
1199
    """Rotate the image by angle.
1200
    If the image is torch Tensor, it is expected
1201
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
1202
1203

    Args:
1204
        degrees (sequence or number): Range of degrees to select from.
1205
1206
            If degrees is a number instead of sequence like (min, max), the range of degrees
            will be (-degrees, +degrees).
1207
1208
1209
        interpolation (InterpolationMode): Desired interpolation enum defined by
            :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``.
            If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported.
1210
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
1211
1212
1213
1214
        expand (bool, optional): Optional expansion flag.
            If true, expands the output 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.
1215
        center (sequence, optional): Optional center of rotation, (x, y). Origin is the upper left corner.
1216
            Default is the center of the image.
1217
1218
        fill (sequence or number): Pixel fill value for the area outside the rotated
            image. Default is ``0``. If given a number, the value is used for all bands respectively.
1219
            If input is PIL Image, the options is only available for ``Pillow>=5.2.0``.
1220
        resample (int, optional): deprecated argument and will be removed since v0.10.0.
1221
            Please use the ``interpolation`` parameter instead.
1222
1223
1224

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

1225
1226
    """

1227
    def __init__(
1228
        self, degrees, interpolation=InterpolationMode.NEAREST, expand=False, center=None, fill=0, resample=None
1229
    ):
1230
        super().__init__()
1231
1232
1233
1234
1235
1236
1237
1238
1239
        if resample is not None:
            warnings.warn(
                "Argument resample is deprecated and will be removed since v0.10.0. Please, use interpolation instead"
            )
            interpolation = _interpolation_modes_from_int(resample)

        # Backward compatibility with integer value
        if isinstance(interpolation, int):
            warnings.warn(
1240
1241
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
1242
1243
1244
            )
            interpolation = _interpolation_modes_from_int(interpolation)

1245
        self.degrees = _setup_angle(degrees, name="degrees", req_sizes=(2, ))
1246
1247

        if center is not None:
1248
            _check_sequence_input(center, "center", req_sizes=(2, ))
1249
1250

        self.center = center
1251

1252
        self.resample = self.interpolation = interpolation
1253
        self.expand = expand
1254
1255
1256
1257
1258
1259

        if fill is None:
            fill = 0
        elif not isinstance(fill, (Sequence, numbers.Number)):
            raise TypeError("Fill should be either a sequence or a number.")

1260
        self.fill = fill
1261
1262

    @staticmethod
1263
    def get_params(degrees: List[float]) -> float:
1264
1265
1266
        """Get parameters for ``rotate`` for a random rotation.

        Returns:
1267
            float: angle parameter to be passed to ``rotate`` for random rotation.
1268
        """
1269
        angle = float(torch.empty(1).uniform_(float(degrees[0]), float(degrees[1])).item())
1270
1271
        return angle

1272
    def forward(self, img):
1273
        """
1274
        Args:
1275
            img (PIL Image or Tensor): Image to be rotated.
1276
1277

        Returns:
1278
            PIL Image or Tensor: Rotated image.
1279
        """
1280
1281
1282
1283
1284
1285
        fill = self.fill
        if isinstance(img, Tensor):
            if isinstance(fill, (int, float)):
                fill = [float(fill)] * F._get_image_num_channels(img)
            else:
                fill = [float(f) for f in fill]
1286
        angle = self.get_params(self.degrees)
1287
1288

        return F.rotate(img, angle, self.resample, self.expand, self.center, fill)
1289

1290
    def __repr__(self):
1291
        interpolate_str = self.interpolation.value
1292
        format_string = self.__class__.__name__ + '(degrees={0}'.format(self.degrees)
1293
        format_string += ', interpolation={0}'.format(interpolate_str)
1294
1295
1296
        format_string += ', expand={0}'.format(self.expand)
        if self.center is not None:
            format_string += ', center={0}'.format(self.center)
1297
1298
        if self.fill is not None:
            format_string += ', fill={0}'.format(self.fill)
1299
1300
        format_string += ')'
        return format_string
1301

1302

1303
1304
class RandomAffine(torch.nn.Module):
    """Random affine transformation of the image keeping center invariant.
1305
    If the image is torch Tensor, it is expected
1306
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
1307
1308

    Args:
1309
        degrees (sequence or number): Range of degrees to select from.
1310
            If degrees is a number instead of sequence like (min, max), the range of degrees
1311
            will be (-degrees, +degrees). Set to 0 to deactivate rotations.
1312
1313
1314
1315
1316
1317
        translate (tuple, optional): tuple of maximum absolute fraction for horizontal
            and vertical translations. For example translate=(a, b), then horizontal shift
            is randomly sampled in the range -img_width * a < dx < img_width * a and vertical shift is
            randomly sampled in the range -img_height * b < dy < img_height * b. Will not translate by default.
        scale (tuple, optional): scaling factor interval, e.g (a, b), then scale is
            randomly sampled from the range a <= scale <= b. Will keep original scale by default.
1318
        shear (sequence or number, optional): Range of degrees to select from.
ptrblck's avatar
ptrblck committed
1319
            If shear is a number, a shear parallel to the x axis in the range (-shear, +shear)
1320
1321
            will be applied. Else if shear is a sequence of 2 values a shear parallel to the x axis in the
            range (shear[0], shear[1]) will be applied. Else if shear is a sequence of 4 values,
ptrblck's avatar
ptrblck committed
1322
            a x-axis shear in (shear[0], shear[1]) and y-axis shear in (shear[2], shear[3]) will be applied.
1323
            Will not apply shear by default.
1324
1325
1326
        interpolation (InterpolationMode): Desired interpolation enum defined by
            :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``.
            If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported.
1327
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
1328
1329
        fill (sequence or number): Pixel fill value for the area outside the transformed
            image. Default is ``0``. If given a number, the value is used for all bands respectively.
1330
            If input is PIL Image, the options is only available for ``Pillow>=5.0.0``.
1331
        fillcolor (sequence or number, optional): deprecated argument and will be removed since v0.10.0.
1332
            Please use the ``fill`` parameter instead.
1333
        resample (int, optional): deprecated argument and will be removed since v0.10.0.
1334
            Please use the ``interpolation`` parameter instead.
1335
1336
1337

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

1338
1339
    """

1340
    def __init__(
1341
        self, degrees, translate=None, scale=None, shear=None, interpolation=InterpolationMode.NEAREST, fill=0,
1342
1343
        fillcolor=None, resample=None
    ):
1344
        super().__init__()
1345
1346
1347
1348
1349
1350
1351
1352
1353
        if resample is not None:
            warnings.warn(
                "Argument resample is deprecated and will be removed since v0.10.0. Please, use interpolation instead"
            )
            interpolation = _interpolation_modes_from_int(resample)

        # Backward compatibility with integer value
        if isinstance(interpolation, int):
            warnings.warn(
1354
1355
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
1356
1357
1358
1359
1360
1361
1362
1363
1364
            )
            interpolation = _interpolation_modes_from_int(interpolation)

        if fillcolor is not None:
            warnings.warn(
                "Argument fillcolor is deprecated and will be removed since v0.10.0. Please, use fill instead"
            )
            fill = fillcolor

1365
        self.degrees = _setup_angle(degrees, name="degrees", req_sizes=(2, ))
1366
1367

        if translate is not None:
1368
            _check_sequence_input(translate, "translate", req_sizes=(2, ))
1369
1370
1371
1372
1373
1374
            for t in translate:
                if not (0.0 <= t <= 1.0):
                    raise ValueError("translation values should be between 0 and 1")
        self.translate = translate

        if scale is not None:
1375
            _check_sequence_input(scale, "scale", req_sizes=(2, ))
1376
1377
1378
1379
1380
1381
            for s in scale:
                if s <= 0:
                    raise ValueError("scale values should be positive")
        self.scale = scale

        if shear is not None:
1382
            self.shear = _setup_angle(shear, name="shear", req_sizes=(2, 4))
1383
1384
1385
        else:
            self.shear = shear

1386
        self.resample = self.interpolation = interpolation
1387
1388
1389
1390
1391
1392

        if fill is None:
            fill = 0
        elif not isinstance(fill, (Sequence, numbers.Number)):
            raise TypeError("Fill should be either a sequence or a number.")

1393
        self.fillcolor = self.fill = fill
1394
1395

    @staticmethod
1396
1397
1398
1399
1400
1401
1402
    def get_params(
            degrees: List[float],
            translate: Optional[List[float]],
            scale_ranges: Optional[List[float]],
            shears: Optional[List[float]],
            img_size: List[int]
    ) -> Tuple[float, Tuple[int, int], float, Tuple[float, float]]:
1403
1404
1405
        """Get parameters for affine transformation

        Returns:
1406
            params to be passed to the affine transformation
1407
        """
1408
        angle = float(torch.empty(1).uniform_(float(degrees[0]), float(degrees[1])).item())
1409
        if translate is not None:
1410
1411
1412
1413
1414
            max_dx = float(translate[0] * img_size[0])
            max_dy = float(translate[1] * img_size[1])
            tx = int(round(torch.empty(1).uniform_(-max_dx, max_dx).item()))
            ty = int(round(torch.empty(1).uniform_(-max_dy, max_dy).item()))
            translations = (tx, ty)
1415
1416
1417
1418
        else:
            translations = (0, 0)

        if scale_ranges is not None:
1419
            scale = float(torch.empty(1).uniform_(scale_ranges[0], scale_ranges[1]).item())
1420
1421
1422
        else:
            scale = 1.0

1423
        shear_x = shear_y = 0.0
1424
        if shears is not None:
1425
1426
1427
1428
1429
            shear_x = float(torch.empty(1).uniform_(shears[0], shears[1]).item())
            if len(shears) == 4:
                shear_y = float(torch.empty(1).uniform_(shears[2], shears[3]).item())

        shear = (shear_x, shear_y)
1430
1431
1432

        return angle, translations, scale, shear

1433
    def forward(self, img):
1434
        """
1435
            img (PIL Image or Tensor): Image to be transformed.
1436
1437

        Returns:
1438
            PIL Image or Tensor: Affine transformed image.
1439
        """
1440
1441
1442
1443
1444
1445
        fill = self.fill
        if isinstance(img, Tensor):
            if isinstance(fill, (int, float)):
                fill = [float(fill)] * F._get_image_num_channels(img)
            else:
                fill = [float(f) for f in fill]
1446
1447
1448
1449

        img_size = F._get_image_size(img)

        ret = self.get_params(self.degrees, self.translate, self.scale, self.shear, img_size)
1450
1451

        return F.affine(img, *ret, interpolation=self.interpolation, fill=fill)
1452
1453
1454
1455
1456
1457
1458
1459
1460

    def __repr__(self):
        s = '{name}(degrees={degrees}'
        if self.translate is not None:
            s += ', translate={translate}'
        if self.scale is not None:
            s += ', scale={scale}'
        if self.shear is not None:
            s += ', shear={shear}'
1461
        if self.interpolation != InterpolationMode.NEAREST:
1462
1463
1464
            s += ', interpolation={interpolation}'
        if self.fill != 0:
            s += ', fill={fill}'
1465
1466
        s += ')'
        d = dict(self.__dict__)
1467
        d['interpolation'] = self.interpolation.value
1468
1469
1470
        return s.format(name=self.__class__.__name__, **d)


1471
class Grayscale(torch.nn.Module):
1472
    """Convert image to grayscale.
1473
1474
    If the image is torch Tensor, it is expected
    to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions
1475

1476
1477
1478
1479
    Args:
        num_output_channels (int): (1 or 3) number of channels desired for output image

    Returns:
1480
        PIL Image: Grayscale version of the input.
1481
1482
1483

        - If ``num_output_channels == 1`` : returned image is single channel
        - If ``num_output_channels == 3`` : returned image is 3 channel with r == g == b
1484
1485
1486
1487

    """

    def __init__(self, num_output_channels=1):
1488
        super().__init__()
1489
1490
        self.num_output_channels = num_output_channels

vfdev's avatar
vfdev committed
1491
    def forward(self, img):
1492
1493
        """
        Args:
1494
            img (PIL Image or Tensor): Image to be converted to grayscale.
1495
1496

        Returns:
1497
            PIL Image or Tensor: Grayscaled image.
1498
        """
1499
        return F.rgb_to_grayscale(img, num_output_channels=self.num_output_channels)
1500

1501
    def __repr__(self):
1502
        return self.__class__.__name__ + '(num_output_channels={0})'.format(self.num_output_channels)
1503

1504

1505
class RandomGrayscale(torch.nn.Module):
1506
    """Randomly convert image to grayscale with a probability of p (default 0.1).
1507
1508
    If the image is torch Tensor, it is expected
    to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions
1509

1510
1511
1512
1513
    Args:
        p (float): probability that image should be converted to grayscale.

    Returns:
1514
        PIL Image or Tensor: Grayscale version of the input image with probability p and unchanged
1515
1516
1517
        with probability (1-p).
        - If input image is 1 channel: grayscale version is 1 channel
        - If input image is 3 channel: grayscale version is 3 channel with r == g == b
1518
1519
1520
1521

    """

    def __init__(self, p=0.1):
1522
        super().__init__()
1523
1524
        self.p = p

vfdev's avatar
vfdev committed
1525
    def forward(self, img):
1526
1527
        """
        Args:
1528
            img (PIL Image or Tensor): Image to be converted to grayscale.
1529
1530

        Returns:
1531
            PIL Image or Tensor: Randomly grayscaled image.
1532
        """
1533
1534
1535
        num_output_channels = F._get_image_num_channels(img)
        if torch.rand(1) < self.p:
            return F.rgb_to_grayscale(img, num_output_channels=num_output_channels)
1536
        return img
1537
1538

    def __repr__(self):
1539
        return self.__class__.__name__ + '(p={0})'.format(self.p)
1540
1541


1542
class RandomErasing(torch.nn.Module):
1543
1544
    """ Randomly selects a rectangle region in an torch Tensor image and erases its pixels.
    This transform does not support PIL Image.
vfdev's avatar
vfdev committed
1545
    'Random Erasing Data Augmentation' by Zhong et al. See https://arxiv.org/abs/1708.04896
1546

1547
1548
1549
1550
1551
1552
1553
1554
    Args:
         p: probability that the random erasing operation will be performed.
         scale: range of proportion of erased area against input image.
         ratio: range of aspect ratio of erased area.
         value: erasing value. Default is 0. If a single int, it is used to
            erase all pixels. If a tuple of length 3, it is used to erase
            R, G, B channels respectively.
            If a str of 'random', erasing each pixel with random values.
Zhun Zhong's avatar
Zhun Zhong committed
1555
         inplace: boolean to make this transform inplace. Default set to False.
1556

1557
1558
    Returns:
        Erased Image.
1559

vfdev's avatar
vfdev committed
1560
    Example:
1561
        >>> transform = transforms.Compose([
1562
1563
1564
1565
        >>>   transforms.RandomHorizontalFlip(),
        >>>   transforms.ToTensor(),
        >>>   transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        >>>   transforms.RandomErasing(),
1566
1567
1568
        >>> ])
    """

Zhun Zhong's avatar
Zhun Zhong committed
1569
    def __init__(self, p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=0, inplace=False):
1570
1571
1572
1573
1574
1575
1576
1577
1578
        super().__init__()
        if not isinstance(value, (numbers.Number, str, tuple, list)):
            raise TypeError("Argument value should be either a number or str or a sequence")
        if isinstance(value, str) and value != "random":
            raise ValueError("If value is str, it should be 'random'")
        if not isinstance(scale, (tuple, list)):
            raise TypeError("Scale should be a sequence")
        if not isinstance(ratio, (tuple, list)):
            raise TypeError("Ratio should be a sequence")
1579
        if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
1580
            warnings.warn("Scale and ratio should be of kind (min, max)")
1581
        if scale[0] < 0 or scale[1] > 1:
1582
            raise ValueError("Scale should be between 0 and 1")
1583
        if p < 0 or p > 1:
1584
            raise ValueError("Random erasing probability should be between 0 and 1")
1585
1586
1587
1588
1589

        self.p = p
        self.scale = scale
        self.ratio = ratio
        self.value = value
1590
        self.inplace = inplace
1591
1592

    @staticmethod
1593
1594
1595
    def get_params(
            img: Tensor, scale: Tuple[float, float], ratio: Tuple[float, float], value: Optional[List[float]] = None
    ) -> Tuple[int, int, int, int, Tensor]:
1596
1597
1598
        """Get parameters for ``erase`` for a random erasing.

        Args:
vfdev's avatar
vfdev committed
1599
            img (Tensor): Tensor image to be erased.
1600
1601
            scale (sequence): range of proportion of erased area against input image.
            ratio (sequence): range of aspect ratio of erased area.
1602
1603
1604
            value (list, optional): erasing value. If None, it is interpreted as "random"
                (erasing each pixel with random values). If ``len(value)`` is 1, it is interpreted as a number,
                i.e. ``value[0]``.
1605
1606
1607
1608

        Returns:
            tuple: params (i, j, h, w, v) to be passed to ``erase`` for random erasing.
        """
vfdev's avatar
vfdev committed
1609
        img_c, img_h, img_w = img.shape[-3], img.shape[-2], img.shape[-1]
1610
        area = img_h * img_w
1611

1612
        log_ratio = torch.log(torch.tensor(ratio))
1613
        for _ in range(10):
1614
            erase_area = area * torch.empty(1).uniform_(scale[0], scale[1]).item()
1615
1616
1617
            aspect_ratio = torch.exp(
                torch.empty(1).uniform_(log_ratio[0], log_ratio[1])
            ).item()
1618
1619
1620

            h = int(round(math.sqrt(erase_area * aspect_ratio)))
            w = int(round(math.sqrt(erase_area / aspect_ratio)))
1621
1622
1623
1624
1625
1626
1627
            if not (h < img_h and w < img_w):
                continue

            if value is None:
                v = torch.empty([img_c, h, w], dtype=torch.float32).normal_()
            else:
                v = torch.tensor(value)[:, None, None]
1628

1629
1630
            i = torch.randint(0, img_h - h + 1, size=(1, )).item()
            j = torch.randint(0, img_w - w + 1, size=(1, )).item()
1631
            return i, j, h, w, v
1632

Zhun Zhong's avatar
Zhun Zhong committed
1633
1634
1635
        # Return original image
        return 0, 0, img_h, img_w, img

1636
    def forward(self, img):
1637
1638
        """
        Args:
vfdev's avatar
vfdev committed
1639
            img (Tensor): Tensor image to be erased.
1640
1641
1642
1643

        Returns:
            img (Tensor): Erased Tensor image.
        """
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
        if torch.rand(1) < self.p:

            # cast self.value to script acceptable type
            if isinstance(self.value, (int, float)):
                value = [self.value, ]
            elif isinstance(self.value, str):
                value = None
            elif isinstance(self.value, tuple):
                value = list(self.value)
            else:
                value = self.value

            if value is not None and not (len(value) in (1, img.shape[-3])):
                raise ValueError(
                    "If value is a sequence, it should have either a single value or "
                    "{} (number of input channels)".format(img.shape[-3])
                )

            x, y, h, w, v = self.get_params(img, scale=self.scale, ratio=self.ratio, value=value)
1663
            return F.erase(img, x, y, h, w, v, self.inplace)
1664
        return img
1665

1666
1667
1668
1669
1670
1671
1672
1673
    def __repr__(self):
        s = '(p={}, '.format(self.p)
        s += 'scale={}, '.format(self.scale)
        s += 'ratio={}, '.format(self.ratio)
        s += 'value={}, '.format(self.value)
        s += 'inplace={})'.format(self.inplace)
        return self.__class__.__name__ + s

1674

1675
1676
class GaussianBlur(torch.nn.Module):
    """Blurs image with randomly chosen Gaussian blur.
1677
1678
    If the image is torch Tensor, it is expected
    to have [..., C, H, W] shape, where ... means an arbitrary number of leading dimensions.
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712

    Args:
        kernel_size (int or sequence): Size of the Gaussian kernel.
        sigma (float or tuple of float (min, max)): Standard deviation to be used for
            creating kernel to perform blurring. If float, sigma is fixed. If it is tuple
            of float (min, max), sigma is chosen uniformly at random to lie in the
            given range.

    Returns:
        PIL Image or Tensor: Gaussian blurred version of the input image.

    """

    def __init__(self, kernel_size, sigma=(0.1, 2.0)):
        super().__init__()
        self.kernel_size = _setup_size(kernel_size, "Kernel size should be a tuple/list of two integers")
        for ks in self.kernel_size:
            if ks <= 0 or ks % 2 == 0:
                raise ValueError("Kernel size value should be an odd and positive number.")

        if isinstance(sigma, numbers.Number):
            if sigma <= 0:
                raise ValueError("If sigma is a single number, it must be positive.")
            sigma = (sigma, sigma)
        elif isinstance(sigma, Sequence) and len(sigma) == 2:
            if not 0. < sigma[0] <= sigma[1]:
                raise ValueError("sigma values should be positive and of the form (min, max).")
        else:
            raise ValueError("sigma should be a single number or a list/tuple with length 2.")

        self.sigma = sigma

    @staticmethod
    def get_params(sigma_min: float, sigma_max: float) -> float:
vfdev's avatar
vfdev committed
1713
        """Choose sigma for random gaussian blurring.
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726

        Args:
            sigma_min (float): Minimum standard deviation that can be chosen for blurring kernel.
            sigma_max (float): Maximum standard deviation that can be chosen for blurring kernel.

        Returns:
            float: Standard deviation to be passed to calculate kernel for gaussian blurring.
        """
        return torch.empty(1).uniform_(sigma_min, sigma_max).item()

    def forward(self, img: Tensor) -> Tensor:
        """
        Args:
vfdev's avatar
vfdev committed
1727
            img (PIL Image or Tensor): image to be blurred.
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740

        Returns:
            PIL Image or Tensor: Gaussian blurred image
        """
        sigma = self.get_params(self.sigma[0], self.sigma[1])
        return F.gaussian_blur(img, self.kernel_size, [sigma, sigma])

    def __repr__(self):
        s = '(kernel_size={}, '.format(self.kernel_size)
        s += 'sigma={})'.format(self.sigma)
        return self.__class__.__name__ + s


1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
def _setup_size(size, error_msg):
    if isinstance(size, numbers.Number):
        return int(size), int(size)

    if isinstance(size, Sequence) and len(size) == 1:
        return size[0], size[0]

    if len(size) != 2:
        raise ValueError(error_msg)

    return size


def _check_sequence_input(x, name, req_sizes):
    msg = req_sizes[0] if len(req_sizes) < 2 else " or ".join([str(s) for s in req_sizes])
    if not isinstance(x, Sequence):
        raise TypeError("{} should be a sequence of length {}.".format(name, msg))
    if len(x) not in req_sizes:
        raise ValueError("{} should be sequence of length {}.".format(name, msg))


def _setup_angle(x, name, req_sizes=(2, )):
    if isinstance(x, numbers.Number):
        if x < 0:
            raise ValueError("If {} is a single number, it must be positive.".format(name))
        x = [-x, x]
    else:
        _check_sequence_input(x, name, req_sizes)

    return [float(d) for d in x]
1771
1772
1773
1774


class RandomInvert(torch.nn.Module):
    """Inverts the colors of the given image randomly with a given probability.
1775
1776
1777
    If img is a Tensor, it is expected to be in [..., 1 or 3, H, W] format,
    where ... means it can have an arbitrary number of leading dimensions.
    If img is PIL Image, it is expected to be in mode "L" or "RGB".
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804

    Args:
        p (float): probability of the image being color inverted. Default value is 0.5
    """

    def __init__(self, p=0.5):
        super().__init__()
        self.p = p

    def forward(self, img):
        """
        Args:
            img (PIL Image or Tensor): Image to be inverted.

        Returns:
            PIL Image or Tensor: Randomly color inverted image.
        """
        if torch.rand(1).item() < self.p:
            return F.invert(img)
        return img

    def __repr__(self):
        return self.__class__.__name__ + '(p={})'.format(self.p)


class RandomPosterize(torch.nn.Module):
    """Posterize the image randomly with a given probability by reducing the
1805
1806
1807
    number of bits for each color channel. If the image is torch Tensor, it should be of type torch.uint8,
    and it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
    If img is PIL Image, it is expected to be in mode "L" or "RGB".
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836

    Args:
        bits (int): number of bits to keep for each channel (0-8)
        p (float): probability of the image being color inverted. Default value is 0.5
    """

    def __init__(self, bits, p=0.5):
        super().__init__()
        self.bits = bits
        self.p = p

    def forward(self, img):
        """
        Args:
            img (PIL Image or Tensor): Image to be posterized.

        Returns:
            PIL Image or Tensor: Randomly posterized image.
        """
        if torch.rand(1).item() < self.p:
            return F.posterize(img, self.bits)
        return img

    def __repr__(self):
        return self.__class__.__name__ + '(bits={},p={})'.format(self.bits, self.p)


class RandomSolarize(torch.nn.Module):
    """Solarize the image randomly with a given probability by inverting all pixel
1837
1838
1839
    values above a threshold. If img is a Tensor, it is expected to be in [..., 1 or 3, H, W] format,
    where ... means it can have an arbitrary number of leading dimensions.
    If img is PIL Image, it is expected to be in mode "L" or "RGB".
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867

    Args:
        threshold (float): all pixels equal or above this value are inverted.
        p (float): probability of the image being color inverted. Default value is 0.5
    """

    def __init__(self, threshold, p=0.5):
        super().__init__()
        self.threshold = threshold
        self.p = p

    def forward(self, img):
        """
        Args:
            img (PIL Image or Tensor): Image to be solarized.

        Returns:
            PIL Image or Tensor: Randomly solarized image.
        """
        if torch.rand(1).item() < self.p:
            return F.solarize(img, self.threshold)
        return img

    def __repr__(self):
        return self.__class__.__name__ + '(threshold={},p={})'.format(self.threshold, self.p)


class RandomAdjustSharpness(torch.nn.Module):
1868
1869
    """Adjust the sharpness of the image randomly with a given probability. If the image is torch Tensor,
    it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900

    Args:
        sharpness_factor (float):  How much to adjust the sharpness. Can be
            any non negative number. 0 gives a blurred image, 1 gives the
            original image while 2 increases the sharpness by a factor of 2.
        p (float): probability of the image being color inverted. Default value is 0.5
    """

    def __init__(self, sharpness_factor, p=0.5):
        super().__init__()
        self.sharpness_factor = sharpness_factor
        self.p = p

    def forward(self, img):
        """
        Args:
            img (PIL Image or Tensor): Image to be sharpened.

        Returns:
            PIL Image or Tensor: Randomly sharpened image.
        """
        if torch.rand(1).item() < self.p:
            return F.adjust_sharpness(img, self.sharpness_factor)
        return img

    def __repr__(self):
        return self.__class__.__name__ + '(sharpness_factor={},p={})'.format(self.sharpness_factor, self.p)


class RandomAutocontrast(torch.nn.Module):
    """Autocontrast the pixels of the given image randomly with a given probability.
1901
1902
1903
    If the image is torch Tensor, it is expected
    to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
    If img is PIL Image, it is expected to be in mode "L" or "RGB".
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930

    Args:
        p (float): probability of the image being autocontrasted. Default value is 0.5
    """

    def __init__(self, p=0.5):
        super().__init__()
        self.p = p

    def forward(self, img):
        """
        Args:
            img (PIL Image or Tensor): Image to be autocontrasted.

        Returns:
            PIL Image or Tensor: Randomly autocontrasted image.
        """
        if torch.rand(1).item() < self.p:
            return F.autocontrast(img)
        return img

    def __repr__(self):
        return self.__class__.__name__ + '(p={})'.format(self.p)


class RandomEqualize(torch.nn.Module):
    """Equalize the histogram of the given image randomly with a given probability.
1931
1932
1933
    If the image is torch Tensor, it is expected
    to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
    If img is PIL Image, it is expected to be in mode "P", "L" or "RGB".
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956

    Args:
        p (float): probability of the image being equalized. Default value is 0.5
    """

    def __init__(self, p=0.5):
        super().__init__()
        self.p = p

    def forward(self, img):
        """
        Args:
            img (PIL Image or Tensor): Image to be equalized.

        Returns:
            PIL Image or Tensor: Randomly equalized image.
        """
        if torch.rand(1).item() < self.p:
            return F.equalize(img)
        return img

    def __repr__(self):
        return self.__class__.__name__ + '(p={})'.format(self.p)