transforms.py 74.3 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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
    """Convert a tensor image to the given ``dtype`` and scale the values accordingly

    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:
142
        super().__init__()
143
144
        self.dtype = dtype

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


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

    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
158
159
160
161
162
            - 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``).
163

csukuangfj's avatar
csukuangfj committed
164
    .. _PIL.Image mode: https://pillow.readthedocs.io/en/latest/handbook/concepts.html#concept-modes
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
    """
    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)

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

187

188
class Normalize(torch.nn.Module):
Fang Gao's avatar
Fang Gao committed
189
    """Normalize a tensor image with mean and standard deviation.
190
191
192
    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
193
    ``output[channel] = (input[channel] - mean[channel]) / std[channel]``
194

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

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

203
204
    """

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

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

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

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

224

vfdev's avatar
vfdev committed
225
226
227
228
class Resize(torch.nn.Module):
    """Resize the input image to the given size.
    The image can be a PIL Image or a torch Tensor, in which case it is expected
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions
229
230
231
232
233
234

    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
235
236
237
            (size * height / width, size).
            In torchscript mode padding as single int is not supported, use a tuple or
            list of length 1: ``[size, ]``.
238
239
240
241
        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.
242
243
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.

244
245
    """

246
    def __init__(self, size, interpolation=InterpolationMode.BILINEAR):
vfdev's avatar
vfdev committed
247
        super().__init__()
248
249
250
251
252
        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
253
254
255
256

        # Backward compatibility with integer value
        if isinstance(interpolation, int):
            warnings.warn(
257
258
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
259
260
261
            )
            interpolation = _interpolation_modes_from_int(interpolation)

262
263
        self.interpolation = interpolation

vfdev's avatar
vfdev committed
264
    def forward(self, img):
265
266
        """
        Args:
vfdev's avatar
vfdev committed
267
            img (PIL Image or Tensor): Image to be scaled.
268
269

        Returns:
vfdev's avatar
vfdev committed
270
            PIL Image or Tensor: Rescaled image.
271
272
273
        """
        return F.resize(img, self.size, self.interpolation)

274
    def __repr__(self):
275
        interpolate_str = self.interpolation.value
276
        return self.__class__.__name__ + '(size={0}, interpolation={1})'.format(self.size, interpolate_str)
277

278
279
280
281
282
283
284
285
286
287
288

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
289
290
291
292
class CenterCrop(torch.nn.Module):
    """Crops the given image at the center.
    The image can be a PIL Image or a torch Tensor, in which case it is expected
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions
293
294
295
296

    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
vfdev's avatar
vfdev committed
297
            made. If provided a tuple or list of length 1, it will be interpreted as (size[0], size[0]).
298
299
300
    """

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

vfdev's avatar
vfdev committed
304
    def forward(self, img):
305
306
        """
        Args:
vfdev's avatar
vfdev committed
307
            img (PIL Image or Tensor): Image to be cropped.
308
309

        Returns:
vfdev's avatar
vfdev committed
310
            PIL Image or Tensor: Cropped image.
311
312
313
        """
        return F.center_crop(img, self.size)

314
315
316
    def __repr__(self):
        return self.__class__.__name__ + '(size={0})'.format(self.size)

317

318
319
320
321
class Pad(torch.nn.Module):
    """Pad the given image on all sides with the given "pad" value.
    The image can be a PIL Image or a torch Tensor, in which case it is expected
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions
322
323

    Args:
324
        padding (int or tuple or list): Padding on each border. If a single int is provided this
325
326
            is used to pad all borders. If tuple of length 2 is provided this is the padding
            on left/right and top/bottom respectively. If a tuple of length 4 is provided
327
328
329
            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, ]``.
330
        fill (int or tuple): Pixel fill value for constant fill. Default is 0. If a tuple of
331
            length 3, it is used to fill R, G, B channels respectively.
332
            This value is only used when the padding_mode is constant
333
        padding_mode (str): Type of padding. Should be: constant, edge, reflect or symmetric.
vfdev's avatar
vfdev committed
334
            Default is constant. Mode symmetric is not yet supported for Tensor inputs.
335
336
337
338
339
340
341
342

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

            - edge: pads with the last value at the edge of the image

            - 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
343
                will result in [3, 2, 1, 2, 3, 4, 3, 2]
344
345
346
347

            - 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
348
                will result in [2, 1, 1, 2, 3, 4, 4, 3]
349
350
    """

351
352
353
354
355
356
357
358
359
360
361
362
363
    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 " +
364
365
366
367
                             "{} element tuple".format(len(padding)))

        self.padding = padding
        self.fill = fill
368
        self.padding_mode = padding_mode
369

370
    def forward(self, img):
371
372
        """
        Args:
373
            img (PIL Image or Tensor): Image to be padded.
374
375

        Returns:
376
            PIL Image or Tensor: Padded image.
377
        """
378
        return F.pad(img, self.padding, self.fill, self.padding_mode)
379

380
    def __repr__(self):
381
382
        return self.__class__.__name__ + '(padding={0}, fill={1}, padding_mode={2})'.\
            format(self.padding, self.fill, self.padding_mode)
383

384

385
class Lambda:
386
    """Apply a user-defined lambda as a transform. This transform does not support torchscript.
387
388
389
390
391
392

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

    def __init__(self, lambd):
393
394
        if not callable(lambd):
            raise TypeError("Argument lambd should be callable, got {}".format(repr(type(lambd).__name__)))
395
396
397
398
399
        self.lambd = lambd

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

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

403

404
class RandomTransforms:
405
406
407
408
409
410
411
    """Base class for a list of transformations with randomness

    Args:
        transforms (list or tuple): list of transformations
    """

    def __init__(self, transforms):
412
413
        if not isinstance(transforms, Sequence):
            raise TypeError("Argument transforms should be a sequence")
414
415
416
417
418
419
420
421
422
423
424
425
426
427
        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


428
class RandomApply(torch.nn.Module):
429
    """Apply randomly a list of transformations with a given probability.
430
431
432
433
434
435
436
437
438
439
440
441

    .. 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``.
442
443

    Args:
444
        transforms (list or tuple or torch.nn.Module): list of transformations
445
446
447
448
        p (float): probability
    """

    def __init__(self, transforms, p=0.5):
449
450
        super().__init__()
        self.transforms = transforms
451
452
        self.p = p

453
454
    def forward(self, img):
        if self.p < torch.rand(1):
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
            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):
471
    """Apply a list of transformations in a random order. This transform does not support torchscript.
472
473
474
475
476
477
478
479
480
481
    """
    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):
482
    """Apply single transformation randomly picked from a list. This transform does not support torchscript.
483
484
485
486
487
488
    """
    def __call__(self, img):
        t = random.choice(self.transforms)
        return t(img)


vfdev's avatar
vfdev committed
489
490
491
492
493
class RandomCrop(torch.nn.Module):
    """Crop the given image at a random location.
    The image can be a PIL Image or a Tensor, in which case it is expected
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
494
495
496
497

    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
vfdev's avatar
vfdev committed
498
            made. If provided a tuple or list of length 1, it will be interpreted as (size[0], size[0]).
499
        padding (int or sequence, optional): Optional padding on each border
vfdev's avatar
vfdev committed
500
501
502
503
504
505
            of the image. Default is None. If a single int is provided this
            is used to pad all borders. If tuple of length 2 is provided this is the padding
            on left/right and top/bottom respectively. If a tuple 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, ]``.
506
        pad_if_needed (boolean): It will pad the image if smaller than the
ekka's avatar
ekka committed
507
            desired size to avoid raising an exception. Since cropping is done
508
            after padding, the padding seems to be done at a random offset.
vfdev's avatar
vfdev committed
509
        fill (int or tuple): Pixel fill value for constant fill. Default is 0. If a tuple of
510
511
            length 3, it is used to fill R, G, B channels respectively.
            This value is only used when the padding_mode is constant
vfdev's avatar
vfdev committed
512
        padding_mode (str): Type of padding. Should be: constant, edge, reflect or symmetric. Default is constant.
vfdev's avatar
vfdev committed
513
            Mode symmetric is not yet supported for Tensor inputs.
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528

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

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

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

529
530
531
    """

    @staticmethod
vfdev's avatar
vfdev committed
532
    def get_params(img: Tensor, output_size: Tuple[int, int]) -> Tuple[int, int, int, int]:
533
534
535
        """Get parameters for ``crop`` for a random crop.

        Args:
vfdev's avatar
vfdev committed
536
            img (PIL Image or Tensor): Image to be cropped.
537
538
539
540
541
            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
542
        w, h = F._get_image_size(img)
543
        th, tw = output_size
vfdev's avatar
vfdev committed
544
545
546
547
548
549

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

550
551
552
        if w == tw and h == th:
            return 0, 0, h, w

553
554
        i = torch.randint(0, h - th + 1, size=(1, )).item()
        j = torch.randint(0, w - tw + 1, size=(1, )).item()
555
556
        return i, j, th, tw

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

560
561
562
563
        self.size = tuple(_setup_size(
            size, error_msg="Please provide only two dimensions (h, w) for size."
        ))

vfdev's avatar
vfdev committed
564
565
566
567
568
569
        self.padding = padding
        self.pad_if_needed = pad_if_needed
        self.fill = fill
        self.padding_mode = padding_mode

    def forward(self, img):
570
571
        """
        Args:
vfdev's avatar
vfdev committed
572
            img (PIL Image or Tensor): Image to be cropped.
573
574

        Returns:
vfdev's avatar
vfdev committed
575
            PIL Image or Tensor: Cropped image.
576
        """
577
578
        if self.padding is not None:
            img = F.pad(img, self.padding, self.fill, self.padding_mode)
579

vfdev's avatar
vfdev committed
580
        width, height = F._get_image_size(img)
581
        # pad the width if needed
vfdev's avatar
vfdev committed
582
583
584
        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)
585
        # pad the height if needed
vfdev's avatar
vfdev committed
586
587
588
        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)
589

590
591
592
593
        i, j, h, w = self.get_params(img, self.size)

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

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

597

598
599
600
601
602
class RandomHorizontalFlip(torch.nn.Module):
    """Horizontally flip the given image randomly with a given probability.
    The image can be a PIL Image or a torch Tensor, in which case it is expected
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
603
604
605
606
607
608

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

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

612
    def forward(self, img):
613
614
        """
        Args:
615
            img (PIL Image or Tensor): Image to be flipped.
616
617

        Returns:
618
            PIL Image or Tensor: Randomly flipped image.
619
        """
620
        if torch.rand(1) < self.p:
621
622
623
            return F.hflip(img)
        return img

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

627

628
class RandomVerticalFlip(torch.nn.Module):
vfdev's avatar
vfdev committed
629
    """Vertically flip the given image randomly with a given probability.
630
631
632
    The image can be a PIL Image or a torch Tensor, in which case it is expected
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
633
634
635
636
637
638

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

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

642
    def forward(self, img):
643
644
        """
        Args:
645
            img (PIL Image or Tensor): Image to be flipped.
646
647

        Returns:
648
            PIL Image or Tensor: Randomly flipped image.
649
        """
650
        if torch.rand(1) < self.p:
651
652
653
            return F.vflip(img)
        return img

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

657

658
659
660
661
class RandomPerspective(torch.nn.Module):
    """Performs a random perspective transformation of the given image with a given probability.
    The image can be a PIL Image or a Tensor, in which case it is expected
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
662
663

    Args:
664
665
666
        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.
667
668
669
        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.
670
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
671
672
673
674
        fill (sequence or int or float, optional): Pixel fill value for the area outside the transformed
            image. If int or float, the value is used for all bands respectively.
            This option is supported for PIL image and Tensor inputs.
            If input is PIL Image, the options is only available for ``Pillow>=5.0.0``.
675
676
    """

677
    def __init__(self, distortion_scale=0.5, p=0.5, interpolation=InterpolationMode.BILINEAR, fill=0):
678
        super().__init__()
679
        self.p = p
680
681
682
683

        # Backward compatibility with integer value
        if isinstance(interpolation, int):
            warnings.warn(
684
685
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
686
687
688
            )
            interpolation = _interpolation_modes_from_int(interpolation)

689
690
        self.interpolation = interpolation
        self.distortion_scale = distortion_scale
691
        self.fill = fill
692

693
    def forward(self, img):
694
695
        """
        Args:
696
            img (PIL Image or Tensor): Image to be Perspectively transformed.
697
698

        Returns:
699
            PIL Image or Tensor: Randomly transformed image.
700
        """
701
702
703
704
705
706
707
708

        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]

709
710
        if torch.rand(1) < self.p:
            width, height = F._get_image_size(img)
711
            startpoints, endpoints = self.get_params(width, height, self.distortion_scale)
712
            return F.perspective(img, startpoints, endpoints, self.interpolation, fill)
713
714
715
        return img

    @staticmethod
716
    def get_params(width: int, height: int, distortion_scale: float) -> Tuple[List[List[int]], List[List[int]]]:
717
718
719
        """Get parameters for ``perspective`` for a random perspective transform.

        Args:
720
721
722
            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.
723
724

        Returns:
725
            List containing [top-left, top-right, bottom-right, bottom-left] of the original image,
726
727
            List containing [top-left, top-right, bottom-right, bottom-left] of the transformed image.
        """
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
        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]]
747
748
749
750
751
752
753
        endpoints = [topleft, topright, botright, botleft]
        return startpoints, endpoints

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


754
755
756
757
class RandomResizedCrop(torch.nn.Module):
    """Crop the given image to random size and aspect ratio.
    The image can be a PIL Image or a Tensor, in which case it is expected
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions
758

759
760
    A crop of random size (default: of 0.08 to 1.0) of the original size and a random
    aspect ratio (default: of 3/4 to 4/3) of the original aspect ratio is made. This crop
761
762
763
764
    is finally resized to given size.
    This is popularly used to train the Inception networks.

    Args:
765
766
767
        size (int or sequence): expected output size of each edge. If size is an
            int instead of sequence like (h, w), a square output size ``(size, size)`` is
            made. If provided a tuple or list of length 1, it will be interpreted as (size[0], size[0]).
768
769
        scale (tuple of float): scale range of the cropped image before resizing, relatively to the origin image.
        ratio (tuple of float): aspect ratio range of the cropped image before resizing.
770
771
772
773
        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.
774
775
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.

776
777
    """

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

782
        if not isinstance(scale, Sequence):
783
            raise TypeError("Scale should be a sequence")
784
        if not isinstance(ratio, Sequence):
785
            raise TypeError("Ratio should be a sequence")
786
        if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
787
            warnings.warn("Scale and ratio should be of kind (min, max)")
788

789
790
791
        # Backward compatibility with integer value
        if isinstance(interpolation, int):
            warnings.warn(
792
793
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
794
795
796
            )
            interpolation = _interpolation_modes_from_int(interpolation)

797
        self.interpolation = interpolation
798
799
        self.scale = scale
        self.ratio = ratio
800
801

    @staticmethod
802
    def get_params(
803
            img: Tensor, scale: List[float], ratio: List[float]
804
    ) -> Tuple[int, int, int, int]:
805
806
807
        """Get parameters for ``crop`` for a random sized crop.

        Args:
808
            img (PIL Image or Tensor): Input image.
809
810
            scale (list): range of scale of the origin size cropped
            ratio (list): range of aspect ratio of the origin aspect ratio cropped
811
812
813
814
815

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

819
        for _ in range(10):
820
            target_area = area * torch.empty(1).uniform_(scale[0], scale[1]).item()
821
822
823
824
            log_ratio = torch.log(torch.tensor(ratio))
            aspect_ratio = torch.exp(
                torch.empty(1).uniform_(log_ratio[0], log_ratio[1])
            ).item()
825
826
827
828

            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
829
            if 0 < w <= width and 0 < h <= height:
830
831
                i = torch.randint(0, height - h + 1, size=(1,)).item()
                j = torch.randint(0, width - w + 1, size=(1,)).item()
832
833
                return i, j, h, w

834
        # Fallback to central crop
Zhicheng Yan's avatar
Zhicheng Yan committed
835
        in_ratio = float(width) / float(height)
836
        if in_ratio < min(ratio):
Zhicheng Yan's avatar
Zhicheng Yan committed
837
            w = width
838
            h = int(round(w / min(ratio)))
839
        elif in_ratio > max(ratio):
Zhicheng Yan's avatar
Zhicheng Yan committed
840
            h = height
841
            w = int(round(h * max(ratio)))
842
        else:  # whole image
Zhicheng Yan's avatar
Zhicheng Yan committed
843
844
845
846
            w = width
            h = height
        i = (height - h) // 2
        j = (width - w) // 2
847
        return i, j, h, w
848

849
    def forward(self, img):
850
851
        """
        Args:
852
            img (PIL Image or Tensor): Image to be cropped and resized.
853
854

        Returns:
855
            PIL Image or Tensor: Randomly cropped and resized image.
856
        """
857
        i, j, h, w = self.get_params(img, self.scale, self.ratio)
858
859
        return F.resized_crop(img, i, j, h, w, self.size, self.interpolation)

860
    def __repr__(self):
861
        interpolate_str = self.interpolation.value
862
        format_string = self.__class__.__name__ + '(size={0}'.format(self.size)
863
864
        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))
865
866
        format_string += ', interpolation={0})'.format(interpolate_str)
        return format_string
867

868
869
870
871
872
873
874
875
876
877
878

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
879
880
881
882
883
class FiveCrop(torch.nn.Module):
    """Crop the given image into four corners and the central crop.
    The image can be a PIL Image or a Tensor, in which case it is expected
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
884
885
886
887
888
889
890
891
892

    .. 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.
vfdev's avatar
vfdev committed
893
            If provided a tuple or list of length 1, it will be interpreted as (size[0], size[0]).
894
895
896
897
898
899
900
901
902
903
904
905
906
907

    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
908
        super().__init__()
909
        self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
910

vfdev's avatar
vfdev committed
911
912
913
914
915
916
917
918
    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
        """
919
920
        return F.five_crop(img, self.size)

921
922
923
    def __repr__(self):
        return self.__class__.__name__ + '(size={0})'.format(self.size)

924

vfdev's avatar
vfdev committed
925
926
927
928
929
930
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).
    The image can be a PIL Image or a Tensor, in which case it is expected
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
931
932
933
934
935
936
937
938
939

    .. 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
vfdev's avatar
vfdev committed
940
            made. If provided a tuple or list of length 1, it will be interpreted as (size[0], size[0]).
941
        vertical_flip (bool): Use vertical flipping instead of horizontal
942
943
944
945
946
947
948
949
950
951
952
953
954
955

    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
956
        super().__init__()
957
        self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
958
959
        self.vertical_flip = vertical_flip

vfdev's avatar
vfdev committed
960
961
962
963
964
965
966
967
    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
        """
968
969
        return F.ten_crop(img, self.size, self.vertical_flip)

970
    def __repr__(self):
971
        return self.__class__.__name__ + '(size={0}, vertical_flip={1})'.format(self.size, self.vertical_flip)
972

973

974
class LinearTransformation(torch.nn.Module):
ekka's avatar
ekka committed
975
    """Transform a tensor image with a square transformation matrix and a mean_vector computed
976
    offline.
ekka's avatar
ekka committed
977
978
979
    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
980
    original shape.
981

982
    Applications:
983
        whitening transformation: Suppose X is a column vector zero-centered data.
984
985
986
        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.

987
988
    Args:
        transformation_matrix (Tensor): tensor [D x D], D = C x H x W
ekka's avatar
ekka committed
989
        mean_vector (Tensor): tensor [D], D = C x H x W
990
991
    """

ekka's avatar
ekka committed
992
    def __init__(self, transformation_matrix, mean_vector):
993
        super().__init__()
994
995
996
        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
997
998
999

        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
1000
1001
                             " as any one of the dimensions of the transformation_matrix [{}]"
                             .format(tuple(transformation_matrix.size())))
ekka's avatar
ekka committed
1002

1003
1004
1005
1006
        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))

1007
        self.transformation_matrix = transformation_matrix
ekka's avatar
ekka committed
1008
        self.mean_vector = mean_vector
1009

1010
    def forward(self, tensor: Tensor) -> Tensor:
1011
1012
        """
        Args:
vfdev's avatar
vfdev committed
1013
            tensor (Tensor): Tensor image to be whitened.
1014
1015
1016
1017

        Returns:
            Tensor: Transformed image.
        """
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
        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
1030
        transformed_tensor = torch.mm(flat_tensor, self.transformation_matrix)
1031
        tensor = transformed_tensor.view(shape)
1032
1033
        return tensor

1034
    def __repr__(self):
ekka's avatar
ekka committed
1035
1036
1037
        format_string = self.__class__.__name__ + '(transformation_matrix='
        format_string += (str(self.transformation_matrix.tolist()) + ')')
        format_string += (", (mean_vector=" + str(self.mean_vector.tolist()) + ')')
1038
1039
        return format_string

1040

1041
class ColorJitter(torch.nn.Module):
1042
    """Randomly change the brightness, contrast, saturation and hue of an image.
1043
1044

    Args:
yaox12's avatar
yaox12 committed
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
        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.
1057
    """
1058

1059
    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):
1060
        super().__init__()
yaox12's avatar
yaox12 committed
1061
1062
1063
1064
1065
1066
        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)

1067
    @torch.jit.unused
yaox12's avatar
yaox12 committed
1068
1069
1070
1071
    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))
1072
            value = [center - float(value), center + float(value)]
yaox12's avatar
yaox12 committed
1073
            if clip_first_on_zero:
1074
                value[0] = max(value[0], 0.0)
yaox12's avatar
yaox12 committed
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
        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:
            raise TypeError("{} should be a single number or a list/tuple with lenght 2.".format(name))

        # 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
1086
1087

    @staticmethod
1088
1089
1090
1091
1092
1093
    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.
1094

1095
1096
1097
1098
1099
1100
1101
1102
1103
        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.
1104
1105

        Returns:
1106
1107
            tuple: The parameters used to apply the randomized transform
            along with their random order.
1108
        """
1109
        fn_idx = torch.randperm(4)
1110

1111
1112
1113
1114
        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]))
1115

1116
        return fn_idx, b, c, s, h
1117

1118
    def forward(self, img):
1119
1120
        """
        Args:
1121
            img (PIL Image or Tensor): Input image.
1122
1123

        Returns:
1124
1125
            PIL Image or Tensor: Color jittered image.
        """
1126
1127
1128
        fn_idx, brightness_factor, contrast_factor, saturation_factor, hue_factor = \
            self.get_params(self.brightness, self.contrast, self.saturation, self.hue)

1129
        for fn_id in fn_idx:
1130
            if fn_id == 0 and brightness_factor is not None:
1131
                img = F.adjust_brightness(img, brightness_factor)
1132
            elif fn_id == 1 and contrast_factor is not None:
1133
                img = F.adjust_contrast(img, contrast_factor)
1134
            elif fn_id == 2 and saturation_factor is not None:
1135
                img = F.adjust_saturation(img, saturation_factor)
1136
            elif fn_id == 3 and hue_factor is not None:
1137
1138
1139
                img = F.adjust_hue(img, hue_factor)

        return img
1140

1141
    def __repr__(self):
1142
1143
1144
1145
1146
1147
        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
1148

1149

1150
class RandomRotation(torch.nn.Module):
1151
    """Rotate the image by angle.
1152
1153
    The image can be a PIL Image or a Tensor, in which case it is expected
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
1154
1155
1156
1157
1158

    Args:
        degrees (sequence or float or int): Range of degrees to select from.
            If degrees is a number instead of sequence like (min, max), the range of degrees
            will be (-degrees, +degrees).
1159
1160
1161
        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.
1162
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
1163
1164
1165
1166
        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.
1167
        center (list or tuple, optional): Optional center of rotation, (x, y). Origin is the upper left corner.
1168
            Default is the center of the image.
1169
        fill (sequence or int or float, optional): Pixel fill value for the area outside the rotated
Philip Meier's avatar
Philip Meier committed
1170
            image. If int or float, the value is used for all bands respectively.
1171
1172
            This option is supported for PIL image and Tensor inputs.
            If input is PIL Image, the options is only available for ``Pillow>=5.2.0``.
1173
1174
        resample (int, optional): deprecated argument and will be removed since v0.10.0.
            Please use `arg`:interpolation: instead.
1175
1176
1177

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

1178
1179
    """

1180
    def __init__(
1181
        self, degrees, interpolation=InterpolationMode.NEAREST, expand=False, center=None, fill=None, resample=None
1182
    ):
1183
        super().__init__()
1184
1185
1186
1187
1188
1189
1190
1191
1192
        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(
1193
1194
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
1195
1196
1197
            )
            interpolation = _interpolation_modes_from_int(interpolation)

1198
        self.degrees = _setup_angle(degrees, name="degrees", req_sizes=(2, ))
1199
1200

        if center is not None:
1201
            _check_sequence_input(center, "center", req_sizes=(2, ))
1202
1203

        self.center = center
1204

1205
        self.resample = self.interpolation = interpolation
1206
        self.expand = expand
1207
        self.fill = fill
1208
1209

    @staticmethod
1210
    def get_params(degrees: List[float]) -> float:
1211
1212
1213
        """Get parameters for ``rotate`` for a random rotation.

        Returns:
1214
            float: angle parameter to be passed to ``rotate`` for random rotation.
1215
        """
1216
        angle = float(torch.empty(1).uniform_(float(degrees[0]), float(degrees[1])).item())
1217
1218
        return angle

1219
    def forward(self, img):
1220
        """
1221
        Args:
1222
            img (PIL Image or Tensor): Image to be rotated.
1223
1224

        Returns:
1225
            PIL Image or Tensor: Rotated image.
1226
        """
1227
1228
1229
1230
1231
1232
        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]
1233
        angle = self.get_params(self.degrees)
1234
1235

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

1237
    def __repr__(self):
1238
        interpolate_str = self.interpolation.value
1239
        format_string = self.__class__.__name__ + '(degrees={0}'.format(self.degrees)
1240
        format_string += ', interpolation={0}'.format(interpolate_str)
1241
1242
1243
        format_string += ', expand={0}'.format(self.expand)
        if self.center is not None:
            format_string += ', center={0}'.format(self.center)
1244
1245
        if self.fill is not None:
            format_string += ', fill={0}'.format(self.fill)
1246
1247
        format_string += ')'
        return format_string
1248

1249

1250
1251
1252
1253
class RandomAffine(torch.nn.Module):
    """Random affine transformation of the image keeping center invariant.
    The image can be a PIL Image or a Tensor, in which case it is expected
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
1254
1255
1256
1257

    Args:
        degrees (sequence or float or int): Range of degrees to select from.
            If degrees is a number instead of sequence like (min, max), the range of degrees
1258
            will be (-degrees, +degrees). Set to 0 to deactivate rotations.
1259
1260
1261
1262
1263
1264
1265
        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.
        shear (sequence or float or int, optional): Range of degrees to select from.
ptrblck's avatar
ptrblck committed
1266
            If shear is a number, a shear parallel to the x axis in the range (-shear, +shear)
1267
            will be applied. Else if shear is a tuple or list of 2 values a shear parallel to the x axis in the
ptrblck's avatar
ptrblck committed
1268
1269
            range (shear[0], shear[1]) will be applied. Else if shear is a tuple or list of 4 values,
            a x-axis shear in (shear[0], shear[1]) and y-axis shear in (shear[2], shear[3]) will be applied.
1270
            Will not apply shear by default.
1271
1272
1273
        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.
1274
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
1275
1276
1277
1278
1279
        fill (sequence or int or float, optional): Pixel fill value for the area outside the transformed
            image. If int or float, the value is used for all bands respectively.
            This option is supported for PIL image and Tensor inputs.
            If input is PIL Image, the options is only available for ``Pillow>=5.0.0``.
        fillcolor (sequence or int or float, optional): deprecated argument and will be removed since v0.10.0.
1280
1281
1282
            Please use `arg`:fill: instead.
        resample (int, optional): deprecated argument and will be removed since v0.10.0.
            Please use `arg`:interpolation: instead.
1283
1284
1285

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

1286
1287
    """

1288
    def __init__(
1289
        self, degrees, translate=None, scale=None, shear=None, interpolation=InterpolationMode.NEAREST, fill=0,
1290
1291
        fillcolor=None, resample=None
    ):
1292
        super().__init__()
1293
1294
1295
1296
1297
1298
1299
1300
1301
        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(
1302
1303
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
1304
1305
1306
1307
1308
1309
1310
1311
1312
            )
            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

1313
        self.degrees = _setup_angle(degrees, name="degrees", req_sizes=(2, ))
1314
1315

        if translate is not None:
1316
            _check_sequence_input(translate, "translate", req_sizes=(2, ))
1317
1318
1319
1320
1321
1322
            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:
1323
            _check_sequence_input(scale, "scale", req_sizes=(2, ))
1324
1325
1326
1327
1328
1329
            for s in scale:
                if s <= 0:
                    raise ValueError("scale values should be positive")
        self.scale = scale

        if shear is not None:
1330
            self.shear = _setup_angle(shear, name="shear", req_sizes=(2, 4))
1331
1332
1333
        else:
            self.shear = shear

1334
1335
        self.resample = self.interpolation = interpolation
        self.fillcolor = self.fill = fill
1336
1337

    @staticmethod
1338
1339
1340
1341
1342
1343
1344
    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]]:
1345
1346
1347
        """Get parameters for affine transformation

        Returns:
1348
            params to be passed to the affine transformation
1349
        """
1350
        angle = float(torch.empty(1).uniform_(float(degrees[0]), float(degrees[1])).item())
1351
        if translate is not None:
1352
1353
1354
1355
1356
            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)
1357
1358
1359
1360
        else:
            translations = (0, 0)

        if scale_ranges is not None:
1361
            scale = float(torch.empty(1).uniform_(scale_ranges[0], scale_ranges[1]).item())
1362
1363
1364
        else:
            scale = 1.0

1365
        shear_x = shear_y = 0.0
1366
        if shears is not None:
1367
1368
1369
1370
1371
            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)
1372
1373
1374

        return angle, translations, scale, shear

1375
    def forward(self, img):
1376
        """
1377
            img (PIL Image or Tensor): Image to be transformed.
1378
1379

        Returns:
1380
            PIL Image or Tensor: Affine transformed image.
1381
        """
1382
1383
1384
1385
1386
1387
        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]
1388
1389
1390
1391

        img_size = F._get_image_size(img)

        ret = self.get_params(self.degrees, self.translate, self.scale, self.shear, img_size)
1392
1393

        return F.affine(img, *ret, interpolation=self.interpolation, fill=fill)
1394
1395
1396
1397
1398
1399
1400
1401
1402

    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}'
1403
        if self.interpolation != InterpolationMode.NEAREST:
1404
1405
1406
            s += ', interpolation={interpolation}'
        if self.fill != 0:
            s += ', fill={fill}'
1407
1408
        s += ')'
        d = dict(self.__dict__)
1409
        d['interpolation'] = self.interpolation.value
1410
1411
1412
        return s.format(name=self.__class__.__name__, **d)


1413
class Grayscale(torch.nn.Module):
1414
    """Convert image to grayscale.
1415
1416
1417
    The image can be a PIL Image or a Tensor, in which case it is expected
    to have [..., 3, H, W] shape, where ... means an arbitrary number of leading
    dimensions
1418

1419
1420
1421
1422
    Args:
        num_output_channels (int): (1 or 3) number of channels desired for output image

    Returns:
1423
        PIL Image: Grayscale version of the input.
1424
1425
         - If ``num_output_channels == 1`` : returned image is single channel
         - If ``num_output_channels == 3`` : returned image is 3 channel with r == g == b
1426
1427
1428
1429

    """

    def __init__(self, num_output_channels=1):
1430
        super().__init__()
1431
1432
        self.num_output_channels = num_output_channels

vfdev's avatar
vfdev committed
1433
    def forward(self, img):
1434
1435
        """
        Args:
1436
            img (PIL Image or Tensor): Image to be converted to grayscale.
1437
1438

        Returns:
1439
            PIL Image or Tensor: Grayscaled image.
1440
        """
1441
        return F.rgb_to_grayscale(img, num_output_channels=self.num_output_channels)
1442

1443
    def __repr__(self):
1444
        return self.__class__.__name__ + '(num_output_channels={0})'.format(self.num_output_channels)
1445

1446

1447
class RandomGrayscale(torch.nn.Module):
1448
    """Randomly convert image to grayscale with a probability of p (default 0.1).
1449
1450
1451
    The image can be a PIL Image or a Tensor, in which case it is expected
    to have [..., 3, H, W] shape, where ... means an arbitrary number of leading
    dimensions
1452

1453
1454
1455
1456
    Args:
        p (float): probability that image should be converted to grayscale.

    Returns:
1457
        PIL Image or Tensor: Grayscale version of the input image with probability p and unchanged
1458
1459
1460
        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
1461
1462
1463
1464

    """

    def __init__(self, p=0.1):
1465
        super().__init__()
1466
1467
        self.p = p

vfdev's avatar
vfdev committed
1468
    def forward(self, img):
1469
1470
        """
        Args:
1471
            img (PIL Image or Tensor): Image to be converted to grayscale.
1472
1473

        Returns:
1474
            PIL Image or Tensor: Randomly grayscaled image.
1475
        """
1476
1477
1478
        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)
1479
        return img
1480
1481

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


1485
class RandomErasing(torch.nn.Module):
1486
    """ Randomly selects a rectangle region in an image and erases its pixels.
vfdev's avatar
vfdev committed
1487
    'Random Erasing Data Augmentation' by Zhong et al. See https://arxiv.org/abs/1708.04896
1488

1489
1490
1491
1492
1493
1494
1495
1496
    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
1497
         inplace: boolean to make this transform inplace. Default set to False.
1498

1499
1500
    Returns:
        Erased Image.
1501

vfdev's avatar
vfdev committed
1502
    Example:
1503
        >>> transform = transforms.Compose([
1504
1505
1506
1507
        >>>   transforms.RandomHorizontalFlip(),
        >>>   transforms.ToTensor(),
        >>>   transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        >>>   transforms.RandomErasing(),
1508
1509
1510
        >>> ])
    """

Zhun Zhong's avatar
Zhun Zhong committed
1511
    def __init__(self, p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=0, inplace=False):
1512
1513
1514
1515
1516
1517
1518
1519
1520
        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")
1521
        if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
1522
            warnings.warn("Scale and ratio should be of kind (min, max)")
1523
        if scale[0] < 0 or scale[1] > 1:
1524
            raise ValueError("Scale should be between 0 and 1")
1525
        if p < 0 or p > 1:
1526
            raise ValueError("Random erasing probability should be between 0 and 1")
1527
1528
1529
1530
1531

        self.p = p
        self.scale = scale
        self.ratio = ratio
        self.value = value
1532
        self.inplace = inplace
1533
1534

    @staticmethod
1535
1536
1537
    def get_params(
            img: Tensor, scale: Tuple[float, float], ratio: Tuple[float, float], value: Optional[List[float]] = None
    ) -> Tuple[int, int, int, int, Tensor]:
1538
1539
1540
        """Get parameters for ``erase`` for a random erasing.

        Args:
vfdev's avatar
vfdev committed
1541
            img (Tensor): Tensor image to be erased.
1542
1543
1544
1545
1546
            scale (tuple or list): range of proportion of erased area against input image.
            ratio (tuple or list): range of aspect ratio of erased area.
            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]``.
1547
1548
1549
1550

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

1554
        for _ in range(10):
1555
1556
            erase_area = area * torch.empty(1).uniform_(scale[0], scale[1]).item()
            aspect_ratio = torch.empty(1).uniform_(ratio[0], ratio[1]).item()
1557
1558
1559

            h = int(round(math.sqrt(erase_area * aspect_ratio)))
            w = int(round(math.sqrt(erase_area / aspect_ratio)))
1560
1561
1562
1563
1564
1565
1566
            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]
1567

1568
1569
            i = torch.randint(0, img_h - h + 1, size=(1, )).item()
            j = torch.randint(0, img_w - w + 1, size=(1, )).item()
1570
            return i, j, h, w, v
1571

Zhun Zhong's avatar
Zhun Zhong committed
1572
1573
1574
        # Return original image
        return 0, 0, img_h, img_w, img

1575
    def forward(self, img):
1576
1577
        """
        Args:
vfdev's avatar
vfdev committed
1578
            img (Tensor): Tensor image to be erased.
1579
1580
1581
1582

        Returns:
            img (Tensor): Erased Tensor image.
        """
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
        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)
1602
            return F.erase(img, x, y, h, w, v, self.inplace)
1603
        return img
1604
1605


1606
1607
1608
class GaussianBlur(torch.nn.Module):
    """Blurs image with randomly chosen Gaussian blur.
    The image can be a PIL Image or a Tensor, in which case it is expected
vfdev's avatar
vfdev committed
1609
    to have [..., C, H, W] shape, where ... means an arbitrary number of leading
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
    dimensions

    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
1645
        """Choose sigma for random gaussian blurring.
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658

        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
1659
            img (PIL Image or Tensor): image to be blurred.
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672

        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


1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
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]
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
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
1771
1772
1773
1774
1775
1776
1777
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
1805
1806
1807
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
1837
1838
1839
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
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889


class RandomInvert(torch.nn.Module):
    """Inverts the colors of the given image randomly with a given probability.
    The image can be a PIL Image or a torch Tensor, in which case it is expected
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions.

    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
    number of bits for each color channel. The image can be a PIL Image or a torch
    Tensor, in which case it is expected to have [..., H, W] shape, where ... means
    an arbitrary number of leading dimensions.

    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
    values above a threshold. The image can be a PIL Image or a torch Tensor, in
    which case it is expected to have [..., H, W] shape, where ... means an arbitrary
    number of leading dimensions.

    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):
    """Adjust the sharpness of the image randomly with a given probability. The image
    can be a PIL Image or a torch Tensor, in which case it is expected to have [..., H, W]
    shape, where ... means an arbitrary number of leading dimensions.

    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.
    The image can be a PIL Image or a torch Tensor, in which case it is expected
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions.

    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.
    The image can be a PIL Image or a torch Tensor, in which case it is expected
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions.

    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)