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

import torch
from torch import Tensor

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

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

19

20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
__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",
    "RandomPerspective",
    "RandomErasing",
    "GaussianBlur",
    "InterpolationMode",
    "RandomInvert",
    "RandomPosterize",
    "RandomSolarize",
    "RandomAdjustSharpness",
    "RandomAutocontrast",
    "RandomEqualize",
]
59

60

61
class Compose:
62
63
    """Composes several transforms together. This transform does not support torchscript.
    Please, see the note below.
64
65
66
67
68
69
70

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

    Example:
        >>> transforms.Compose([
        >>>     transforms.CenterCrop(10),
71
72
        >>>     transforms.PILToTensor(),
        >>>     transforms.ConvertImageDtype(torch.float),
73
        >>> ])
74
75
76
77
78
79
80
81
82
83
84
85
86

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

87
88
89
90
91
92
93
94
95
96
    """

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

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

97
    def __repr__(self):
98
        format_string = self.__class__.__name__ + "("
99
        for t in self.transforms:
100
101
102
            format_string += "\n"
            format_string += "    {0}".format(t)
        format_string += "\n)"
103
104
        return format_string

105

106
class ToTensor:
107
    """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor. This transform does not support torchscript.
108
109

    Converts a PIL Image or numpy.ndarray (H x W x C) in the range
surgan12's avatar
surgan12 committed
110
111
112
113
114
    [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.
115
116
117
118
119

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

120
    .. _references: https://github.com/pytorch/vision/tree/main/references/segmentation
121
122
123
124
125
126
127
128
129
130
131
132
    """

    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)

133
    def __repr__(self):
134
        return self.__class__.__name__ + "()"
135

136

137
class PILToTensor:
138
    """Convert a ``PIL Image`` to a tensor of the same type. This transform does not support torchscript.
139

vfdev's avatar
vfdev committed
140
    Converts a PIL Image (H x W x C) to a Tensor of shape (C x H x W).
141
142
143
144
145
146
147
148
149
150
151
152
153
    """

    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):
154
        return self.__class__.__name__ + "()"
155
156


157
class ConvertImageDtype(torch.nn.Module):
158
    """Convert a tensor image to the given ``dtype`` and scale the values accordingly
159
    This function does not support PIL Image.
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176

    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:
177
        super().__init__()
178
179
        self.dtype = dtype

vfdev's avatar
vfdev committed
180
    def forward(self, image):
181
182
183
        return F.convert_image_dtype(image, self.dtype)


184
class ToPILImage:
185
    """Convert a tensor or an ndarray to PIL Image. This transform does not support torchscript.
186
187
188
189
190
191
192

    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
193
194
195
196
197
            - 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``).
198

csukuangfj's avatar
csukuangfj committed
199
    .. _PIL.Image mode: https://pillow.readthedocs.io/en/latest/handbook/concepts.html#concept-modes
200
    """
201

202
203
204
205
206
207
208
209
210
211
212
213
214
215
    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)

216
    def __repr__(self):
217
        format_string = self.__class__.__name__ + "("
218
        if self.mode is not None:
219
220
            format_string += "mode={0}".format(self.mode)
        format_string += ")"
221
        return format_string
222

223

224
class Normalize(torch.nn.Module):
Fang Gao's avatar
Fang Gao committed
225
    """Normalize a tensor image with mean and standard deviation.
226
    This transform does not support PIL Image.
227
228
229
    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
230
    ``output[channel] = (input[channel] - mean[channel]) / std[channel]``
231

232
    .. note::
233
        This transform acts out of place, i.e., it does not mutate the input tensor.
234

235
236
237
    Args:
        mean (sequence): Sequence of means for each channel.
        std (sequence): Sequence of standard deviations for each channel.
238
239
        inplace(bool,optional): Bool to make this operation in-place.

240
241
    """

surgan12's avatar
surgan12 committed
242
    def __init__(self, mean, std, inplace=False):
243
        super().__init__()
244
245
        self.mean = mean
        self.std = std
surgan12's avatar
surgan12 committed
246
        self.inplace = inplace
247

248
    def forward(self, tensor: Tensor) -> Tensor:
249
250
        """
        Args:
vfdev's avatar
vfdev committed
251
            tensor (Tensor): Tensor image to be normalized.
252
253
254
255

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

258
    def __repr__(self):
259
        return self.__class__.__name__ + "(mean={0}, std={1})".format(self.mean, self.std)
260

261

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

267
268
269
270
    .. 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
271
272
        types. See also below the ``antialias`` parameter, which can help making the output of PIL images and tensors
        closer.
273

274
275
276
277
278
    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
279
            (size * height / width, size).
280
281
282

            .. note::
                In torchscript mode size as single int is not supported, use a sequence of length 1: ``[size, ]``.
283
284
285
286
        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.
287
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
288
289
290
291
        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
292
            ``max_size``. As a result, ``size`` might be overruled, i.e the
293
294
295
            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).
296
        antialias (bool, optional): antialias flag. If ``img`` is PIL Image, the flag is ignored and anti-alias
297
298
299
            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.
300
301
302

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

304
305
    """

306
    def __init__(self, size, interpolation=InterpolationMode.BILINEAR, max_size=None, antialias=None):
vfdev's avatar
vfdev committed
307
        super().__init__()
308
309
310
311
312
        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
313
        self.max_size = max_size
314
315
316
317

        # Backward compatibility with integer value
        if isinstance(interpolation, int):
            warnings.warn(
318
319
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
320
321
322
            )
            interpolation = _interpolation_modes_from_int(interpolation)

323
        self.interpolation = interpolation
324
        self.antialias = antialias
325

vfdev's avatar
vfdev committed
326
    def forward(self, img):
327
328
        """
        Args:
vfdev's avatar
vfdev committed
329
            img (PIL Image or Tensor): Image to be scaled.
330
331

        Returns:
vfdev's avatar
vfdev committed
332
            PIL Image or Tensor: Rescaled image.
333
        """
334
        return F.resize(img, self.size, self.interpolation, self.max_size, self.antialias)
335

336
    def __repr__(self):
337
        interpolate_str = self.interpolation.value
338
339
340
        return self.__class__.__name__ + "(size={0}, interpolation={1}, max_size={2}, antialias={3})".format(
            self.size, interpolate_str, self.max_size, self.antialias
        )
341

342
343
344
345
346

class Scale(Resize):
    """
    Note: This transform is deprecated in favor of Resize.
    """
347

348
    def __init__(self, *args, **kwargs):
349
350
351
        warnings.warn(
            "The use of the transforms.Scale transform is deprecated, " + "please use transforms.Resize instead."
        )
352
353
354
        super(Scale, self).__init__(*args, **kwargs)


vfdev's avatar
vfdev committed
355
356
class CenterCrop(torch.nn.Module):
    """Crops the given image at the center.
357
    If the image is torch Tensor, it is expected
358
359
    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.
360
361
362
363

    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
364
            made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
365
366
367
    """

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

vfdev's avatar
vfdev committed
371
    def forward(self, img):
372
373
        """
        Args:
vfdev's avatar
vfdev committed
374
            img (PIL Image or Tensor): Image to be cropped.
375
376

        Returns:
vfdev's avatar
vfdev committed
377
            PIL Image or Tensor: Cropped image.
378
379
380
        """
        return F.center_crop(img, self.size)

381
    def __repr__(self):
382
        return self.__class__.__name__ + "(size={0})".format(self.size)
383

384

385
386
class Pad(torch.nn.Module):
    """Pad the given image on all sides with the given "pad" value.
387
    If the image is torch Tensor, it is expected
388
389
390
    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
391
392

    Args:
393
394
395
        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
396
            this is the padding for the left, top, right and bottom borders respectively.
397
398
399
400

            .. note::
                In torchscript mode padding as single int is not supported, use a sequence of
                length 1: ``[padding, ]``.
401
        fill (number or str or tuple): Pixel fill value for constant fill. Default is 0. If a tuple of
402
            length 3, it is used to fill R, G, B channels respectively.
403
404
405
            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.
406
        padding_mode (str): Type of padding. Should be: constant, edge, reflect or symmetric.
407
            Default is constant.
408
409
410

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

411
412
            - 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
413

414
415
416
            - 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]
417

418
419
420
            - 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]
421
422
    """

423
424
425
426
427
428
429
430
431
432
433
434
    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]:
435
436
437
            raise ValueError(
                "Padding must be an int or a 1, 2, or 4 element tuple, not a " + "{} element tuple".format(len(padding))
            )
438
439
440

        self.padding = padding
        self.fill = fill
441
        self.padding_mode = padding_mode
442

443
    def forward(self, img):
444
445
        """
        Args:
446
            img (PIL Image or Tensor): Image to be padded.
447
448

        Returns:
449
            PIL Image or Tensor: Padded image.
450
        """
451
        return F.pad(img, self.padding, self.fill, self.padding_mode)
452

453
    def __repr__(self):
454
455
456
        return self.__class__.__name__ + "(padding={0}, fill={1}, padding_mode={2})".format(
            self.padding, self.fill, self.padding_mode
        )
457

458

459
class Lambda:
460
    """Apply a user-defined lambda as a transform. This transform does not support torchscript.
461
462
463
464
465
466

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

    def __init__(self, lambd):
467
468
        if not callable(lambd):
            raise TypeError("Argument lambd should be callable, got {}".format(repr(type(lambd).__name__)))
469
470
471
472
473
        self.lambd = lambd

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

474
    def __repr__(self):
475
        return self.__class__.__name__ + "()"
476

477

478
class RandomTransforms:
479
480
481
    """Base class for a list of transformations with randomness

    Args:
482
        transforms (sequence): list of transformations
483
484
485
    """

    def __init__(self, transforms):
486
487
        if not isinstance(transforms, Sequence):
            raise TypeError("Argument transforms should be a sequence")
488
489
490
491
492
493
        self.transforms = transforms

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

    def __repr__(self):
494
        format_string = self.__class__.__name__ + "("
495
        for t in self.transforms:
496
497
498
            format_string += "\n"
            format_string += "    {0}".format(t)
        format_string += "\n)"
499
500
501
        return format_string


502
class RandomApply(torch.nn.Module):
503
    """Apply randomly a list of transformations with a given probability.
504
505
506
507
508
509
510
511
512
513
514
515

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

    Args:
518
        transforms (sequence or torch.nn.Module): list of transformations
519
520
521
522
        p (float): probability
    """

    def __init__(self, transforms, p=0.5):
523
524
        super().__init__()
        self.transforms = transforms
525
526
        self.p = p

527
528
    def forward(self, img):
        if self.p < torch.rand(1):
529
530
531
532
533
534
            return img
        for t in self.transforms:
            img = t(img)
        return img

    def __repr__(self):
535
536
        format_string = self.__class__.__name__ + "("
        format_string += "\n    p={}".format(self.p)
537
        for t in self.transforms:
538
539
540
            format_string += "\n"
            format_string += "    {0}".format(t)
        format_string += "\n)"
541
542
543
544
        return format_string


class RandomOrder(RandomTransforms):
545
546
    """Apply a list of transformations in a random order. This transform does not support torchscript."""

547
548
549
550
551
552
553
554
555
    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):
556
557
    """Apply single transformation randomly picked from a list. This transform does not support torchscript."""

558
559
560
561
562
563
564
565
566
567
568
569
    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__()
570
        format_string += "(p={0})".format(self.p)
571
        return format_string
572
573


vfdev's avatar
vfdev committed
574
575
class RandomCrop(torch.nn.Module):
    """Crop the given image at a random location.
576
    If the image is torch Tensor, it is expected
577
578
    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
579
580
581
582

    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
583
            made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
584
        padding (int or sequence, optional): Optional padding on each border
vfdev's avatar
vfdev committed
585
            of the image. Default is None. If a single int is provided this
586
587
            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
588
            this is the padding for the left, top, right and bottom borders respectively.
589
590
591
592

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

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

606
607
            - 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
608

609
610
611
            - 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]
612

613
614
615
            - 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]
616
617
618
    """

    @staticmethod
vfdev's avatar
vfdev committed
619
    def get_params(img: Tensor, output_size: Tuple[int, int]) -> Tuple[int, int, int, int]:
620
621
622
        """Get parameters for ``crop`` for a random crop.

        Args:
vfdev's avatar
vfdev committed
623
            img (PIL Image or Tensor): Image to be cropped.
624
625
626
627
628
            output_size (tuple): Expected output size of the crop.

        Returns:
            tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
        """
629
        w, h = F.get_image_size(img)
630
        th, tw = output_size
vfdev's avatar
vfdev committed
631
632

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

635
636
637
        if w == tw and h == th:
            return 0, 0, h, w

638
639
        i = torch.randint(0, h - th + 1, size=(1,)).item()
        j = torch.randint(0, w - tw + 1, size=(1,)).item()
640
641
        return i, j, th, tw

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

645
        self.size = tuple(_setup_size(size, error_msg="Please provide only two dimensions (h, w) for size."))
646

vfdev's avatar
vfdev committed
647
648
649
650
651
652
        self.padding = padding
        self.pad_if_needed = pad_if_needed
        self.fill = fill
        self.padding_mode = padding_mode

    def forward(self, img):
653
654
        """
        Args:
vfdev's avatar
vfdev committed
655
            img (PIL Image or Tensor): Image to be cropped.
656
657

        Returns:
vfdev's avatar
vfdev committed
658
            PIL Image or Tensor: Cropped image.
659
        """
660
661
        if self.padding is not None:
            img = F.pad(img, self.padding, self.fill, self.padding_mode)
662

663
        width, height = F.get_image_size(img)
664
        # pad the width if needed
vfdev's avatar
vfdev committed
665
666
667
        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)
668
        # pad the height if needed
vfdev's avatar
vfdev committed
669
670
671
        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)
672

673
674
675
676
        i, j, h, w = self.get_params(img, self.size)

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

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

680

681
682
class RandomHorizontalFlip(torch.nn.Module):
    """Horizontally flip the given image randomly with a given probability.
683
    If the image is torch Tensor, it is expected
684
685
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
686
687
688
689
690
691

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

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

695
    def forward(self, img):
696
697
        """
        Args:
698
            img (PIL Image or Tensor): Image to be flipped.
699
700

        Returns:
701
            PIL Image or Tensor: Randomly flipped image.
702
        """
703
        if torch.rand(1) < self.p:
704
705
706
            return F.hflip(img)
        return img

707
    def __repr__(self):
708
        return self.__class__.__name__ + "(p={})".format(self.p)
709

710

711
class RandomVerticalFlip(torch.nn.Module):
vfdev's avatar
vfdev committed
712
    """Vertically flip the given image randomly with a given probability.
713
    If the image is torch Tensor, it is expected
714
715
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
716
717
718
719
720
721

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

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

725
    def forward(self, img):
726
727
        """
        Args:
728
            img (PIL Image or Tensor): Image to be flipped.
729
730

        Returns:
731
            PIL Image or Tensor: Randomly flipped image.
732
        """
733
        if torch.rand(1) < self.p:
734
735
736
            return F.vflip(img)
        return img

737
    def __repr__(self):
738
        return self.__class__.__name__ + "(p={})".format(self.p)
739

740

741
742
class RandomPerspective(torch.nn.Module):
    """Performs a random perspective transformation of the given image with a given probability.
743
    If the image is torch Tensor, it is expected
744
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
745
746

    Args:
747
748
749
        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.
750
751
752
        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.
753
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
754
755
        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.
756
757
    """

758
    def __init__(self, distortion_scale=0.5, p=0.5, interpolation=InterpolationMode.BILINEAR, fill=0):
759
        super().__init__()
760
        self.p = p
761
762
763
764

        # Backward compatibility with integer value
        if isinstance(interpolation, int):
            warnings.warn(
765
766
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
767
768
769
            )
            interpolation = _interpolation_modes_from_int(interpolation)

770
771
        self.interpolation = interpolation
        self.distortion_scale = distortion_scale
772
773
774
775
776
777

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

778
        self.fill = fill
779

780
    def forward(self, img):
781
782
        """
        Args:
783
            img (PIL Image or Tensor): Image to be Perspectively transformed.
784
785

        Returns:
786
            PIL Image or Tensor: Randomly transformed image.
787
        """
788
789
790
791

        fill = self.fill
        if isinstance(img, Tensor):
            if isinstance(fill, (int, float)):
792
                fill = [float(fill)] * F.get_image_num_channels(img)
793
794
795
            else:
                fill = [float(f) for f in fill]

796
        if torch.rand(1) < self.p:
797
            width, height = F.get_image_size(img)
798
            startpoints, endpoints = self.get_params(width, height, self.distortion_scale)
799
            return F.perspective(img, startpoints, endpoints, self.interpolation, fill)
800
801
802
        return img

    @staticmethod
803
    def get_params(width: int, height: int, distortion_scale: float) -> Tuple[List[List[int]], List[List[int]]]:
804
805
806
        """Get parameters for ``perspective`` for a random perspective transform.

        Args:
807
808
809
            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.
810
811

        Returns:
812
            List containing [top-left, top-right, bottom-right, bottom-left] of the original image,
813
814
            List containing [top-left, top-right, bottom-right, bottom-left] of the transformed image.
        """
815
816
817
        half_height = height // 2
        half_width = width // 2
        topleft = [
818
819
            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()),
820
821
        ]
        topright = [
822
823
            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()),
824
825
        ]
        botright = [
826
827
            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()),
828
829
        ]
        botleft = [
830
831
            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()),
832
833
        ]
        startpoints = [[0, 0], [width - 1, 0], [width - 1, height - 1], [0, height - 1]]
834
835
836
837
        endpoints = [topleft, topright, botright, botleft]
        return startpoints, endpoints

    def __repr__(self):
838
        return self.__class__.__name__ + "(p={})".format(self.p)
839
840


841
class RandomResizedCrop(torch.nn.Module):
842
843
    """Crop a random portion of image and resize it to a given size.

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

847
848
849
    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.
850
851

    Args:
852
        size (int or sequence): expected output size of the crop, for each edge. If size is an
853
            int instead of sequence like (h, w), a square output size ``(size, size)`` is
854
            made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
855
856
857

            .. 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
858
859
        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.
860
861
        ratio (tuple of float): lower and upper bounds for the random aspect ratio of the crop, before
            resizing.
862
863
864
865
        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.
866
867
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.

868
869
    """

870
    def __init__(self, size, scale=(0.08, 1.0), ratio=(3.0 / 4.0, 4.0 / 3.0), interpolation=InterpolationMode.BILINEAR):
871
        super().__init__()
872
        self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
873

874
        if not isinstance(scale, Sequence):
875
            raise TypeError("Scale should be a sequence")
876
        if not isinstance(ratio, Sequence):
877
            raise TypeError("Ratio should be a sequence")
878
        if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
879
            warnings.warn("Scale and ratio should be of kind (min, max)")
880

881
882
883
        # Backward compatibility with integer value
        if isinstance(interpolation, int):
            warnings.warn(
884
885
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
886
887
888
            )
            interpolation = _interpolation_modes_from_int(interpolation)

889
        self.interpolation = interpolation
890
891
        self.scale = scale
        self.ratio = ratio
892
893

    @staticmethod
894
    def get_params(img: Tensor, scale: List[float], ratio: List[float]) -> Tuple[int, int, int, int]:
895
896
897
        """Get parameters for ``crop`` for a random sized crop.

        Args:
898
            img (PIL Image or Tensor): Input image.
899
900
            scale (list): range of scale of the origin size cropped
            ratio (list): range of aspect ratio of the origin aspect ratio cropped
901
902
903

        Returns:
            tuple: params (i, j, h, w) to be passed to ``crop`` for a random
904
            sized crop.
905
        """
906
        width, height = F.get_image_size(img)
Zhicheng Yan's avatar
Zhicheng Yan committed
907
        area = height * width
908

909
        log_ratio = torch.log(torch.tensor(ratio))
910
        for _ in range(10):
911
            target_area = area * torch.empty(1).uniform_(scale[0], scale[1]).item()
912
            aspect_ratio = torch.exp(torch.empty(1).uniform_(log_ratio[0], log_ratio[1])).item()
913
914
915
916

            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
917
            if 0 < w <= width and 0 < h <= height:
918
919
                i = torch.randint(0, height - h + 1, size=(1,)).item()
                j = torch.randint(0, width - w + 1, size=(1,)).item()
920
921
                return i, j, h, w

922
        # Fallback to central crop
Zhicheng Yan's avatar
Zhicheng Yan committed
923
        in_ratio = float(width) / float(height)
924
        if in_ratio < min(ratio):
Zhicheng Yan's avatar
Zhicheng Yan committed
925
            w = width
926
            h = int(round(w / min(ratio)))
927
        elif in_ratio > max(ratio):
Zhicheng Yan's avatar
Zhicheng Yan committed
928
            h = height
929
            w = int(round(h * max(ratio)))
930
        else:  # whole image
Zhicheng Yan's avatar
Zhicheng Yan committed
931
932
933
934
            w = width
            h = height
        i = (height - h) // 2
        j = (width - w) // 2
935
        return i, j, h, w
936

937
    def forward(self, img):
938
939
        """
        Args:
940
            img (PIL Image or Tensor): Image to be cropped and resized.
941
942

        Returns:
943
            PIL Image or Tensor: Randomly cropped and resized image.
944
        """
945
        i, j, h, w = self.get_params(img, self.scale, self.ratio)
946
947
        return F.resized_crop(img, i, j, h, w, self.size, self.interpolation)

948
    def __repr__(self):
949
        interpolate_str = self.interpolation.value
950
951
952
953
        format_string = self.__class__.__name__ + "(size={0}".format(self.size)
        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))
        format_string += ", interpolation={0})".format(interpolate_str)
954
        return format_string
955

956
957
958
959
960

class RandomSizedCrop(RandomResizedCrop):
    """
    Note: This transform is deprecated in favor of RandomResizedCrop.
    """
961

962
    def __init__(self, *args, **kwargs):
963
964
965
966
        warnings.warn(
            "The use of the transforms.RandomSizedCrop transform is deprecated, "
            + "please use transforms.RandomResizedCrop instead."
        )
967
968
969
        super(RandomSizedCrop, self).__init__(*args, **kwargs)


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

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

    Args:
         size (sequence or int): Desired output size of the crop. If size is an ``int``
            instead of sequence like (h, w), a square crop of size (size, size) is made.
984
            If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
985
986
987
988
989
990
991
992
993
994
995
996
997
998

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

vfdev's avatar
vfdev committed
1002
1003
1004
1005
1006
1007
1008
1009
    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
        """
1010
1011
        return F.five_crop(img, self.size)

1012
    def __repr__(self):
1013
        return self.__class__.__name__ + "(size={0})".format(self.size)
1014

1015

vfdev's avatar
vfdev committed
1016
1017
1018
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).
1019
    If the image is torch Tensor, it is expected
vfdev's avatar
vfdev committed
1020
1021
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
1022
1023
1024
1025
1026
1027
1028
1029
1030

    .. 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
1031
            made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
1032
        vertical_flip (bool): Use vertical flipping instead of horizontal
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046

    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
1047
        super().__init__()
1048
        self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
1049
1050
        self.vertical_flip = vertical_flip

vfdev's avatar
vfdev committed
1051
1052
1053
1054
1055
1056
1057
1058
    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
        """
1059
1060
        return F.ten_crop(img, self.size, self.vertical_flip)

1061
    def __repr__(self):
1062
        return self.__class__.__name__ + "(size={0}, vertical_flip={1})".format(self.size, self.vertical_flip)
1063

1064

1065
class LinearTransformation(torch.nn.Module):
ekka's avatar
ekka committed
1066
    """Transform a tensor image with a square transformation matrix and a mean_vector computed
1067
    offline.
1068
    This transform does not support PIL Image.
ekka's avatar
ekka committed
1069
1070
1071
    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
1072
    original shape.
1073

1074
    Applications:
1075
        whitening transformation: Suppose X is a column vector zero-centered data.
1076
1077
1078
        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.

1079
1080
    Args:
        transformation_matrix (Tensor): tensor [D x D], D = C x H x W
ekka's avatar
ekka committed
1081
        mean_vector (Tensor): tensor [D], D = C x H x W
1082
1083
    """

ekka's avatar
ekka committed
1084
    def __init__(self, transformation_matrix, mean_vector):
1085
        super().__init__()
1086
        if transformation_matrix.size(0) != transformation_matrix.size(1):
1087
1088
1089
1090
            raise ValueError(
                "transformation_matrix should be square. Got "
                + "[{} x {}] rectangular matrix.".format(*transformation_matrix.size())
            )
ekka's avatar
ekka committed
1091
1092

        if mean_vector.size(0) != transformation_matrix.size(0):
1093
1094
1095
1096
1097
1098
            raise ValueError(
                "mean_vector should have the same length {}".format(mean_vector.size(0))
                + " as any one of the dimensions of the transformation_matrix [{}]".format(
                    tuple(transformation_matrix.size())
                )
            )
ekka's avatar
ekka committed
1099

1100
        if transformation_matrix.device != mean_vector.device:
1101
1102
1103
1104
1105
            raise ValueError(
                "Input tensors should be on the same device. Got {} and {}".format(
                    transformation_matrix.device, mean_vector.device
                )
            )
1106

1107
        self.transformation_matrix = transformation_matrix
ekka's avatar
ekka committed
1108
        self.mean_vector = mean_vector
1109

1110
    def forward(self, tensor: Tensor) -> Tensor:
1111
1112
        """
        Args:
vfdev's avatar
vfdev committed
1113
            tensor (Tensor): Tensor image to be whitened.
1114
1115
1116
1117

        Returns:
            Tensor: Transformed image.
        """
1118
1119
1120
        shape = tensor.shape
        n = shape[-3] * shape[-2] * shape[-1]
        if n != self.transformation_matrix.shape[0]:
1121
1122
1123
1124
1125
            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])
            )
1126
1127

        if tensor.device.type != self.mean_vector.device.type:
1128
1129
1130
1131
            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)
            )
1132
1133

        flat_tensor = tensor.view(-1, n) - self.mean_vector
1134
        transformed_tensor = torch.mm(flat_tensor, self.transformation_matrix)
1135
        tensor = transformed_tensor.view(shape)
1136
1137
        return tensor

1138
    def __repr__(self):
1139
1140
1141
        format_string = self.__class__.__name__ + "(transformation_matrix="
        format_string += str(self.transformation_matrix.tolist()) + ")"
        format_string += ", (mean_vector=" + str(self.mean_vector.tolist()) + ")"
1142
1143
        return format_string

1144

1145
class ColorJitter(torch.nn.Module):
1146
    """Randomly change the brightness, contrast, saturation and hue of an image.
1147
    If the image is torch Tensor, it is expected
1148
1149
    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.
1150
1151

    Args:
yaox12's avatar
yaox12 committed
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
        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.
1164
    """
1165

1166
    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):
1167
        super().__init__()
1168
1169
1170
1171
        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)
yaox12's avatar
yaox12 committed
1172

1173
    @torch.jit.unused
1174
    def _check_input(self, value, name, center=1, bound=(0, float("inf")), clip_first_on_zero=True):
yaox12's avatar
yaox12 committed
1175
1176
1177
        if isinstance(value, numbers.Number):
            if value < 0:
                raise ValueError("If {} is a single number, it must be non negative.".format(name))
1178
            value = [center - float(value), center + float(value)]
yaox12's avatar
yaox12 committed
1179
            if clip_first_on_zero:
1180
                value[0] = max(value[0], 0.0)
yaox12's avatar
yaox12 committed
1181
1182
1183
1184
        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:
1185
            raise TypeError("{} should be a single number or a list/tuple with length 2.".format(name))
yaox12's avatar
yaox12 committed
1186
1187
1188
1189
1190
1191

        # 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
1192
1193

    @staticmethod
1194
1195
1196
1197
1198
1199
    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]]:
1200
        """Get the parameters for the randomized transform to be applied on image.
1201

1202
1203
1204
1205
1206
1207
1208
1209
1210
        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.
1211
1212

        Returns:
1213
1214
            tuple: The parameters used to apply the randomized transform
            along with their random order.
1215
        """
1216
        fn_idx = torch.randperm(4)
1217

1218
1219
1220
1221
        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]))
1222

1223
        return fn_idx, b, c, s, h
1224

1225
    def forward(self, img):
1226
1227
        """
        Args:
1228
            img (PIL Image or Tensor): Input image.
1229
1230

        Returns:
1231
1232
            PIL Image or Tensor: Color jittered image.
        """
1233
1234
1235
        fn_idx, brightness_factor, contrast_factor, saturation_factor, hue_factor = self.get_params(
            self.brightness, self.contrast, self.saturation, self.hue
        )
1236

1237
        for fn_id in fn_idx:
1238
            if fn_id == 0 and brightness_factor is not None:
1239
                img = F.adjust_brightness(img, brightness_factor)
1240
            elif fn_id == 1 and contrast_factor is not None:
1241
                img = F.adjust_contrast(img, contrast_factor)
1242
            elif fn_id == 2 and saturation_factor is not None:
1243
                img = F.adjust_saturation(img, saturation_factor)
1244
            elif fn_id == 3 and hue_factor is not None:
1245
1246
1247
                img = F.adjust_hue(img, hue_factor)

        return img
1248

1249
    def __repr__(self):
1250
1251
1252
1253
1254
        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)
1255
        return format_string
1256

1257

1258
class RandomRotation(torch.nn.Module):
1259
    """Rotate the image by angle.
1260
    If the image is torch Tensor, it is expected
1261
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
1262
1263

    Args:
1264
        degrees (sequence or number): Range of degrees to select from.
1265
1266
            If degrees is a number instead of sequence like (min, max), the range of degrees
            will be (-degrees, +degrees).
1267
1268
1269
        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.
1270
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
1271
1272
1273
1274
        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.
1275
        center (sequence, optional): Optional center of rotation, (x, y). Origin is the upper left corner.
1276
            Default is the center of the image.
1277
1278
        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.
1279
        resample (int, optional): deprecated argument and will be removed since v0.10.0.
1280
            Please use the ``interpolation`` parameter instead.
1281
1282
1283

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

1284
1285
    """

1286
    def __init__(
1287
        self, degrees, interpolation=InterpolationMode.NEAREST, expand=False, center=None, fill=0, resample=None
1288
    ):
1289
        super().__init__()
1290
1291
1292
1293
1294
1295
1296
1297
1298
        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(
1299
1300
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
1301
1302
1303
            )
            interpolation = _interpolation_modes_from_int(interpolation)

1304
        self.degrees = _setup_angle(degrees, name="degrees", req_sizes=(2,))
1305
1306

        if center is not None:
1307
            _check_sequence_input(center, "center", req_sizes=(2,))
1308
1309

        self.center = center
1310

1311
        self.resample = self.interpolation = interpolation
1312
        self.expand = expand
1313
1314
1315
1316
1317
1318

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

1319
        self.fill = fill
1320
1321

    @staticmethod
1322
    def get_params(degrees: List[float]) -> float:
1323
1324
1325
        """Get parameters for ``rotate`` for a random rotation.

        Returns:
1326
            float: angle parameter to be passed to ``rotate`` for random rotation.
1327
        """
1328
        angle = float(torch.empty(1).uniform_(float(degrees[0]), float(degrees[1])).item())
1329
1330
        return angle

1331
    def forward(self, img):
1332
        """
1333
        Args:
1334
            img (PIL Image or Tensor): Image to be rotated.
1335
1336

        Returns:
1337
            PIL Image or Tensor: Rotated image.
1338
        """
1339
1340
1341
        fill = self.fill
        if isinstance(img, Tensor):
            if isinstance(fill, (int, float)):
1342
                fill = [float(fill)] * F.get_image_num_channels(img)
1343
1344
            else:
                fill = [float(f) for f in fill]
1345
        angle = self.get_params(self.degrees)
1346
1347

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

1349
    def __repr__(self):
1350
        interpolate_str = self.interpolation.value
1351
1352
1353
        format_string = self.__class__.__name__ + "(degrees={0}".format(self.degrees)
        format_string += ", interpolation={0}".format(interpolate_str)
        format_string += ", expand={0}".format(self.expand)
1354
        if self.center is not None:
1355
            format_string += ", center={0}".format(self.center)
1356
        if self.fill is not None:
1357
1358
            format_string += ", fill={0}".format(self.fill)
        format_string += ")"
1359
        return format_string
1360

1361

1362
1363
class RandomAffine(torch.nn.Module):
    """Random affine transformation of the image keeping center invariant.
1364
    If the image is torch Tensor, it is expected
1365
    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
1366
1367

    Args:
1368
        degrees (sequence or number): Range of degrees to select from.
1369
            If degrees is a number instead of sequence like (min, max), the range of degrees
1370
            will be (-degrees, +degrees). Set to 0 to deactivate rotations.
1371
1372
1373
1374
1375
1376
        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.
1377
        shear (sequence or number, optional): Range of degrees to select from.
ptrblck's avatar
ptrblck committed
1378
            If shear is a number, a shear parallel to the x axis in the range (-shear, +shear)
1379
1380
            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
1381
            a x-axis shear in (shear[0], shear[1]) and y-axis shear in (shear[2], shear[3]) will be applied.
1382
            Will not apply shear by default.
1383
1384
1385
        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.
1386
            For backward compatibility integer values (e.g. ``PIL.Image.NEAREST``) are still acceptable.
1387
1388
        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.
1389
        fillcolor (sequence or number, optional): deprecated argument and will be removed since v0.10.0.
1390
            Please use the ``fill`` parameter instead.
1391
        resample (int, optional): deprecated argument and will be removed since v0.10.0.
1392
            Please use the ``interpolation`` parameter instead.
1393
1394
1395

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

1396
1397
    """

1398
    def __init__(
1399
1400
1401
1402
1403
1404
1405
1406
1407
        self,
        degrees,
        translate=None,
        scale=None,
        shear=None,
        interpolation=InterpolationMode.NEAREST,
        fill=0,
        fillcolor=None,
        resample=None,
1408
    ):
1409
        super().__init__()
1410
1411
1412
1413
1414
1415
1416
1417
1418
        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(
1419
1420
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
1421
1422
1423
1424
1425
1426
1427
1428
1429
            )
            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

1430
        self.degrees = _setup_angle(degrees, name="degrees", req_sizes=(2,))
1431
1432

        if translate is not None:
1433
            _check_sequence_input(translate, "translate", req_sizes=(2,))
1434
1435
1436
1437
1438
1439
            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:
1440
            _check_sequence_input(scale, "scale", req_sizes=(2,))
1441
1442
1443
1444
1445
1446
            for s in scale:
                if s <= 0:
                    raise ValueError("scale values should be positive")
        self.scale = scale

        if shear is not None:
1447
            self.shear = _setup_angle(shear, name="shear", req_sizes=(2, 4))
1448
1449
1450
        else:
            self.shear = shear

1451
        self.resample = self.interpolation = interpolation
1452
1453
1454
1455
1456
1457

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

1458
        self.fillcolor = self.fill = fill
1459
1460

    @staticmethod
1461
    def get_params(
1462
1463
1464
1465
1466
        degrees: List[float],
        translate: Optional[List[float]],
        scale_ranges: Optional[List[float]],
        shears: Optional[List[float]],
        img_size: List[int],
1467
    ) -> Tuple[float, Tuple[int, int], float, Tuple[float, float]]:
1468
1469
1470
        """Get parameters for affine transformation

        Returns:
1471
            params to be passed to the affine transformation
1472
        """
1473
        angle = float(torch.empty(1).uniform_(float(degrees[0]), float(degrees[1])).item())
1474
        if translate is not None:
1475
1476
1477
1478
1479
            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)
1480
1481
1482
1483
        else:
            translations = (0, 0)

        if scale_ranges is not None:
1484
            scale = float(torch.empty(1).uniform_(scale_ranges[0], scale_ranges[1]).item())
1485
1486
1487
        else:
            scale = 1.0

1488
        shear_x = shear_y = 0.0
1489
        if shears is not None:
1490
1491
1492
1493
1494
            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)
1495
1496
1497

        return angle, translations, scale, shear

1498
    def forward(self, img):
1499
        """
1500
            img (PIL Image or Tensor): Image to be transformed.
1501
1502

        Returns:
1503
            PIL Image or Tensor: Affine transformed image.
1504
        """
1505
1506
1507
        fill = self.fill
        if isinstance(img, Tensor):
            if isinstance(fill, (int, float)):
1508
                fill = [float(fill)] * F.get_image_num_channels(img)
1509
1510
            else:
                fill = [float(f) for f in fill]
1511

1512
        img_size = F.get_image_size(img)
1513
1514

        ret = self.get_params(self.degrees, self.translate, self.scale, self.shear, img_size)
1515
1516

        return F.affine(img, *ret, interpolation=self.interpolation, fill=fill)
1517
1518

    def __repr__(self):
1519
        s = "{name}(degrees={degrees}"
1520
        if self.translate is not None:
1521
            s += ", translate={translate}"
1522
        if self.scale is not None:
1523
            s += ", scale={scale}"
1524
        if self.shear is not None:
1525
            s += ", shear={shear}"
1526
        if self.interpolation != InterpolationMode.NEAREST:
1527
            s += ", interpolation={interpolation}"
1528
        if self.fill != 0:
1529
1530
            s += ", fill={fill}"
        s += ")"
1531
        d = dict(self.__dict__)
1532
        d["interpolation"] = self.interpolation.value
1533
1534
1535
        return s.format(name=self.__class__.__name__, **d)


1536
class Grayscale(torch.nn.Module):
1537
    """Convert image to grayscale.
1538
1539
    If the image is torch Tensor, it is expected
    to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions
1540

1541
1542
1543
1544
    Args:
        num_output_channels (int): (1 or 3) number of channels desired for output image

    Returns:
1545
        PIL Image: Grayscale version of the input.
1546
1547
1548

        - If ``num_output_channels == 1`` : returned image is single channel
        - If ``num_output_channels == 3`` : returned image is 3 channel with r == g == b
1549
1550
1551
1552

    """

    def __init__(self, num_output_channels=1):
1553
        super().__init__()
1554
1555
        self.num_output_channels = num_output_channels

vfdev's avatar
vfdev committed
1556
    def forward(self, img):
1557
1558
        """
        Args:
1559
            img (PIL Image or Tensor): Image to be converted to grayscale.
1560
1561

        Returns:
1562
            PIL Image or Tensor: Grayscaled image.
1563
        """
1564
        return F.rgb_to_grayscale(img, num_output_channels=self.num_output_channels)
1565

1566
    def __repr__(self):
1567
        return self.__class__.__name__ + "(num_output_channels={0})".format(self.num_output_channels)
1568

1569

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

1575
1576
1577
1578
    Args:
        p (float): probability that image should be converted to grayscale.

    Returns:
1579
        PIL Image or Tensor: Grayscale version of the input image with probability p and unchanged
1580
1581
1582
        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
1583
1584
1585
1586

    """

    def __init__(self, p=0.1):
1587
        super().__init__()
1588
1589
        self.p = p

vfdev's avatar
vfdev committed
1590
    def forward(self, img):
1591
1592
        """
        Args:
1593
            img (PIL Image or Tensor): Image to be converted to grayscale.
1594
1595

        Returns:
1596
            PIL Image or Tensor: Randomly grayscaled image.
1597
        """
1598
        num_output_channels = F.get_image_num_channels(img)
1599
1600
        if torch.rand(1) < self.p:
            return F.rgb_to_grayscale(img, num_output_channels=num_output_channels)
1601
        return img
1602
1603

    def __repr__(self):
1604
        return self.__class__.__name__ + "(p={0})".format(self.p)
1605
1606


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

1612
1613
1614
1615
1616
1617
1618
1619
    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
1620
         inplace: boolean to make this transform inplace. Default set to False.
1621

1622
1623
    Returns:
        Erased Image.
1624

vfdev's avatar
vfdev committed
1625
    Example:
1626
        >>> transform = transforms.Compose([
1627
        >>>   transforms.RandomHorizontalFlip(),
1628
1629
        >>>   transforms.PILToTensor(),
        >>>   transforms.ConvertImageDtype(torch.float),
1630
1631
        >>>   transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        >>>   transforms.RandomErasing(),
1632
1633
1634
        >>> ])
    """

Zhun Zhong's avatar
Zhun Zhong committed
1635
    def __init__(self, p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=0, inplace=False):
1636
1637
1638
1639
1640
1641
1642
1643
1644
        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")
1645
        if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
1646
            warnings.warn("Scale and ratio should be of kind (min, max)")
1647
        if scale[0] < 0 or scale[1] > 1:
1648
            raise ValueError("Scale should be between 0 and 1")
1649
        if p < 0 or p > 1:
1650
            raise ValueError("Random erasing probability should be between 0 and 1")
1651
1652
1653
1654
1655

        self.p = p
        self.scale = scale
        self.ratio = ratio
        self.value = value
1656
        self.inplace = inplace
1657
1658

    @staticmethod
1659
    def get_params(
1660
        img: Tensor, scale: Tuple[float, float], ratio: Tuple[float, float], value: Optional[List[float]] = None
1661
    ) -> Tuple[int, int, int, int, Tensor]:
1662
1663
1664
        """Get parameters for ``erase`` for a random erasing.

        Args:
vfdev's avatar
vfdev committed
1665
            img (Tensor): Tensor image to be erased.
1666
1667
            scale (sequence): range of proportion of erased area against input image.
            ratio (sequence): range of aspect ratio of erased area.
1668
1669
1670
            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]``.
1671
1672
1673
1674

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

1678
        log_ratio = torch.log(torch.tensor(ratio))
1679
        for _ in range(10):
1680
            erase_area = area * torch.empty(1).uniform_(scale[0], scale[1]).item()
1681
            aspect_ratio = torch.exp(torch.empty(1).uniform_(log_ratio[0], log_ratio[1])).item()
1682
1683
1684

            h = int(round(math.sqrt(erase_area * aspect_ratio)))
            w = int(round(math.sqrt(erase_area / aspect_ratio)))
1685
1686
1687
1688
1689
1690
1691
            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]
1692

1693
1694
            i = torch.randint(0, img_h - h + 1, size=(1,)).item()
            j = torch.randint(0, img_w - w + 1, size=(1,)).item()
1695
            return i, j, h, w, v
1696

Zhun Zhong's avatar
Zhun Zhong committed
1697
1698
1699
        # Return original image
        return 0, 0, img_h, img_w, img

1700
    def forward(self, img):
1701
1702
        """
        Args:
vfdev's avatar
vfdev committed
1703
            img (Tensor): Tensor image to be erased.
1704
1705
1706
1707

        Returns:
            img (Tensor): Erased Tensor image.
        """
1708
1709
1710
1711
        if torch.rand(1) < self.p:

            # cast self.value to script acceptable type
            if isinstance(self.value, (int, float)):
1712
1713
1714
                value = [
                    self.value,
                ]
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
            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)
1729
            return F.erase(img, x, y, h, w, v, self.inplace)
1730
        return img
1731

1732
    def __repr__(self):
1733
1734
1735
1736
1737
        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)
1738
1739
        return self.__class__.__name__ + s

1740

1741
1742
class GaussianBlur(torch.nn.Module):
    """Blurs image with randomly chosen Gaussian blur.
1743
1744
    If the image is torch Tensor, it is expected
    to have [..., C, H, W] shape, where ... means an arbitrary number of leading dimensions.
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769

    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:
1770
            if not 0.0 < sigma[0] <= sigma[1]:
1771
1772
1773
1774
1775
1776
1777
1778
                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
1779
        """Choose sigma for random gaussian blurring.
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792

        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
1793
            img (PIL Image or Tensor): image to be blurred.
1794
1795
1796
1797
1798
1799
1800
1801

        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):
1802
1803
        s = "(kernel_size={}, ".format(self.kernel_size)
        s += "sigma={})".format(self.sigma)
1804
1805
1806
        return self.__class__.__name__ + s


1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
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))


1828
def _setup_angle(x, name, req_sizes=(2,)):
1829
1830
1831
1832
1833
1834
1835
1836
    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]
1837
1838
1839
1840


class RandomInvert(torch.nn.Module):
    """Inverts the colors of the given image randomly with a given probability.
1841
1842
1843
    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".
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865

    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):
1866
        return self.__class__.__name__ + "(p={})".format(self.p)
1867
1868
1869
1870


class RandomPosterize(torch.nn.Module):
    """Posterize the image randomly with a given probability by reducing the
1871
1872
1873
    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".
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897

    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):
1898
        return self.__class__.__name__ + "(bits={},p={})".format(self.bits, self.p)
1899
1900
1901
1902


class RandomSolarize(torch.nn.Module):
    """Solarize the image randomly with a given probability by inverting all pixel
1903
1904
1905
    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".
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929

    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):
1930
        return self.__class__.__name__ + "(threshold={},p={})".format(self.threshold, self.p)
1931
1932
1933


class RandomAdjustSharpness(torch.nn.Module):
1934
1935
    """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.
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961

    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):
1962
        return self.__class__.__name__ + "(sharpness_factor={},p={})".format(self.sharpness_factor, self.p)
1963
1964
1965
1966


class RandomAutocontrast(torch.nn.Module):
    """Autocontrast the pixels of the given image randomly with a given probability.
1967
1968
1969
    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".
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991

    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):
1992
        return self.__class__.__name__ + "(p={})".format(self.p)
1993
1994
1995
1996


class RandomEqualize(torch.nn.Module):
    """Equalize the histogram of the given image randomly with a given probability.
1997
1998
1999
    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".
2000
2001
2002
2003
2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
2021

    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):
2022
        return self.__class__.__name__ + "(p={})".format(self.p)