transforms.py 76.5 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

    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
237
            (size * height / width, size).
238
            In torchscript mode size as single int is not supported, use a sequence of length 1: ``[size, ]``.
239
240
241
242
        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.
243
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
244
245
246
247
248
249
250
251
        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
            ``max_size``. As a result, ```size` might be overruled, i.e the
            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).
252

253
254
    """

255
    def __init__(self, size, interpolation=InterpolationMode.BILINEAR, max_size=None):
vfdev's avatar
vfdev committed
256
        super().__init__()
257
258
259
260
261
        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
262
        self.max_size = max_size
263
264
265
266

        # Backward compatibility with integer value
        if isinstance(interpolation, int):
            warnings.warn(
267
268
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
269
270
271
            )
            interpolation = _interpolation_modes_from_int(interpolation)

272
273
        self.interpolation = interpolation

vfdev's avatar
vfdev committed
274
    def forward(self, img):
275
276
        """
        Args:
vfdev's avatar
vfdev committed
277
            img (PIL Image or Tensor): Image to be scaled.
278
279

        Returns:
vfdev's avatar
vfdev committed
280
            PIL Image or Tensor: Rescaled image.
281
        """
282
        return F.resize(img, self.size, self.interpolation, self.max_size)
283

284
    def __repr__(self):
285
        interpolate_str = self.interpolation.value
286
287
        return self.__class__.__name__ + '(size={0}, interpolation={1}, max_size={2})'.format(
            self.size, interpolate_str, self.max_size)
288

289
290
291
292
293
294
295
296
297
298
299

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
300
301
class CenterCrop(torch.nn.Module):
    """Crops the given image at the center.
302
    If the image is torch Tensor, it is expected
303
304
    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.
305
306
307
308

    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
309
            made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
310
311
312
    """

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

vfdev's avatar
vfdev committed
316
    def forward(self, img):
317
318
        """
        Args:
vfdev's avatar
vfdev committed
319
            img (PIL Image or Tensor): Image to be cropped.
320
321

        Returns:
vfdev's avatar
vfdev committed
322
            PIL Image or Tensor: Cropped image.
323
324
325
        """
        return F.center_crop(img, self.size)

326
327
328
    def __repr__(self):
        return self.__class__.__name__ + '(size={0})'.format(self.size)

329

330
331
class Pad(torch.nn.Module):
    """Pad the given image on all sides with the given "pad" value.
332
    If the image is torch Tensor, it is expected
333
334
335
    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
336
337

    Args:
338
339
340
        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
341
            this is the padding for the left, top, right and bottom borders respectively.
342
343
            In torchscript mode padding as single int is not supported, use a sequence of length 1: ``[padding, ]``.
        fill (number or str or tuple): Pixel fill value for constant fill. Default is 0. If a tuple of
344
            length 3, it is used to fill R, G, B channels respectively.
345
346
347
            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.
348
        padding_mode (str): Type of padding. Should be: constant, edge, reflect or symmetric.
349
            Default is constant.
350
351
352

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

353
354
            - 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
355
356
357
358

            - 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
359
                will result in [3, 2, 1, 2, 3, 4, 3, 2]
360
361
362
363

            - 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
364
                will result in [2, 1, 1, 2, 3, 4, 4, 3]
365
366
    """

367
368
369
370
371
372
373
374
375
376
377
378
379
    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 " +
380
381
382
383
                             "{} element tuple".format(len(padding)))

        self.padding = padding
        self.fill = fill
384
        self.padding_mode = padding_mode
385

386
    def forward(self, img):
387
388
        """
        Args:
389
            img (PIL Image or Tensor): Image to be padded.
390
391

        Returns:
392
            PIL Image or Tensor: Padded image.
393
        """
394
        return F.pad(img, self.padding, self.fill, self.padding_mode)
395

396
    def __repr__(self):
397
398
        return self.__class__.__name__ + '(padding={0}, fill={1}, padding_mode={2})'.\
            format(self.padding, self.fill, self.padding_mode)
399

400

401
class Lambda:
402
    """Apply a user-defined lambda as a transform. This transform does not support torchscript.
403
404
405
406
407
408

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

    def __init__(self, lambd):
409
410
        if not callable(lambd):
            raise TypeError("Argument lambd should be callable, got {}".format(repr(type(lambd).__name__)))
411
412
413
414
415
        self.lambd = lambd

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

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

419

420
class RandomTransforms:
421
422
423
    """Base class for a list of transformations with randomness

    Args:
424
        transforms (sequence): list of transformations
425
426
427
    """

    def __init__(self, transforms):
428
429
        if not isinstance(transforms, Sequence):
            raise TypeError("Argument transforms should be a sequence")
430
431
432
433
434
435
436
437
438
439
440
441
442
443
        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


444
class RandomApply(torch.nn.Module):
445
    """Apply randomly a list of transformations with a given probability.
446
447
448
449
450
451
452
453
454
455
456
457

    .. 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``.
458
459

    Args:
460
        transforms (sequence or torch.nn.Module): list of transformations
461
462
463
464
        p (float): probability
    """

    def __init__(self, transforms, p=0.5):
465
466
        super().__init__()
        self.transforms = transforms
467
468
        self.p = p

469
470
    def forward(self, img):
        if self.p < torch.rand(1):
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
            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):
487
    """Apply a list of transformations in a random order. This transform does not support torchscript.
488
489
490
491
492
493
494
495
496
497
    """
    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):
498
    """Apply single transformation randomly picked from a list. This transform does not support torchscript.
499
500
501
502
503
504
    """
    def __call__(self, img):
        t = random.choice(self.transforms)
        return t(img)


vfdev's avatar
vfdev committed
505
506
class RandomCrop(torch.nn.Module):
    """Crop the given image at a random location.
507
    If the image is torch Tensor, it is expected
508
509
    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
510
511
512
513

    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
514
            made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
515
        padding (int or sequence, optional): Optional padding on each border
vfdev's avatar
vfdev committed
516
            of the image. Default is None. If a single int is provided this
517
518
            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
519
            this is the padding for the left, top, right and bottom borders respectively.
520
            In torchscript mode padding as single int is not supported, use a sequence of length 1: ``[padding, ]``.
521
        pad_if_needed (boolean): It will pad the image if smaller than the
ekka's avatar
ekka committed
522
            desired size to avoid raising an exception. Since cropping is done
523
            after padding, the padding seems to be done at a random offset.
524
        fill (number or str or tuple): Pixel fill value for constant fill. Default is 0. If a tuple of
525
            length 3, it is used to fill R, G, B channels respectively.
526
527
528
            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.
vfdev's avatar
vfdev committed
529
        padding_mode (str): Type of padding. Should be: constant, edge, reflect or symmetric. Default is constant.
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544

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

545
546
547
    """

    @staticmethod
vfdev's avatar
vfdev committed
548
    def get_params(img: Tensor, output_size: Tuple[int, int]) -> Tuple[int, int, int, int]:
549
550
551
        """Get parameters for ``crop`` for a random crop.

        Args:
vfdev's avatar
vfdev committed
552
            img (PIL Image or Tensor): Image to be cropped.
553
554
555
556
557
            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
558
        w, h = F._get_image_size(img)
559
        th, tw = output_size
vfdev's avatar
vfdev committed
560
561
562
563
564
565

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

566
567
568
        if w == tw and h == th:
            return 0, 0, h, w

569
570
        i = torch.randint(0, h - th + 1, size=(1, )).item()
        j = torch.randint(0, w - tw + 1, size=(1, )).item()
571
572
        return i, j, th, tw

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

576
577
578
579
        self.size = tuple(_setup_size(
            size, error_msg="Please provide only two dimensions (h, w) for size."
        ))

vfdev's avatar
vfdev committed
580
581
582
583
584
585
        self.padding = padding
        self.pad_if_needed = pad_if_needed
        self.fill = fill
        self.padding_mode = padding_mode

    def forward(self, img):
586
587
        """
        Args:
vfdev's avatar
vfdev committed
588
            img (PIL Image or Tensor): Image to be cropped.
589
590

        Returns:
vfdev's avatar
vfdev committed
591
            PIL Image or Tensor: Cropped image.
592
        """
593
594
        if self.padding is not None:
            img = F.pad(img, self.padding, self.fill, self.padding_mode)
595

vfdev's avatar
vfdev committed
596
        width, height = F._get_image_size(img)
597
        # pad the width if needed
vfdev's avatar
vfdev committed
598
599
600
        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)
601
        # pad the height if needed
vfdev's avatar
vfdev committed
602
603
604
        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)
605

606
607
608
609
        i, j, h, w = self.get_params(img, self.size)

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

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

613

614
615
class RandomHorizontalFlip(torch.nn.Module):
    """Horizontally flip the given image randomly with a given probability.
616
    If the image is torch Tensor, it is expected
617
618
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
619
620
621
622
623
624

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

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

628
    def forward(self, img):
629
630
        """
        Args:
631
            img (PIL Image or Tensor): Image to be flipped.
632
633

        Returns:
634
            PIL Image or Tensor: Randomly flipped image.
635
        """
636
        if torch.rand(1) < self.p:
637
638
639
            return F.hflip(img)
        return img

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

643

644
class RandomVerticalFlip(torch.nn.Module):
vfdev's avatar
vfdev committed
645
    """Vertically flip the given image randomly with a given probability.
646
    If the image is torch Tensor, it is expected
647
648
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
649
650
651
652
653
654

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

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

658
    def forward(self, img):
659
660
        """
        Args:
661
            img (PIL Image or Tensor): Image to be flipped.
662
663

        Returns:
664
            PIL Image or Tensor: Randomly flipped image.
665
        """
666
        if torch.rand(1) < self.p:
667
668
669
            return F.vflip(img)
        return img

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

673

674
675
class RandomPerspective(torch.nn.Module):
    """Performs a random perspective transformation of the given image with a given probability.
676
    If the image is torch Tensor, it is expected
677
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
678
679

    Args:
680
681
682
        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.
683
684
685
        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.
686
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
687
688
        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.
689
            If input is PIL Image, the options is only available for ``Pillow>=5.0.0``.
690
691
    """

692
    def __init__(self, distortion_scale=0.5, p=0.5, interpolation=InterpolationMode.BILINEAR, fill=0):
693
        super().__init__()
694
        self.p = p
695
696
697
698

        # Backward compatibility with integer value
        if isinstance(interpolation, int):
            warnings.warn(
699
700
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
701
702
703
            )
            interpolation = _interpolation_modes_from_int(interpolation)

704
705
        self.interpolation = interpolation
        self.distortion_scale = distortion_scale
706
707
708
709
710
711

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

712
        self.fill = fill
713

714
    def forward(self, img):
715
716
        """
        Args:
717
            img (PIL Image or Tensor): Image to be Perspectively transformed.
718
719

        Returns:
720
            PIL Image or Tensor: Randomly transformed image.
721
        """
722
723
724
725
726
727
728
729

        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]

730
731
        if torch.rand(1) < self.p:
            width, height = F._get_image_size(img)
732
            startpoints, endpoints = self.get_params(width, height, self.distortion_scale)
733
            return F.perspective(img, startpoints, endpoints, self.interpolation, fill)
734
735
736
        return img

    @staticmethod
737
    def get_params(width: int, height: int, distortion_scale: float) -> Tuple[List[List[int]], List[List[int]]]:
738
739
740
        """Get parameters for ``perspective`` for a random perspective transform.

        Args:
741
742
743
            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.
744
745

        Returns:
746
            List containing [top-left, top-right, bottom-right, bottom-left] of the original image,
747
748
            List containing [top-left, top-right, bottom-right, bottom-left] of the transformed image.
        """
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
        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]]
768
769
770
771
772
773
774
        endpoints = [topleft, topright, botright, botleft]
        return startpoints, endpoints

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


775
class RandomResizedCrop(torch.nn.Module):
776
777
    """Crop a random portion of image and resize it to a given size.

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

781
782
783
    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.
784
785

    Args:
786
        size (int or sequence): expected output size of the crop, for each edge. If size is an
787
            int instead of sequence like (h, w), a square output size ``(size, size)`` is
788
789
            made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
            In torchscript mode size as single int is not supported, use a sequence of length 1: ``[size, ]``.
790
791
792
        scale (tuple of float): lower and upper bounds for the random area of the crop, before resizing.
        ratio (tuple of float): lower and upper bounds for the random aspect ratio of the crop, before
            resizing.
793
794
795
796
        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.
797
798
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.

799
800
    """

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

805
        if not isinstance(scale, Sequence):
806
            raise TypeError("Scale should be a sequence")
807
        if not isinstance(ratio, Sequence):
808
            raise TypeError("Ratio should be a sequence")
809
        if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
810
            warnings.warn("Scale and ratio should be of kind (min, max)")
811

812
813
814
        # Backward compatibility with integer value
        if isinstance(interpolation, int):
            warnings.warn(
815
816
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
817
818
819
            )
            interpolation = _interpolation_modes_from_int(interpolation)

820
        self.interpolation = interpolation
821
822
        self.scale = scale
        self.ratio = ratio
823
824

    @staticmethod
825
    def get_params(
826
            img: Tensor, scale: List[float], ratio: List[float]
827
    ) -> Tuple[int, int, int, int]:
828
829
830
        """Get parameters for ``crop`` for a random sized crop.

        Args:
831
            img (PIL Image or Tensor): Input image.
832
833
            scale (list): range of scale of the origin size cropped
            ratio (list): range of aspect ratio of the origin aspect ratio cropped
834
835
836
837
838

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

842
        log_ratio = torch.log(torch.tensor(ratio))
843
        for _ in range(10):
844
            target_area = area * torch.empty(1).uniform_(scale[0], scale[1]).item()
845
846
847
            aspect_ratio = torch.exp(
                torch.empty(1).uniform_(log_ratio[0], log_ratio[1])
            ).item()
848
849
850
851

            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
852
            if 0 < w <= width and 0 < h <= height:
853
854
                i = torch.randint(0, height - h + 1, size=(1,)).item()
                j = torch.randint(0, width - w + 1, size=(1,)).item()
855
856
                return i, j, h, w

857
        # Fallback to central crop
Zhicheng Yan's avatar
Zhicheng Yan committed
858
        in_ratio = float(width) / float(height)
859
        if in_ratio < min(ratio):
Zhicheng Yan's avatar
Zhicheng Yan committed
860
            w = width
861
            h = int(round(w / min(ratio)))
862
        elif in_ratio > max(ratio):
Zhicheng Yan's avatar
Zhicheng Yan committed
863
            h = height
864
            w = int(round(h * max(ratio)))
865
        else:  # whole image
Zhicheng Yan's avatar
Zhicheng Yan committed
866
867
868
869
            w = width
            h = height
        i = (height - h) // 2
        j = (width - w) // 2
870
        return i, j, h, w
871

872
    def forward(self, img):
873
874
        """
        Args:
875
            img (PIL Image or Tensor): Image to be cropped and resized.
876
877

        Returns:
878
            PIL Image or Tensor: Randomly cropped and resized image.
879
        """
880
        i, j, h, w = self.get_params(img, self.scale, self.ratio)
881
882
        return F.resized_crop(img, i, j, h, w, self.size, self.interpolation)

883
    def __repr__(self):
884
        interpolate_str = self.interpolation.value
885
        format_string = self.__class__.__name__ + '(size={0}'.format(self.size)
886
887
        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))
888
889
        format_string += ', interpolation={0})'.format(interpolate_str)
        return format_string
890

891
892
893
894
895
896
897
898
899
900
901

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
902
903
class FiveCrop(torch.nn.Module):
    """Crop the given image into four corners and the central crop.
904
    If the image is torch Tensor, it is expected
vfdev's avatar
vfdev committed
905
906
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
907
908
909
910
911
912
913
914
915

    .. 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.
916
            If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
917
918
919
920
921
922
923
924
925
926
927
928
929
930

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

vfdev's avatar
vfdev committed
934
935
936
937
938
939
940
941
    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
        """
942
943
        return F.five_crop(img, self.size)

944
945
946
    def __repr__(self):
        return self.__class__.__name__ + '(size={0})'.format(self.size)

947

vfdev's avatar
vfdev committed
948
949
950
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).
951
    If the image is torch Tensor, it is expected
vfdev's avatar
vfdev committed
952
953
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
954
955
956
957
958
959
960
961
962

    .. 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
963
            made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
964
        vertical_flip (bool): Use vertical flipping instead of horizontal
965
966
967
968
969
970
971
972
973
974
975
976
977
978

    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
979
        super().__init__()
980
        self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
981
982
        self.vertical_flip = vertical_flip

vfdev's avatar
vfdev committed
983
984
985
986
987
988
989
990
    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
        """
991
992
        return F.ten_crop(img, self.size, self.vertical_flip)

993
    def __repr__(self):
994
        return self.__class__.__name__ + '(size={0}, vertical_flip={1})'.format(self.size, self.vertical_flip)
995

996

997
class LinearTransformation(torch.nn.Module):
ekka's avatar
ekka committed
998
    """Transform a tensor image with a square transformation matrix and a mean_vector computed
999
    offline.
1000
    This transform does not support PIL Image.
ekka's avatar
ekka committed
1001
1002
1003
    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
1004
    original shape.
1005

1006
    Applications:
1007
        whitening transformation: Suppose X is a column vector zero-centered data.
1008
1009
1010
        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.

1011
1012
    Args:
        transformation_matrix (Tensor): tensor [D x D], D = C x H x W
ekka's avatar
ekka committed
1013
        mean_vector (Tensor): tensor [D], D = C x H x W
1014
1015
    """

ekka's avatar
ekka committed
1016
    def __init__(self, transformation_matrix, mean_vector):
1017
        super().__init__()
1018
1019
1020
        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
1021
1022
1023

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

1027
1028
1029
1030
        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))

1031
        self.transformation_matrix = transformation_matrix
ekka's avatar
ekka committed
1032
        self.mean_vector = mean_vector
1033

1034
    def forward(self, tensor: Tensor) -> Tensor:
1035
1036
        """
        Args:
vfdev's avatar
vfdev committed
1037
            tensor (Tensor): Tensor image to be whitened.
1038
1039
1040
1041

        Returns:
            Tensor: Transformed image.
        """
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
        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
1054
        transformed_tensor = torch.mm(flat_tensor, self.transformation_matrix)
1055
        tensor = transformed_tensor.view(shape)
1056
1057
        return tensor

1058
    def __repr__(self):
ekka's avatar
ekka committed
1059
1060
1061
        format_string = self.__class__.__name__ + '(transformation_matrix='
        format_string += (str(self.transformation_matrix.tolist()) + ')')
        format_string += (", (mean_vector=" + str(self.mean_vector.tolist()) + ')')
1062
1063
        return format_string

1064

1065
class ColorJitter(torch.nn.Module):
1066
    """Randomly change the brightness, contrast, saturation and hue of an image.
1067
    If the image is torch Tensor, it is expected
1068
1069
    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.
1070
1071

    Args:
yaox12's avatar
yaox12 committed
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
        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.
1084
    """
1085

1086
    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):
1087
        super().__init__()
yaox12's avatar
yaox12 committed
1088
1089
1090
1091
1092
1093
        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)

1094
    @torch.jit.unused
yaox12's avatar
yaox12 committed
1095
1096
1097
1098
    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))
1099
            value = [center - float(value), center + float(value)]
yaox12's avatar
yaox12 committed
1100
            if clip_first_on_zero:
1101
                value[0] = max(value[0], 0.0)
yaox12's avatar
yaox12 committed
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
        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
1113
1114

    @staticmethod
1115
1116
1117
1118
1119
1120
    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.
1121

1122
1123
1124
1125
1126
1127
1128
1129
1130
        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.
1131
1132

        Returns:
1133
1134
            tuple: The parameters used to apply the randomized transform
            along with their random order.
1135
        """
1136
        fn_idx = torch.randperm(4)
1137

1138
1139
1140
1141
        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]))
1142

1143
        return fn_idx, b, c, s, h
1144

1145
    def forward(self, img):
1146
1147
        """
        Args:
1148
            img (PIL Image or Tensor): Input image.
1149
1150

        Returns:
1151
1152
            PIL Image or Tensor: Color jittered image.
        """
1153
1154
1155
        fn_idx, brightness_factor, contrast_factor, saturation_factor, hue_factor = \
            self.get_params(self.brightness, self.contrast, self.saturation, self.hue)

1156
        for fn_id in fn_idx:
1157
            if fn_id == 0 and brightness_factor is not None:
1158
                img = F.adjust_brightness(img, brightness_factor)
1159
            elif fn_id == 1 and contrast_factor is not None:
1160
                img = F.adjust_contrast(img, contrast_factor)
1161
            elif fn_id == 2 and saturation_factor is not None:
1162
                img = F.adjust_saturation(img, saturation_factor)
1163
            elif fn_id == 3 and hue_factor is not None:
1164
1165
1166
                img = F.adjust_hue(img, hue_factor)

        return img
1167

1168
    def __repr__(self):
1169
1170
1171
1172
1173
1174
        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
1175

1176

1177
class RandomRotation(torch.nn.Module):
1178
    """Rotate the image by angle.
1179
    If the image is torch Tensor, it is expected
1180
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
1181
1182

    Args:
1183
        degrees (sequence or number): Range of degrees to select from.
1184
1185
            If degrees is a number instead of sequence like (min, max), the range of degrees
            will be (-degrees, +degrees).
1186
1187
1188
        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.
1189
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
1190
1191
1192
1193
        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.
1194
        center (sequence, optional): Optional center of rotation, (x, y). Origin is the upper left corner.
1195
            Default is the center of the image.
1196
1197
        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.
1198
            If input is PIL Image, the options is only available for ``Pillow>=5.2.0``.
1199
        resample (int, optional): deprecated argument and will be removed since v0.10.0.
1200
            Please use the ``interpolation`` parameter instead.
1201
1202
1203

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

1204
1205
    """

1206
    def __init__(
1207
        self, degrees, interpolation=InterpolationMode.NEAREST, expand=False, center=None, fill=0, resample=None
1208
    ):
1209
        super().__init__()
1210
1211
1212
1213
1214
1215
1216
1217
1218
        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(
1219
1220
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
1221
1222
1223
            )
            interpolation = _interpolation_modes_from_int(interpolation)

1224
        self.degrees = _setup_angle(degrees, name="degrees", req_sizes=(2, ))
1225
1226

        if center is not None:
1227
            _check_sequence_input(center, "center", req_sizes=(2, ))
1228
1229

        self.center = center
1230

1231
        self.resample = self.interpolation = interpolation
1232
        self.expand = expand
1233
1234
1235
1236
1237
1238

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

1239
        self.fill = fill
1240
1241

    @staticmethod
1242
    def get_params(degrees: List[float]) -> float:
1243
1244
1245
        """Get parameters for ``rotate`` for a random rotation.

        Returns:
1246
            float: angle parameter to be passed to ``rotate`` for random rotation.
1247
        """
1248
        angle = float(torch.empty(1).uniform_(float(degrees[0]), float(degrees[1])).item())
1249
1250
        return angle

1251
    def forward(self, img):
1252
        """
1253
        Args:
1254
            img (PIL Image or Tensor): Image to be rotated.
1255
1256

        Returns:
1257
            PIL Image or Tensor: Rotated image.
1258
        """
1259
1260
1261
1262
1263
1264
        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]
1265
        angle = self.get_params(self.degrees)
1266
1267

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

1269
    def __repr__(self):
1270
        interpolate_str = self.interpolation.value
1271
        format_string = self.__class__.__name__ + '(degrees={0}'.format(self.degrees)
1272
        format_string += ', interpolation={0}'.format(interpolate_str)
1273
1274
1275
        format_string += ', expand={0}'.format(self.expand)
        if self.center is not None:
            format_string += ', center={0}'.format(self.center)
1276
1277
        if self.fill is not None:
            format_string += ', fill={0}'.format(self.fill)
1278
1279
        format_string += ')'
        return format_string
1280

1281

1282
1283
class RandomAffine(torch.nn.Module):
    """Random affine transformation of the image keeping center invariant.
1284
    If the image is torch Tensor, it is expected
1285
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
1286
1287

    Args:
1288
        degrees (sequence or number): Range of degrees to select from.
1289
            If degrees is a number instead of sequence like (min, max), the range of degrees
1290
            will be (-degrees, +degrees). Set to 0 to deactivate rotations.
1291
1292
1293
1294
1295
1296
        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.
1297
        shear (sequence or number, optional): Range of degrees to select from.
ptrblck's avatar
ptrblck committed
1298
            If shear is a number, a shear parallel to the x axis in the range (-shear, +shear)
1299
1300
            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
1301
            a x-axis shear in (shear[0], shear[1]) and y-axis shear in (shear[2], shear[3]) will be applied.
1302
            Will not apply shear by default.
1303
1304
1305
        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.
1306
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
1307
1308
        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.
1309
            If input is PIL Image, the options is only available for ``Pillow>=5.0.0``.
1310
        fillcolor (sequence or number, optional): deprecated argument and will be removed since v0.10.0.
1311
            Please use the ``fill`` parameter instead.
1312
        resample (int, optional): deprecated argument and will be removed since v0.10.0.
1313
            Please use the ``interpolation`` parameter instead.
1314
1315
1316

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

1317
1318
    """

1319
    def __init__(
1320
        self, degrees, translate=None, scale=None, shear=None, interpolation=InterpolationMode.NEAREST, fill=0,
1321
1322
        fillcolor=None, resample=None
    ):
1323
        super().__init__()
1324
1325
1326
1327
1328
1329
1330
1331
1332
        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(
1333
1334
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
1335
1336
1337
1338
1339
1340
1341
1342
1343
            )
            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

1344
        self.degrees = _setup_angle(degrees, name="degrees", req_sizes=(2, ))
1345
1346

        if translate is not None:
1347
            _check_sequence_input(translate, "translate", req_sizes=(2, ))
1348
1349
1350
1351
1352
1353
            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:
1354
            _check_sequence_input(scale, "scale", req_sizes=(2, ))
1355
1356
1357
1358
1359
1360
            for s in scale:
                if s <= 0:
                    raise ValueError("scale values should be positive")
        self.scale = scale

        if shear is not None:
1361
            self.shear = _setup_angle(shear, name="shear", req_sizes=(2, 4))
1362
1363
1364
        else:
            self.shear = shear

1365
        self.resample = self.interpolation = interpolation
1366
1367
1368
1369
1370
1371

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

1372
        self.fillcolor = self.fill = fill
1373
1374

    @staticmethod
1375
1376
1377
1378
1379
1380
1381
    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]]:
1382
1383
1384
        """Get parameters for affine transformation

        Returns:
1385
            params to be passed to the affine transformation
1386
        """
1387
        angle = float(torch.empty(1).uniform_(float(degrees[0]), float(degrees[1])).item())
1388
        if translate is not None:
1389
1390
1391
1392
1393
            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)
1394
1395
1396
1397
        else:
            translations = (0, 0)

        if scale_ranges is not None:
1398
            scale = float(torch.empty(1).uniform_(scale_ranges[0], scale_ranges[1]).item())
1399
1400
1401
        else:
            scale = 1.0

1402
        shear_x = shear_y = 0.0
1403
        if shears is not None:
1404
1405
1406
1407
1408
            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)
1409
1410
1411

        return angle, translations, scale, shear

1412
    def forward(self, img):
1413
        """
1414
            img (PIL Image or Tensor): Image to be transformed.
1415
1416

        Returns:
1417
            PIL Image or Tensor: Affine transformed image.
1418
        """
1419
1420
1421
1422
1423
1424
        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]
1425
1426
1427
1428

        img_size = F._get_image_size(img)

        ret = self.get_params(self.degrees, self.translate, self.scale, self.shear, img_size)
1429
1430

        return F.affine(img, *ret, interpolation=self.interpolation, fill=fill)
1431
1432
1433
1434
1435
1436
1437
1438
1439

    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}'
1440
        if self.interpolation != InterpolationMode.NEAREST:
1441
1442
1443
            s += ', interpolation={interpolation}'
        if self.fill != 0:
            s += ', fill={fill}'
1444
1445
        s += ')'
        d = dict(self.__dict__)
1446
        d['interpolation'] = self.interpolation.value
1447
1448
1449
        return s.format(name=self.__class__.__name__, **d)


1450
class Grayscale(torch.nn.Module):
1451
    """Convert image to grayscale.
1452
1453
    If the image is torch Tensor, it is expected
    to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions
1454

1455
1456
1457
1458
    Args:
        num_output_channels (int): (1 or 3) number of channels desired for output image

    Returns:
1459
        PIL Image: Grayscale version of the input.
1460
1461
         - If ``num_output_channels == 1`` : returned image is single channel
         - If ``num_output_channels == 3`` : returned image is 3 channel with r == g == b
1462
1463
1464
1465

    """

    def __init__(self, num_output_channels=1):
1466
        super().__init__()
1467
1468
        self.num_output_channels = num_output_channels

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

        Returns:
1475
            PIL Image or Tensor: Grayscaled image.
1476
        """
1477
        return F.rgb_to_grayscale(img, num_output_channels=self.num_output_channels)
1478

1479
    def __repr__(self):
1480
        return self.__class__.__name__ + '(num_output_channels={0})'.format(self.num_output_channels)
1481

1482

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

1488
1489
1490
1491
    Args:
        p (float): probability that image should be converted to grayscale.

    Returns:
1492
        PIL Image or Tensor: Grayscale version of the input image with probability p and unchanged
1493
1494
1495
        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
1496
1497
1498
1499

    """

    def __init__(self, p=0.1):
1500
        super().__init__()
1501
1502
        self.p = p

vfdev's avatar
vfdev committed
1503
    def forward(self, img):
1504
1505
        """
        Args:
1506
            img (PIL Image or Tensor): Image to be converted to grayscale.
1507
1508

        Returns:
1509
            PIL Image or Tensor: Randomly grayscaled image.
1510
        """
1511
1512
1513
        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)
1514
        return img
1515
1516

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


1520
class RandomErasing(torch.nn.Module):
1521
1522
    """ 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
1523
    'Random Erasing Data Augmentation' by Zhong et al. See https://arxiv.org/abs/1708.04896
1524

1525
1526
1527
1528
1529
1530
1531
1532
    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
1533
         inplace: boolean to make this transform inplace. Default set to False.
1534

1535
1536
    Returns:
        Erased Image.
1537

vfdev's avatar
vfdev committed
1538
    Example:
1539
        >>> transform = transforms.Compose([
1540
1541
1542
1543
        >>>   transforms.RandomHorizontalFlip(),
        >>>   transforms.ToTensor(),
        >>>   transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        >>>   transforms.RandomErasing(),
1544
1545
1546
        >>> ])
    """

Zhun Zhong's avatar
Zhun Zhong committed
1547
    def __init__(self, p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=0, inplace=False):
1548
1549
1550
1551
1552
1553
1554
1555
1556
        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")
1557
        if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
1558
            warnings.warn("Scale and ratio should be of kind (min, max)")
1559
        if scale[0] < 0 or scale[1] > 1:
1560
            raise ValueError("Scale should be between 0 and 1")
1561
        if p < 0 or p > 1:
1562
            raise ValueError("Random erasing probability should be between 0 and 1")
1563
1564
1565
1566
1567

        self.p = p
        self.scale = scale
        self.ratio = ratio
        self.value = value
1568
        self.inplace = inplace
1569
1570

    @staticmethod
1571
1572
1573
    def get_params(
            img: Tensor, scale: Tuple[float, float], ratio: Tuple[float, float], value: Optional[List[float]] = None
    ) -> Tuple[int, int, int, int, Tensor]:
1574
1575
1576
        """Get parameters for ``erase`` for a random erasing.

        Args:
vfdev's avatar
vfdev committed
1577
            img (Tensor): Tensor image to be erased.
1578
1579
            scale (sequence): range of proportion of erased area against input image.
            ratio (sequence): range of aspect ratio of erased area.
1580
1581
1582
            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]``.
1583
1584
1585
1586

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

1590
        log_ratio = torch.log(torch.tensor(ratio))
1591
        for _ in range(10):
1592
            erase_area = area * torch.empty(1).uniform_(scale[0], scale[1]).item()
1593
1594
1595
            aspect_ratio = torch.exp(
                torch.empty(1).uniform_(log_ratio[0], log_ratio[1])
            ).item()
1596
1597
1598

            h = int(round(math.sqrt(erase_area * aspect_ratio)))
            w = int(round(math.sqrt(erase_area / aspect_ratio)))
1599
1600
1601
1602
1603
1604
1605
            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]
1606

1607
1608
            i = torch.randint(0, img_h - h + 1, size=(1, )).item()
            j = torch.randint(0, img_w - w + 1, size=(1, )).item()
1609
            return i, j, h, w, v
1610

Zhun Zhong's avatar
Zhun Zhong committed
1611
1612
1613
        # Return original image
        return 0, 0, img_h, img_w, img

1614
    def forward(self, img):
1615
1616
        """
        Args:
vfdev's avatar
vfdev committed
1617
            img (Tensor): Tensor image to be erased.
1618
1619
1620
1621

        Returns:
            img (Tensor): Erased Tensor image.
        """
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
        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)
1641
            return F.erase(img, x, y, h, w, v, self.inplace)
1642
        return img
1643

1644
1645
1646
1647
1648
1649
1650
1651
    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

1652

1653
1654
class GaussianBlur(torch.nn.Module):
    """Blurs image with randomly chosen Gaussian blur.
1655
1656
    If the image is torch Tensor, it is expected
    to have [..., C, H, W] shape, where ... means an arbitrary number of leading dimensions.
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690

    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
1691
        """Choose sigma for random gaussian blurring.
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704

        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
1705
            img (PIL Image or Tensor): image to be blurred.
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718

        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


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
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]
1749
1750
1751
1752


class RandomInvert(torch.nn.Module):
    """Inverts the colors of the given image randomly with a given probability.
1753
1754
1755
    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".
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

    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
1783
1784
1785
    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".
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

    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
1815
1816
1817
    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".
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

    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):
1846
1847
    """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.
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

    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.
1879
1880
1881
    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".
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908

    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.
1909
1910
1911
    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".
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934

    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)