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

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

86
    .. _references: https://github.com/pytorch/vision/tree/main/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
    .. 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
236
237
        types. See also below the ``antialias`` parameter, which can help making the output of PIL images and tensors
        closer.
238

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

            .. note::
                In torchscript mode size as single int is not supported, use a sequence of length 1: ``[size, ]``.
248
249
250
251
        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.
252
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
253
254
255
256
        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
257
            ``max_size``. As a result, ``size`` might be overruled, i.e the
258
259
260
            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).
261
        antialias (bool, optional): antialias flag. If ``img`` is PIL Image, the flag is ignored and anti-alias
262
263
264
            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.
265
266
267

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

269
270
    """

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

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

288
        self.interpolation = interpolation
289
        self.antialias = antialias
290

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

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

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

306
307
308
309
310
311
312
313
314
315
316

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

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

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

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

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

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

346

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

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

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

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

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

376
377
378
            - 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]
379

380
381
382
            - 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]
383
384
    """

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

        self.padding = padding
        self.fill = fill
402
        self.padding_mode = padding_mode
403

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

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

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

418

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

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

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

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

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

437

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

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

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


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

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

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

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

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


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

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

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

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

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

569
570
571
            - 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]
572

573
574
575
            - 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]
576
577
578
    """

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

644

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

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

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

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

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

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

674

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

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

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

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

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

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

704

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

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

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

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

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

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

742
        self.fill = fill
743

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

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

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

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

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

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

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

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


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

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

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

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

            .. 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
822
823
        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.
824
825
        ratio (tuple of float): lower and upper bounds for the random aspect ratio of the crop, before
            resizing.
826
827
828
829
        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.
830
831
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.

832
833
    """

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

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

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

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

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

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

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

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

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

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

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

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

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

924
925
926
927
928
929
930
931
932
933
934

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

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

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

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

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

980

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

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

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

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

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

1029

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

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

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

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

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

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

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

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

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

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

1097

1098
class ColorJitter(torch.nn.Module):
1099
    """Randomly change the brightness, contrast, saturation and hue of an image.
1100
    If the image is torch Tensor, it is expected
1101
1102
    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.
1103
1104

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

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

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

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

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

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

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

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

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

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

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

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

        return img
1200

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

1209

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

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

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

1236
1237
    """

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

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

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

        self.center = center
1262

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

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

1271
        self.fill = fill
1272
1273

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

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

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

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

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

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

1313

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

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

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

1348
1349
    """

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

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

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

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

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

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

1403
        self.fillcolor = self.fill = fill
1404
1405

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

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

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

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

        return angle, translations, scale, shear

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

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

1457
        img_size = F.get_image_size(img)
1458
1459

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

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

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


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

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

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

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

    """

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

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

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

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

1514

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

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

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

    """

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

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

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

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


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

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

1567
1568
    Returns:
        Erased Image.
1569

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

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

        self.p = p
        self.scale = scale
        self.ratio = ratio
        self.value = value
1600
        self.inplace = inplace
1601
1602

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

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

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

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

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

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

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

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

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

1676
1677
1678
1679
1680
1681
1682
1683
    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

1684

1685
1686
class GaussianBlur(torch.nn.Module):
    """Blurs image with randomly chosen Gaussian blur.
1687
1688
    If the image is torch Tensor, it is expected
    to have [..., C, H, W] shape, where ... means an arbitrary number of leading dimensions.
1689
1690
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

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

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

        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


1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
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]
1781
1782
1783
1784


class RandomInvert(torch.nn.Module):
    """Inverts the colors of the given image randomly with a given probability.
1785
1786
1787
    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".
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:
        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
1815
1816
1817
    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".
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846

    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
1847
1848
1849
    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".
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

    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):
1878
1879
    """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.
1880
1881
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

    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.
1911
1912
1913
    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".
1914
1915
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

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

    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)