transforms.py 6.86 KB
Newer Older
1
from __future__ import division
soumith's avatar
soumith committed
2
3
4
import torch
import math
import random
5
from PIL import Image, ImageOps
6
import numpy as np
7
import numbers
Soumith Chintala's avatar
Soumith Chintala committed
8
import types
soumith's avatar
soumith committed
9
10

class Compose(object):
Adam Paszke's avatar
Adam Paszke committed
11
12
13
14
15
16
17
18
19
20
    """Composes several transforms together.

    Args:
        transforms (List[Transform]): list of transforms to compose.

    Example:
        >>> transforms.Compose([
        >>>     transforms.CenterCrop(10),
        >>>     transforms.ToTensor(),
        >>> ])
21
    """
soumith's avatar
soumith committed
22
23
24
25
26
27
28
29
30
31
    def __init__(self, transforms):
        self.transforms = transforms

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


class ToTensor(object):
Adam Paszke's avatar
Adam Paszke committed
32
33
34
    """Converts a PIL.Image (RGB) or numpy.ndarray (H x W x C) in the range
    [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0].
    """
soumith's avatar
soumith committed
35
    def __call__(self, pic):
36
37
38
39
40
41
        if isinstance(pic, np.ndarray):
            # handle numpy array
            img = torch.from_numpy(pic)
        else:
            # handle PIL Image
            img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes()))
42
            img = img.view(pic.size[1], pic.size[0], len(pic.mode))
Soumith Chintala's avatar
Soumith Chintala committed
43
            # put it from HWC to CHW format
44
            # yikes, this transpose takes 80% of the loading time/CPU
Soumith Chintala's avatar
Soumith Chintala committed
45
            img = img.transpose(0, 1).transpose(0, 2).contiguous()
46
47
        return img.float().div(255)

Adam Paszke's avatar
Adam Paszke committed
48

49
class ToPILImage(object):
Adam Paszke's avatar
Adam Paszke committed
50
    """Converts a torch.*Tensor of range [0, 1] and shape C x H x W
51
52
53
54
    or numpy ndarray of dtype=uint8, range[0, 255] and shape H x W x C
    to a PIL.Image of range [0, 255]
    """
    def __call__(self, pic):
55
56
57
        npimg = pic
        mode = None
        if not isinstance(npimg, np.ndarray):
58
59
            npimg = pic.mul(255).byte().numpy()
            npimg = np.transpose(npimg, (1,2,0))
60
61
62
63
64
65
66

        if npimg.shape[2] == 1:
            npimg = npimg[:, :, 0]
            mode = "L"

        return Image.fromarray(npimg, mode=mode)

soumith's avatar
soumith committed
67
68

class Normalize(object):
Adam Paszke's avatar
Adam Paszke committed
69
    """Given mean: (R, G, B) and std: (R, G, B),
70
71
72
    will normalize each channel of the torch.*Tensor, i.e.
    channel = (channel - mean) / std
    """
soumith's avatar
soumith committed
73
74
75
76
77
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, tensor):
78
        # TODO: make efficient
soumith's avatar
soumith committed
79
80
81
82
83
84
        for t, m, s in zip(tensor, self.mean, self.std):
            t.sub_(m).div_(s)
        return tensor


class Scale(object):
Adam Paszke's avatar
Adam Paszke committed
85
    """Rescales the input PIL.Image to the given 'size'.
86
87
88
89
90
91
    'size' will be the size of the smaller edge.
    For example, if height > width, then image will be
    rescaled to (size * height / width, size)
    size: size of the smaller edge
    interpolation: Default: PIL.Image.BILINEAR
    """
soumith's avatar
soumith committed
92
93
94
95
96
97
98
99
100
    def __init__(self, size, interpolation=Image.BILINEAR):
        self.size = size
        self.interpolation = interpolation

    def __call__(self, img):
        w, h = img.size
        if (w <= h and w == self.size) or (h <= w and h == self.size):
            return img
        if w < h:
101
102
103
            ow = self.size
            oh = int(self.size * h / w)
            return img.resize((ow, oh), self.interpolation)
soumith's avatar
soumith committed
104
        else:
105
106
107
            oh = self.size
            ow = int(self.size * w / h)
            return img.resize((ow, oh), self.interpolation)
soumith's avatar
soumith committed
108
109
110


class CenterCrop(object):
111
112
113
114
    """Crops the given PIL.Image at the center to have a region of
    the given size. size can be a tuple (target_height, target_width)
    or an integer, in which case the target will be of a square shape (size, size)
    """
soumith's avatar
soumith committed
115
    def __init__(self, size):
116
117
118
119
        if isinstance(size, numbers.Number):
            self.size = (int(size), int(size))
        else:
            self.size = size
soumith's avatar
soumith committed
120
121
122

    def __call__(self, img):
        w, h = img.size
123
        th, tw = self.size
124
125
        x1 = int(round((w - tw) / 2.))
        y1 = int(round((h - th) / 2.))
126
        return img.crop((x1, y1, x1 + tw, y1 + th))
soumith's avatar
soumith committed
127
128


129
130
131
132
133
134
135
136
137
138
139
class Pad(object):
    """Pads the given PIL.Image on all sides with the given "pad" value"""
    def __init__(self, padding, fill=0):
        assert isinstance(padding, numbers.Number)
        assert isinstance(fill, numbers.Number)
        self.padding = padding
        self.fill = fill

    def __call__(self, img):
        return ImageOps.expand(img, border=self.padding, fill=self.fill)

Soumith Chintala's avatar
Soumith Chintala committed
140
class Lambda(object):
Adam Paszke's avatar
Adam Paszke committed
141
    """Applies a lambda as a transform."""
Soumith Chintala's avatar
Soumith Chintala committed
142
143
144
145
146
147
148
    def __init__(self, lambd):
        assert type(lambd) is types.LambdaType
        self.lambd = lambd

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

149

soumith's avatar
soumith committed
150
class RandomCrop(object):
151
152
153
154
    """Crops the given PIL.Image at a random location to have a region of
    the given size. size can be a tuple (target_height, target_width)
    or an integer, in which case the target will be of a square shape (size, size)
    """
soumith's avatar
soumith committed
155
    def __init__(self, size, padding=0):
156
157
158
159
        if isinstance(size, numbers.Number):
            self.size = (int(size), int(size))
        else:
            self.size = size
soumith's avatar
soumith committed
160
161
162
163
        self.padding = padding

    def __call__(self, img):
        if self.padding > 0:
164
            img = ImageOps.expand(img, border=self.padding, fill=0)
soumith's avatar
soumith committed
165
166

        w, h = img.size
167
168
        th, tw = self.size
        if w == tw and h == th:
soumith's avatar
soumith committed
169
170
            return img

171
172
173
        x1 = random.randint(0, w - tw)
        y1 = random.randint(0, h - th)
        return img.crop((x1, y1, x1 + tw, y1 + th))
soumith's avatar
soumith committed
174
175
176


class RandomHorizontalFlip(object):
177
178
    """Randomly horizontally flips the given PIL.Image with a probability of 0.5
    """
soumith's avatar
soumith committed
179
180
181
182
183
184
185
    def __call__(self, img):
        if random.random() < 0.5:
            return img.transpose(Image.FLIP_LEFT_RIGHT)
        return img


class RandomSizedCrop(object):
186
187
188
189
190
191
    """Random crop the given PIL.Image to a random size of (0.08 to 1.0) of the original size
    and and a random aspect ratio of 3/4 to 4/3 of the original aspect ratio
    This is popularly used to train the Inception networks
    size: size of the smaller edge
    interpolation: Default: PIL.Image.BILINEAR
    """
soumith's avatar
soumith committed
192
193
194
195
196
197
198
199
    def __init__(self, size, interpolation=Image.BILINEAR):
        self.size = size
        self.interpolation = interpolation

    def __call__(self, img):
        for attempt in range(10):
            area = img.size[0] * img.size[1]
            target_area = random.uniform(0.08, 1.0) * area
200
            aspect_ratio = random.uniform(3. / 4, 4. / 3)
soumith's avatar
soumith committed
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220

            w = int(round(math.sqrt(target_area * aspect_ratio)))
            h = int(round(math.sqrt(target_area / aspect_ratio)))

            if random.random() < 0.5:
                w, h = h, w

            if w <= img.size[0] and h <= img.size[1]:
                x1 = random.randint(0, img.size[0] - w)
                y1 = random.randint(0, img.size[1] - h)

                img = img.crop((x1, y1, x1 + w, y1 + h))
                assert(img.size == (w, h))

                return img.resize((self.size, self.size), self.interpolation)

        # Fallback
        scale = Scale(self.size, interpolation=self.interpolation)
        crop = CenterCrop(self.size)
        return crop(scale(img))