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

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

    Example:
        >>> transforms.Compose([
        >>>     transforms.CenterCrop(10),
38
39
        >>>     transforms.PILToTensor(),
        >>>     transforms.ConvertImageDtype(torch.float),
40
        >>> ])
41
42
43
44
45
46
47
48
49
50
51
52
53

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

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

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

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

64
65
66
67
68
69
70
71
    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

72

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

    Converts a PIL Image or numpy.ndarray (H x W x C) in the range
surgan12's avatar
surgan12 committed
77
78
79
80
81
    [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.
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.

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

    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)

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

103

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

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

    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__ + '()'


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

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

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


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

    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
160
161
162
163
164
            - 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``).
165

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

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

189

190
class Normalize(torch.nn.Module):
Fang Gao's avatar
Fang Gao committed
191
    """Normalize a tensor image with mean and standard deviation.
192
    This transform does not support PIL Image.
193
194
195
    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
196
    ``output[channel] = (input[channel] - mean[channel]) / std[channel]``
197

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

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

206
207
    """

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

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

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

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

227

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

233
234
235
236
    .. warning::
        The output image might be different depending on its type: when downsampling, the interpolation of PIL images
        and tensors is slightly different, because PIL applies antialiasing. This may lead to significant differences
        in the performance of a network. Therefore, it is preferable to train and serve a model with the same input
237
238
        types. See also below the ``antialias`` parameter, which can help making the output of PIL images and tensors
        closer.
239

240
241
242
243
244
    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
245
            (size * height / width, size).
246
247
248

            .. note::
                In torchscript mode size as single int is not supported, use a sequence of length 1: ``[size, ]``.
249
250
251
252
        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.
253
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
254
255
256
257
        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
258
            ``max_size``. As a result, ``size`` might be overruled, i.e the
259
260
261
            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).
262
        antialias (bool, optional): antialias flag. If ``img`` is PIL Image, the flag is ignored and anti-alias
263
264
265
            is always used. If ``img`` is Tensor, the flag is False by default and can be set to True for
            ``InterpolationMode.BILINEAR`` only mode. This can help making the output for PIL images and tensors
            closer.
266
267
268

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

270
271
    """

272
    def __init__(self, size, interpolation=InterpolationMode.BILINEAR, max_size=None, antialias=None):
vfdev's avatar
vfdev committed
273
        super().__init__()
274
275
276
277
278
        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
279
        self.max_size = max_size
280
281
282
283

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

289
        self.interpolation = interpolation
290
        self.antialias = antialias
291

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

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

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

307
308
309
310
311
312
313
314
315
316
317

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
318
319
class CenterCrop(torch.nn.Module):
    """Crops the given image at the center.
320
    If the image is torch Tensor, it is expected
321
322
    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.
323
324
325
326

    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
327
            made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
328
329
330
    """

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

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

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

344
345
346
    def __repr__(self):
        return self.__class__.__name__ + '(size={0})'.format(self.size)

347

348
349
class Pad(torch.nn.Module):
    """Pad the given image on all sides with the given "pad" value.
350
    If the image is torch Tensor, it is expected
351
352
353
    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
354
355

    Args:
356
357
358
        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
359
            this is the padding for the left, top, right and bottom borders respectively.
360
361
362
363

            .. note::
                In torchscript mode padding as single int is not supported, use a sequence of
                length 1: ``[padding, ]``.
364
        fill (number or str or tuple): Pixel fill value for constant fill. Default is 0. If a tuple of
365
            length 3, it is used to fill R, G, B channels respectively.
366
367
368
            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.
369
        padding_mode (str): Type of padding. Should be: constant, edge, reflect or symmetric.
370
            Default is constant.
371
372
373

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

374
375
            - 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
376

377
378
379
            - reflect: pads with reflection of image without repeating the last value on the edge.
              For example, padding [1, 2, 3, 4] with 2 elements on both sides in reflect mode
              will result in [3, 2, 1, 2, 3, 4, 3, 2]
380

381
382
383
            - symmetric: pads with reflection of image repeating the last value on the edge.
              For example, padding [1, 2, 3, 4] with 2 elements on both sides in symmetric mode
              will result in [2, 1, 1, 2, 3, 4, 4, 3]
384
385
    """

386
387
388
389
390
391
392
393
394
395
396
397
398
    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 " +
399
400
401
402
                             "{} element tuple".format(len(padding)))

        self.padding = padding
        self.fill = fill
403
        self.padding_mode = padding_mode
404

405
    def forward(self, img):
406
407
        """
        Args:
408
            img (PIL Image or Tensor): Image to be padded.
409
410

        Returns:
411
            PIL Image or Tensor: Padded image.
412
        """
413
        return F.pad(img, self.padding, self.fill, self.padding_mode)
414

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

419

420
class Lambda:
421
    """Apply a user-defined lambda as a transform. This transform does not support torchscript.
422
423
424
425
426
427

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

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

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

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

438

439
class RandomTransforms:
440
441
442
    """Base class for a list of transformations with randomness

    Args:
443
        transforms (sequence): list of transformations
444
445
446
    """

    def __init__(self, transforms):
447
448
        if not isinstance(transforms, Sequence):
            raise TypeError("Argument transforms should be a sequence")
449
450
451
452
453
454
455
456
457
458
459
460
461
462
        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


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

    .. 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``.
477
478

    Args:
479
        transforms (sequence or torch.nn.Module): list of transformations
480
481
482
483
        p (float): probability
    """

    def __init__(self, transforms, p=0.5):
484
485
        super().__init__()
        self.transforms = transforms
486
487
        self.p = p

488
489
    def forward(self, img):
        if self.p < torch.rand(1):
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
            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):
506
    """Apply a list of transformations in a random order. This transform does not support torchscript.
507
508
509
510
511
512
513
514
515
516
    """
    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):
517
    """Apply single transformation randomly picked from a list. This transform does not support torchscript.
518
    """
519
520
521
522
523
524
525
526
527
528
529
530
531
532
    def __init__(self, transforms, p=None):
        super().__init__(transforms)
        if p is not None and not isinstance(p, Sequence):
            raise TypeError("Argument transforms should be a sequence")
        self.p = p

    def __call__(self, *args):
        t = random.choices(self.transforms, weights=self.p)[0]
        return t(*args)

    def __repr__(self):
        format_string = super().__repr__()
        format_string += '(p={0})'.format(self.p)
        return format_string
533
534


vfdev's avatar
vfdev committed
535
536
class RandomCrop(torch.nn.Module):
    """Crop the given image at a random location.
537
    If the image is torch Tensor, it is expected
538
539
    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
540
541
542
543

    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
544
            made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
545
        padding (int or sequence, optional): Optional padding on each border
vfdev's avatar
vfdev committed
546
            of the image. Default is None. If a single int is provided this
547
548
            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
549
            this is the padding for the left, top, right and bottom borders respectively.
550
551
552
553

            .. note::
                In torchscript mode padding as single int is not supported, use a sequence of
                length 1: ``[padding, ]``.
554
        pad_if_needed (boolean): It will pad the image if smaller than the
ekka's avatar
ekka committed
555
            desired size to avoid raising an exception. Since cropping is done
556
            after padding, the padding seems to be done at a random offset.
557
        fill (number or str or tuple): Pixel fill value for constant fill. Default is 0. If a tuple of
558
            length 3, it is used to fill R, G, B channels respectively.
559
560
561
            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.
562
563
        padding_mode (str): Type of padding. Should be: constant, edge, reflect or symmetric.
            Default is constant.
564

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

567
568
            - 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
569

570
571
572
            - reflect: pads with reflection of image without repeating the last value on the edge.
              For example, padding [1, 2, 3, 4] with 2 elements on both sides in reflect mode
              will result in [3, 2, 1, 2, 3, 4, 3, 2]
573

574
575
576
            - symmetric: pads with reflection of image repeating the last value on the edge.
              For example, padding [1, 2, 3, 4] with 2 elements on both sides in symmetric mode
              will result in [2, 1, 1, 2, 3, 4, 4, 3]
577
578
579
    """

    @staticmethod
vfdev's avatar
vfdev committed
580
    def get_params(img: Tensor, output_size: Tuple[int, int]) -> Tuple[int, int, int, int]:
581
582
583
        """Get parameters for ``crop`` for a random crop.

        Args:
vfdev's avatar
vfdev committed
584
            img (PIL Image or Tensor): Image to be cropped.
585
586
587
588
589
            output_size (tuple): Expected output size of the crop.

        Returns:
            tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
        """
590
        w, h = F.get_image_size(img)
591
        th, tw = output_size
vfdev's avatar
vfdev committed
592
593
594
595
596
597

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

598
599
600
        if w == tw and h == th:
            return 0, 0, h, w

601
602
        i = torch.randint(0, h - th + 1, size=(1, )).item()
        j = torch.randint(0, w - tw + 1, size=(1, )).item()
603
604
        return i, j, th, tw

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

608
609
610
611
        self.size = tuple(_setup_size(
            size, error_msg="Please provide only two dimensions (h, w) for size."
        ))

vfdev's avatar
vfdev committed
612
613
614
615
616
617
        self.padding = padding
        self.pad_if_needed = pad_if_needed
        self.fill = fill
        self.padding_mode = padding_mode

    def forward(self, img):
618
619
        """
        Args:
vfdev's avatar
vfdev committed
620
            img (PIL Image or Tensor): Image to be cropped.
621
622

        Returns:
vfdev's avatar
vfdev committed
623
            PIL Image or Tensor: Cropped image.
624
        """
625
626
        if self.padding is not None:
            img = F.pad(img, self.padding, self.fill, self.padding_mode)
627

628
        width, height = F.get_image_size(img)
629
        # pad the width if needed
vfdev's avatar
vfdev committed
630
631
632
        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)
633
        # pad the height if needed
vfdev's avatar
vfdev committed
634
635
636
        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)
637

638
639
640
641
        i, j, h, w = self.get_params(img, self.size)

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

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

645

646
647
class RandomHorizontalFlip(torch.nn.Module):
    """Horizontally flip the given image randomly with a given probability.
648
    If the image is torch Tensor, it is expected
649
650
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
651
652
653
654
655
656

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

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

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

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

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

675

676
class RandomVerticalFlip(torch.nn.Module):
vfdev's avatar
vfdev committed
677
    """Vertically flip the given image randomly with a given probability.
678
    If the image is torch Tensor, it is expected
679
680
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
681
682
683
684
685
686

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

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

690
    def forward(self, img):
691
692
        """
        Args:
693
            img (PIL Image or Tensor): Image to be flipped.
694
695

        Returns:
696
            PIL Image or Tensor: Randomly flipped image.
697
        """
698
        if torch.rand(1) < self.p:
699
700
701
            return F.vflip(img)
        return img

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

705

706
707
class RandomPerspective(torch.nn.Module):
    """Performs a random perspective transformation of the given image with a given probability.
708
    If the image is torch Tensor, it is expected
709
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
710
711

    Args:
712
713
714
        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.
715
716
717
        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.
718
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
719
720
        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.
721
722
    """

723
    def __init__(self, distortion_scale=0.5, p=0.5, interpolation=InterpolationMode.BILINEAR, fill=0):
724
        super().__init__()
725
        self.p = p
726
727
728
729

        # Backward compatibility with integer value
        if isinstance(interpolation, int):
            warnings.warn(
730
731
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
732
733
734
            )
            interpolation = _interpolation_modes_from_int(interpolation)

735
736
        self.interpolation = interpolation
        self.distortion_scale = distortion_scale
737
738
739
740
741
742

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

743
        self.fill = fill
744

745
    def forward(self, img):
746
747
        """
        Args:
748
            img (PIL Image or Tensor): Image to be Perspectively transformed.
749
750

        Returns:
751
            PIL Image or Tensor: Randomly transformed image.
752
        """
753
754
755
756

        fill = self.fill
        if isinstance(img, Tensor):
            if isinstance(fill, (int, float)):
757
                fill = [float(fill)] * F.get_image_num_channels(img)
758
759
760
            else:
                fill = [float(f) for f in fill]

761
        if torch.rand(1) < self.p:
762
            width, height = F.get_image_size(img)
763
            startpoints, endpoints = self.get_params(width, height, self.distortion_scale)
764
            return F.perspective(img, startpoints, endpoints, self.interpolation, fill)
765
766
767
        return img

    @staticmethod
768
    def get_params(width: int, height: int, distortion_scale: float) -> Tuple[List[List[int]], List[List[int]]]:
769
770
771
        """Get parameters for ``perspective`` for a random perspective transform.

        Args:
772
773
774
            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.
775
776

        Returns:
777
            List containing [top-left, top-right, bottom-right, bottom-left] of the original image,
778
779
            List containing [top-left, top-right, bottom-right, bottom-left] of the transformed image.
        """
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
        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]]
799
800
801
802
803
804
805
        endpoints = [topleft, topright, botright, botleft]
        return startpoints, endpoints

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


806
class RandomResizedCrop(torch.nn.Module):
807
808
    """Crop a random portion of image and resize it to a given size.

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

812
813
814
    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.
815
816

    Args:
817
        size (int or sequence): expected output size of the crop, for each edge. If size is an
818
            int instead of sequence like (h, w), a square output size ``(size, size)`` is
819
            made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
820
821
822

            .. note::
                In torchscript mode size as single int is not supported, use a sequence of length 1: ``[size, ]``.
Nicolas Hug's avatar
Nicolas Hug committed
823
824
        scale (tuple of float): Specifies the lower and upper bounds for the random area of the crop,
            before resizing. The scale is defined with respect to the area of the original image.
825
826
        ratio (tuple of float): lower and upper bounds for the random aspect ratio of the crop, before
            resizing.
827
828
829
830
        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.
831
832
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.

833
834
    """

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

839
        if not isinstance(scale, Sequence):
840
            raise TypeError("Scale should be a sequence")
841
        if not isinstance(ratio, Sequence):
842
            raise TypeError("Ratio should be a sequence")
843
        if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
844
            warnings.warn("Scale and ratio should be of kind (min, max)")
845

846
847
848
        # Backward compatibility with integer value
        if isinstance(interpolation, int):
            warnings.warn(
849
850
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
851
852
853
            )
            interpolation = _interpolation_modes_from_int(interpolation)

854
        self.interpolation = interpolation
855
856
        self.scale = scale
        self.ratio = ratio
857
858

    @staticmethod
859
    def get_params(
860
            img: Tensor, scale: List[float], ratio: List[float]
861
    ) -> Tuple[int, int, int, int]:
862
863
864
        """Get parameters for ``crop`` for a random sized crop.

        Args:
865
            img (PIL Image or Tensor): Input image.
866
867
            scale (list): range of scale of the origin size cropped
            ratio (list): range of aspect ratio of the origin aspect ratio cropped
868
869
870

        Returns:
            tuple: params (i, j, h, w) to be passed to ``crop`` for a random
871
            sized crop.
872
        """
873
        width, height = F.get_image_size(img)
Zhicheng Yan's avatar
Zhicheng Yan committed
874
        area = height * width
875

876
        log_ratio = torch.log(torch.tensor(ratio))
877
        for _ in range(10):
878
            target_area = area * torch.empty(1).uniform_(scale[0], scale[1]).item()
879
880
881
            aspect_ratio = torch.exp(
                torch.empty(1).uniform_(log_ratio[0], log_ratio[1])
            ).item()
882
883
884
885

            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
886
            if 0 < w <= width and 0 < h <= height:
887
888
                i = torch.randint(0, height - h + 1, size=(1,)).item()
                j = torch.randint(0, width - w + 1, size=(1,)).item()
889
890
                return i, j, h, w

891
        # Fallback to central crop
Zhicheng Yan's avatar
Zhicheng Yan committed
892
        in_ratio = float(width) / float(height)
893
        if in_ratio < min(ratio):
Zhicheng Yan's avatar
Zhicheng Yan committed
894
            w = width
895
            h = int(round(w / min(ratio)))
896
        elif in_ratio > max(ratio):
Zhicheng Yan's avatar
Zhicheng Yan committed
897
            h = height
898
            w = int(round(h * max(ratio)))
899
        else:  # whole image
Zhicheng Yan's avatar
Zhicheng Yan committed
900
901
902
903
            w = width
            h = height
        i = (height - h) // 2
        j = (width - w) // 2
904
        return i, j, h, w
905

906
    def forward(self, img):
907
908
        """
        Args:
909
            img (PIL Image or Tensor): Image to be cropped and resized.
910
911

        Returns:
912
            PIL Image or Tensor: Randomly cropped and resized image.
913
        """
914
        i, j, h, w = self.get_params(img, self.scale, self.ratio)
915
916
        return F.resized_crop(img, i, j, h, w, self.size, self.interpolation)

917
    def __repr__(self):
918
        interpolate_str = self.interpolation.value
919
        format_string = self.__class__.__name__ + '(size={0}'.format(self.size)
920
921
        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))
922
923
        format_string += ', interpolation={0})'.format(interpolate_str)
        return format_string
924

925
926
927
928
929
930
931
932
933
934
935

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
936
937
class FiveCrop(torch.nn.Module):
    """Crop the given image into four corners and the central crop.
938
    If the image is torch Tensor, it is expected
vfdev's avatar
vfdev committed
939
940
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
941
942
943
944
945
946
947
948
949

    .. 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.
950
            If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
951
952
953
954
955
956
957
958
959
960
961
962
963
964

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

vfdev's avatar
vfdev committed
968
969
970
971
972
973
974
975
    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
        """
976
977
        return F.five_crop(img, self.size)

978
979
980
    def __repr__(self):
        return self.__class__.__name__ + '(size={0})'.format(self.size)

981

vfdev's avatar
vfdev committed
982
983
984
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).
985
    If the image is torch Tensor, it is expected
vfdev's avatar
vfdev committed
986
987
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
988
989
990
991
992
993
994
995
996

    .. 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
997
            made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
998
        vertical_flip (bool): Use vertical flipping instead of horizontal
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012

    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
1013
        super().__init__()
1014
        self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
1015
1016
        self.vertical_flip = vertical_flip

vfdev's avatar
vfdev committed
1017
1018
1019
1020
1021
1022
1023
1024
    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
        """
1025
1026
        return F.ten_crop(img, self.size, self.vertical_flip)

1027
    def __repr__(self):
1028
        return self.__class__.__name__ + '(size={0}, vertical_flip={1})'.format(self.size, self.vertical_flip)
1029

1030

1031
class LinearTransformation(torch.nn.Module):
ekka's avatar
ekka committed
1032
    """Transform a tensor image with a square transformation matrix and a mean_vector computed
1033
    offline.
1034
    This transform does not support PIL Image.
ekka's avatar
ekka committed
1035
1036
1037
    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
1038
    original shape.
1039

1040
    Applications:
1041
        whitening transformation: Suppose X is a column vector zero-centered data.
1042
1043
1044
        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.

1045
1046
    Args:
        transformation_matrix (Tensor): tensor [D x D], D = C x H x W
ekka's avatar
ekka committed
1047
        mean_vector (Tensor): tensor [D], D = C x H x W
1048
1049
    """

ekka's avatar
ekka committed
1050
    def __init__(self, transformation_matrix, mean_vector):
1051
        super().__init__()
1052
1053
1054
        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
1055
1056
1057

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

1061
1062
1063
1064
        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))

1065
        self.transformation_matrix = transformation_matrix
ekka's avatar
ekka committed
1066
        self.mean_vector = mean_vector
1067

1068
    def forward(self, tensor: Tensor) -> Tensor:
1069
1070
        """
        Args:
vfdev's avatar
vfdev committed
1071
            tensor (Tensor): Tensor image to be whitened.
1072
1073
1074
1075

        Returns:
            Tensor: Transformed image.
        """
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
        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
1088
        transformed_tensor = torch.mm(flat_tensor, self.transformation_matrix)
1089
        tensor = transformed_tensor.view(shape)
1090
1091
        return tensor

1092
    def __repr__(self):
ekka's avatar
ekka committed
1093
1094
1095
        format_string = self.__class__.__name__ + '(transformation_matrix='
        format_string += (str(self.transformation_matrix.tolist()) + ')')
        format_string += (", (mean_vector=" + str(self.mean_vector.tolist()) + ')')
1096
1097
        return format_string

1098

1099
class ColorJitter(torch.nn.Module):
1100
    """Randomly change the brightness, contrast, saturation and hue of an image.
1101
    If the image is torch Tensor, it is expected
1102
1103
    to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
    If img is PIL Image, mode "1", "I", "F" and modes with transparency (alpha channel) are not supported.
1104
1105

    Args:
yaox12's avatar
yaox12 committed
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
        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.
1118
    """
1119

1120
    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):
1121
        super().__init__()
yaox12's avatar
yaox12 committed
1122
1123
1124
1125
1126
1127
        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)

1128
    @torch.jit.unused
yaox12's avatar
yaox12 committed
1129
1130
1131
1132
    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))
1133
            value = [center - float(value), center + float(value)]
yaox12's avatar
yaox12 committed
1134
            if clip_first_on_zero:
1135
                value[0] = max(value[0], 0.0)
yaox12's avatar
yaox12 committed
1136
1137
1138
1139
        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:
1140
            raise TypeError("{} should be a single number or a list/tuple with length 2.".format(name))
yaox12's avatar
yaox12 committed
1141
1142
1143
1144
1145
1146

        # 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
1147
1148

    @staticmethod
1149
1150
1151
1152
1153
1154
    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.
1155

1156
1157
1158
1159
1160
1161
1162
1163
1164
        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.
1165
1166

        Returns:
1167
1168
            tuple: The parameters used to apply the randomized transform
            along with their random order.
1169
        """
1170
        fn_idx = torch.randperm(4)
1171

1172
1173
1174
1175
        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]))
1176

1177
        return fn_idx, b, c, s, h
1178

1179
    def forward(self, img):
1180
1181
        """
        Args:
1182
            img (PIL Image or Tensor): Input image.
1183
1184

        Returns:
1185
1186
            PIL Image or Tensor: Color jittered image.
        """
1187
1188
1189
        fn_idx, brightness_factor, contrast_factor, saturation_factor, hue_factor = \
            self.get_params(self.brightness, self.contrast, self.saturation, self.hue)

1190
        for fn_id in fn_idx:
1191
            if fn_id == 0 and brightness_factor is not None:
1192
                img = F.adjust_brightness(img, brightness_factor)
1193
            elif fn_id == 1 and contrast_factor is not None:
1194
                img = F.adjust_contrast(img, contrast_factor)
1195
            elif fn_id == 2 and saturation_factor is not None:
1196
                img = F.adjust_saturation(img, saturation_factor)
1197
            elif fn_id == 3 and hue_factor is not None:
1198
1199
1200
                img = F.adjust_hue(img, hue_factor)

        return img
1201

1202
    def __repr__(self):
1203
1204
1205
1206
1207
1208
        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
1209

1210

1211
class RandomRotation(torch.nn.Module):
1212
    """Rotate the image by angle.
1213
    If the image is torch Tensor, it is expected
1214
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
1215
1216

    Args:
1217
        degrees (sequence or number): Range of degrees to select from.
1218
1219
            If degrees is a number instead of sequence like (min, max), the range of degrees
            will be (-degrees, +degrees).
1220
1221
1222
        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.
1223
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
1224
1225
1226
1227
        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.
1228
        center (sequence, optional): Optional center of rotation, (x, y). Origin is the upper left corner.
1229
            Default is the center of the image.
1230
1231
        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.
1232
        resample (int, optional): deprecated argument and will be removed since v0.10.0.
1233
            Please use the ``interpolation`` parameter instead.
1234
1235
1236

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

1237
1238
    """

1239
    def __init__(
1240
        self, degrees, interpolation=InterpolationMode.NEAREST, expand=False, center=None, fill=0, resample=None
1241
    ):
1242
        super().__init__()
1243
1244
1245
1246
1247
1248
1249
1250
1251
        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(
1252
1253
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
1254
1255
1256
            )
            interpolation = _interpolation_modes_from_int(interpolation)

1257
        self.degrees = _setup_angle(degrees, name="degrees", req_sizes=(2, ))
1258
1259

        if center is not None:
1260
            _check_sequence_input(center, "center", req_sizes=(2, ))
1261
1262

        self.center = center
1263

1264
        self.resample = self.interpolation = interpolation
1265
        self.expand = expand
1266
1267
1268
1269
1270
1271

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

1272
        self.fill = fill
1273
1274

    @staticmethod
1275
    def get_params(degrees: List[float]) -> float:
1276
1277
1278
        """Get parameters for ``rotate`` for a random rotation.

        Returns:
1279
            float: angle parameter to be passed to ``rotate`` for random rotation.
1280
        """
1281
        angle = float(torch.empty(1).uniform_(float(degrees[0]), float(degrees[1])).item())
1282
1283
        return angle

1284
    def forward(self, img):
1285
        """
1286
        Args:
1287
            img (PIL Image or Tensor): Image to be rotated.
1288
1289

        Returns:
1290
            PIL Image or Tensor: Rotated image.
1291
        """
1292
1293
1294
        fill = self.fill
        if isinstance(img, Tensor):
            if isinstance(fill, (int, float)):
1295
                fill = [float(fill)] * F.get_image_num_channels(img)
1296
1297
            else:
                fill = [float(f) for f in fill]
1298
        angle = self.get_params(self.degrees)
1299
1300

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

1302
    def __repr__(self):
1303
        interpolate_str = self.interpolation.value
1304
        format_string = self.__class__.__name__ + '(degrees={0}'.format(self.degrees)
1305
        format_string += ', interpolation={0}'.format(interpolate_str)
1306
1307
1308
        format_string += ', expand={0}'.format(self.expand)
        if self.center is not None:
            format_string += ', center={0}'.format(self.center)
1309
1310
        if self.fill is not None:
            format_string += ', fill={0}'.format(self.fill)
1311
1312
        format_string += ')'
        return format_string
1313

1314

1315
1316
class RandomAffine(torch.nn.Module):
    """Random affine transformation of the image keeping center invariant.
1317
    If the image is torch Tensor, it is expected
1318
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
1319
1320

    Args:
1321
        degrees (sequence or number): Range of degrees to select from.
1322
            If degrees is a number instead of sequence like (min, max), the range of degrees
1323
            will be (-degrees, +degrees). Set to 0 to deactivate rotations.
1324
1325
1326
1327
1328
1329
        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.
1330
        shear (sequence or number, optional): Range of degrees to select from.
ptrblck's avatar
ptrblck committed
1331
            If shear is a number, a shear parallel to the x axis in the range (-shear, +shear)
1332
1333
            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
1334
            a x-axis shear in (shear[0], shear[1]) and y-axis shear in (shear[2], shear[3]) will be applied.
1335
            Will not apply shear by default.
1336
1337
1338
        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.
1339
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
1340
1341
        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.
1342
        fillcolor (sequence or number, optional): deprecated argument and will be removed since v0.10.0.
1343
            Please use the ``fill`` parameter instead.
1344
        resample (int, optional): deprecated argument and will be removed since v0.10.0.
1345
            Please use the ``interpolation`` parameter instead.
1346
1347
1348

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

1349
1350
    """

1351
    def __init__(
1352
        self, degrees, translate=None, scale=None, shear=None, interpolation=InterpolationMode.NEAREST, fill=0,
1353
1354
        fillcolor=None, resample=None
    ):
1355
        super().__init__()
1356
1357
1358
1359
1360
1361
1362
1363
1364
        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(
1365
1366
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
1367
1368
1369
1370
1371
1372
1373
1374
1375
            )
            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

1376
        self.degrees = _setup_angle(degrees, name="degrees", req_sizes=(2, ))
1377
1378

        if translate is not None:
1379
            _check_sequence_input(translate, "translate", req_sizes=(2, ))
1380
1381
1382
1383
1384
1385
            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:
1386
            _check_sequence_input(scale, "scale", req_sizes=(2, ))
1387
1388
1389
1390
1391
1392
            for s in scale:
                if s <= 0:
                    raise ValueError("scale values should be positive")
        self.scale = scale

        if shear is not None:
1393
            self.shear = _setup_angle(shear, name="shear", req_sizes=(2, 4))
1394
1395
1396
        else:
            self.shear = shear

1397
        self.resample = self.interpolation = interpolation
1398
1399
1400
1401
1402
1403

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

1404
        self.fillcolor = self.fill = fill
1405
1406

    @staticmethod
1407
1408
1409
1410
1411
1412
1413
    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]]:
1414
1415
1416
        """Get parameters for affine transformation

        Returns:
1417
            params to be passed to the affine transformation
1418
        """
1419
        angle = float(torch.empty(1).uniform_(float(degrees[0]), float(degrees[1])).item())
1420
        if translate is not None:
1421
1422
1423
1424
1425
            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)
1426
1427
1428
1429
        else:
            translations = (0, 0)

        if scale_ranges is not None:
1430
            scale = float(torch.empty(1).uniform_(scale_ranges[0], scale_ranges[1]).item())
1431
1432
1433
        else:
            scale = 1.0

1434
        shear_x = shear_y = 0.0
1435
        if shears is not None:
1436
1437
1438
1439
1440
            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)
1441
1442
1443

        return angle, translations, scale, shear

1444
    def forward(self, img):
1445
        """
1446
            img (PIL Image or Tensor): Image to be transformed.
1447
1448

        Returns:
1449
            PIL Image or Tensor: Affine transformed image.
1450
        """
1451
1452
1453
        fill = self.fill
        if isinstance(img, Tensor):
            if isinstance(fill, (int, float)):
1454
                fill = [float(fill)] * F.get_image_num_channels(img)
1455
1456
            else:
                fill = [float(f) for f in fill]
1457

1458
        img_size = F.get_image_size(img)
1459
1460

        ret = self.get_params(self.degrees, self.translate, self.scale, self.shear, img_size)
1461
1462

        return F.affine(img, *ret, interpolation=self.interpolation, fill=fill)
1463
1464
1465
1466
1467
1468
1469
1470
1471

    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}'
1472
        if self.interpolation != InterpolationMode.NEAREST:
1473
1474
1475
            s += ', interpolation={interpolation}'
        if self.fill != 0:
            s += ', fill={fill}'
1476
1477
        s += ')'
        d = dict(self.__dict__)
1478
        d['interpolation'] = self.interpolation.value
1479
1480
1481
        return s.format(name=self.__class__.__name__, **d)


1482
class Grayscale(torch.nn.Module):
1483
    """Convert image to grayscale.
1484
1485
    If the image is torch Tensor, it is expected
    to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions
1486

1487
1488
1489
1490
    Args:
        num_output_channels (int): (1 or 3) number of channels desired for output image

    Returns:
1491
        PIL Image: Grayscale version of the input.
1492
1493
1494

        - If ``num_output_channels == 1`` : returned image is single channel
        - If ``num_output_channels == 3`` : returned image is 3 channel with r == g == b
1495
1496
1497
1498

    """

    def __init__(self, num_output_channels=1):
1499
        super().__init__()
1500
1501
        self.num_output_channels = num_output_channels

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

        Returns:
1508
            PIL Image or Tensor: Grayscaled image.
1509
        """
1510
        return F.rgb_to_grayscale(img, num_output_channels=self.num_output_channels)
1511

1512
    def __repr__(self):
1513
        return self.__class__.__name__ + '(num_output_channels={0})'.format(self.num_output_channels)
1514

1515

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

1521
1522
1523
1524
    Args:
        p (float): probability that image should be converted to grayscale.

    Returns:
1525
        PIL Image or Tensor: Grayscale version of the input image with probability p and unchanged
1526
1527
1528
        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
1529
1530
1531
1532

    """

    def __init__(self, p=0.1):
1533
        super().__init__()
1534
1535
        self.p = p

vfdev's avatar
vfdev committed
1536
    def forward(self, img):
1537
1538
        """
        Args:
1539
            img (PIL Image or Tensor): Image to be converted to grayscale.
1540
1541

        Returns:
1542
            PIL Image or Tensor: Randomly grayscaled image.
1543
        """
1544
        num_output_channels = F.get_image_num_channels(img)
1545
1546
        if torch.rand(1) < self.p:
            return F.rgb_to_grayscale(img, num_output_channels=num_output_channels)
1547
        return img
1548
1549

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


1553
class RandomErasing(torch.nn.Module):
1554
1555
    """ 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
1556
    'Random Erasing Data Augmentation' by Zhong et al. See https://arxiv.org/abs/1708.04896
1557

1558
1559
1560
1561
1562
1563
1564
1565
    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
1566
         inplace: boolean to make this transform inplace. Default set to False.
1567

1568
1569
    Returns:
        Erased Image.
1570

vfdev's avatar
vfdev committed
1571
    Example:
1572
        >>> transform = transforms.Compose([
1573
        >>>   transforms.RandomHorizontalFlip(),
1574
1575
        >>>   transforms.PILToTensor(),
        >>>   transforms.ConvertImageDtype(torch.float),
1576
1577
        >>>   transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        >>>   transforms.RandomErasing(),
1578
1579
1580
        >>> ])
    """

Zhun Zhong's avatar
Zhun Zhong committed
1581
    def __init__(self, p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=0, inplace=False):
1582
1583
1584
1585
1586
1587
1588
1589
1590
        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")
1591
        if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
1592
            warnings.warn("Scale and ratio should be of kind (min, max)")
1593
        if scale[0] < 0 or scale[1] > 1:
1594
            raise ValueError("Scale should be between 0 and 1")
1595
        if p < 0 or p > 1:
1596
            raise ValueError("Random erasing probability should be between 0 and 1")
1597
1598
1599
1600
1601

        self.p = p
        self.scale = scale
        self.ratio = ratio
        self.value = value
1602
        self.inplace = inplace
1603
1604

    @staticmethod
1605
1606
1607
    def get_params(
            img: Tensor, scale: Tuple[float, float], ratio: Tuple[float, float], value: Optional[List[float]] = None
    ) -> Tuple[int, int, int, int, Tensor]:
1608
1609
1610
        """Get parameters for ``erase`` for a random erasing.

        Args:
vfdev's avatar
vfdev committed
1611
            img (Tensor): Tensor image to be erased.
1612
1613
            scale (sequence): range of proportion of erased area against input image.
            ratio (sequence): range of aspect ratio of erased area.
1614
1615
1616
            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]``.
1617
1618
1619
1620

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

1624
        log_ratio = torch.log(torch.tensor(ratio))
1625
        for _ in range(10):
1626
            erase_area = area * torch.empty(1).uniform_(scale[0], scale[1]).item()
1627
1628
1629
            aspect_ratio = torch.exp(
                torch.empty(1).uniform_(log_ratio[0], log_ratio[1])
            ).item()
1630
1631
1632

            h = int(round(math.sqrt(erase_area * aspect_ratio)))
            w = int(round(math.sqrt(erase_area / aspect_ratio)))
1633
1634
1635
1636
1637
1638
1639
            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]
1640

1641
1642
            i = torch.randint(0, img_h - h + 1, size=(1, )).item()
            j = torch.randint(0, img_w - w + 1, size=(1, )).item()
1643
            return i, j, h, w, v
1644

Zhun Zhong's avatar
Zhun Zhong committed
1645
1646
1647
        # Return original image
        return 0, 0, img_h, img_w, img

1648
    def forward(self, img):
1649
1650
        """
        Args:
vfdev's avatar
vfdev committed
1651
            img (Tensor): Tensor image to be erased.
1652
1653
1654
1655

        Returns:
            img (Tensor): Erased Tensor image.
        """
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
        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)
1675
            return F.erase(img, x, y, h, w, v, self.inplace)
1676
        return img
1677

1678
1679
1680
1681
1682
1683
1684
1685
    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

1686

1687
1688
class GaussianBlur(torch.nn.Module):
    """Blurs image with randomly chosen Gaussian blur.
1689
1690
    If the image is torch Tensor, it is expected
    to have [..., C, H, W] shape, where ... means an arbitrary number of leading dimensions.
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724

    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
1725
        """Choose sigma for random gaussian blurring.
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738

        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
1739
            img (PIL Image or Tensor): image to be blurred.
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752

        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


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
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]
1783
1784
1785
1786


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

    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
1817
1818
1819
    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".
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

    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
1849
1850
1851
    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".
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

    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):
1880
1881
    """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.
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
1909
1910
1911
1912

    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.
1913
1914
1915
    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".
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942

    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.
1943
1944
1945
    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".
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968

    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)