transforms.py 68.4 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
           "RandomPerspective", "RandomErasing", "GaussianBlur", "InterpolationMode"]
25

26

27
class Compose:
28
29
    """Composes several transforms together. This transform does not support torchscript.
    Please, see the note below.
30
31
32
33
34
35
36
37
38

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

    Example:
        >>> transforms.Compose([
        >>>     transforms.CenterCrop(10),
        >>>     transforms.ToTensor(),
        >>> ])
39
40
41
42
43
44
45
46
47
48
49
50
51

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

52
53
54
55
56
57
58
59
60
61
    """

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

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

62
63
64
65
66
67
68
69
    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

70

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

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

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

    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)

98
99
100
    def __repr__(self):
        return self.__class__.__name__ + '()'

101

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

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

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


122
class ConvertImageDtype(torch.nn.Module):
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
    """Convert a tensor image to the given ``dtype`` and scale the values accordingly

    Args:
        dtype (torch.dtype): Desired data type of the output

    .. note::

        When converting from a smaller to a larger integer ``dtype`` the maximum values are **not** mapped exactly.
        If converted back and forth, this mismatch has no effect.

    Raises:
        RuntimeError: When trying to cast :class:`torch.float32` to :class:`torch.int32` or :class:`torch.int64` as
            well as for trying to cast :class:`torch.float64` to :class:`torch.int64`. These conversions might lead to
            overflow errors since the floating point ``dtype`` cannot store consecutive integers over the whole range
            of the integer ``dtype``.
    """

    def __init__(self, dtype: torch.dtype) -> None:
141
        super().__init__()
142
143
        self.dtype = dtype

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


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

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

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

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

186

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

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

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

202
203
    """

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

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

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

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

223

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

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

243
244
    """

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

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

261
262
        self.interpolation = interpolation

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

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

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

277
278
279
280
281
282
283
284
285
286
287

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

    Args:
        size (sequence or int): Desired output size of the crop. If size is an
            int instead of sequence like (h, w), a square crop (size, size) is
vfdev's avatar
vfdev committed
296
            made. If provided a tuple or list of length 1, it will be interpreted as (size[0], size[0]).
297
298
299
    """

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

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

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

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

316

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

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

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

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

            - reflect: pads with reflection of image without repeating the last value on the edge

                For example, padding [1, 2, 3, 4] with 2 elements on both sides in reflect mode
342
                will result in [3, 2, 1, 2, 3, 4, 3, 2]
343
344
345
346

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

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

        self.padding = padding
        self.fill = fill
367
        self.padding_mode = padding_mode
368

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

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

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

383

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

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

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

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

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

402

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

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

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


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

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

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

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

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


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

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

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

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

             - reflect: pads with reflection of image (without repeating the last value on the edge)

                padding [1, 2, 3, 4] with 2 elements on both sides in reflect mode
                will result in [3, 2, 1, 2, 3, 4, 3, 2]

             - symmetric: pads with reflection of image (repeating the last value on the edge)

                padding [1, 2, 3, 4] with 2 elements on both sides in symmetric mode
                will result in [2, 1, 1, 2, 3, 4, 4, 3]

528
529
530
    """

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

        Args:
vfdev's avatar
vfdev committed
535
            img (PIL Image or Tensor): Image to be cropped.
536
537
538
539
540
            output_size (tuple): Expected output size of the crop.

        Returns:
            tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
        """
vfdev's avatar
vfdev committed
541
        w, h = F._get_image_size(img)
542
        th, tw = output_size
vfdev's avatar
vfdev committed
543
544
545
546
547
548

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

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

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

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

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

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

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

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

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

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

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

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

596

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

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

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

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

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

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

626

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

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

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

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

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

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

656

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

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

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

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

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

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

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

        fill = self.fill
        if isinstance(img, Tensor):
            if isinstance(fill, (int, float)):
                fill = [float(fill)] * F._get_image_num_channels(img)
            else:
                fill = [float(f) for f in fill]

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

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

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

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

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


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

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

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

775
776
    """

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

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

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

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

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

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

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

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

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

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

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

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

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

867
868
869
870
871
872
873
874
875
876
877

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

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

    Args:
         size (sequence or int): Desired output size of the crop. If size is an ``int``
            instead of sequence like (h, w), a square crop of size (size, size) is made.
vfdev's avatar
vfdev committed
892
            If provided a tuple or list of length 1, it will be interpreted as (size[0], size[0]).
893
894
895
896
897
898
899
900
901
902
903
904
905
906

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

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

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

923

vfdev's avatar
vfdev committed
924
925
926
927
928
929
class TenCrop(torch.nn.Module):
    """Crop the given image into four corners and the central crop plus the flipped version of
    these (horizontal flipping is used by default).
    The image can be a PIL Image or a Tensor, in which case it is expected
    to have [..., H, W] shape, where ... means an arbitrary number of leading
    dimensions
930
931
932
933
934
935
936
937
938

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

    Args:
        size (sequence or int): Desired output size of the crop. If size is an
            int instead of sequence like (h, w), a square crop (size, size) is
vfdev's avatar
vfdev committed
939
            made. If provided a tuple or list of length 1, it will be interpreted as (size[0], size[0]).
940
        vertical_flip (bool): Use vertical flipping instead of horizontal
941
942
943
944
945
946
947
948
949
950
951
952
953
954

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

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

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

972

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

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

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

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

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

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

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

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

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

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

1039

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

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

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

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

        # if value is 0 or (1., 1.) for brightness/contrast/saturation
        # or (0., 0.) for hue, do nothing
        if value[0] == value[1] == center:
            value = None
        return value
1085
1086

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

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

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

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

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

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

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

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

        return img
1139

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

1148

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

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

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

1177
1178
    """

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

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

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

        self.center = center
1203

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

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

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

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

        Returns:
1224
            PIL Image or Tensor: Rotated image.
1225
        """
1226
1227
1228
1229
1230
1231
        fill = self.fill
        if isinstance(img, Tensor):
            if isinstance(fill, (int, float)):
                fill = [float(fill)] * F._get_image_num_channels(img)
            else:
                fill = [float(f) for f in fill]
1232
        angle = self.get_params(self.degrees)
1233
1234

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

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

1248

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

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

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

1285
1286
    """

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

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

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

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

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

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

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

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

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

        return angle, translations, scale, shear

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

        Returns:
1379
            PIL Image or Tensor: Affine transformed image.
1380
        """
1381
1382
1383
1384
1385
1386
        fill = self.fill
        if isinstance(img, Tensor):
            if isinstance(fill, (int, float)):
                fill = [float(fill)] * F._get_image_num_channels(img)
            else:
                fill = [float(f) for f in fill]
1387
1388
1389
1390

        img_size = F._get_image_size(img)

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

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

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


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

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

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

    """

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

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

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

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

1445

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

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

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

    """

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

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

        Returns:
1473
            PIL Image or Tensor: Randomly grayscaled image.
1474
        """
1475
1476
1477
        num_output_channels = F._get_image_num_channels(img)
        if torch.rand(1) < self.p:
            return F.rgb_to_grayscale(img, num_output_channels=num_output_channels)
1478
        return img
1479
1480

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


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

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

1498
1499
    Returns:
        Erased Image.
1500

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

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

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

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

        Args:
vfdev's avatar
vfdev committed
1540
            img (Tensor): Tensor image to be erased.
1541
1542
1543
1544
1545
            scale (tuple or list): range of proportion of erased area against input image.
            ratio (tuple or list): range of aspect ratio of erased area.
            value (list, optional): erasing value. If None, it is interpreted as "random"
                (erasing each pixel with random values). If ``len(value)`` is 1, it is interpreted as a number,
                i.e. ``value[0]``.
1546
1547
1548
1549

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

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

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

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

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

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

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


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

    Args:
        kernel_size (int or sequence): Size of the Gaussian kernel.
        sigma (float or tuple of float (min, max)): Standard deviation to be used for
            creating kernel to perform blurring. If float, sigma is fixed. If it is tuple
            of float (min, max), sigma is chosen uniformly at random to lie in the
            given range.

    Returns:
        PIL Image or Tensor: Gaussian blurred version of the input image.

    """

    def __init__(self, kernel_size, sigma=(0.1, 2.0)):
        super().__init__()
        self.kernel_size = _setup_size(kernel_size, "Kernel size should be a tuple/list of two integers")
        for ks in self.kernel_size:
            if ks <= 0 or ks % 2 == 0:
                raise ValueError("Kernel size value should be an odd and positive number.")

        if isinstance(sigma, numbers.Number):
            if sigma <= 0:
                raise ValueError("If sigma is a single number, it must be positive.")
            sigma = (sigma, sigma)
        elif isinstance(sigma, Sequence) and len(sigma) == 2:
            if not 0. < sigma[0] <= sigma[1]:
                raise ValueError("sigma values should be positive and of the form (min, max).")
        else:
            raise ValueError("sigma should be a single number or a list/tuple with length 2.")

        self.sigma = sigma

    @staticmethod
    def get_params(sigma_min: float, sigma_max: float) -> float:
vfdev's avatar
vfdev committed
1644
        """Choose sigma for random gaussian blurring.
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657

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

        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


1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
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]