transforms.py 7.04 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):
32
    """Converts a PIL.Image or numpy.ndarray (H x W x C) in the range
Adam Paszke's avatar
Adam Paszke committed
33
34
    [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
43
44
45
46
47
            # PIL image mode: 1, L, P, I, F, RGB, YCbCr, RGBA, CMYK
            if pic.mode == 'YCbCr':
                nchannel = 3
            else:
                nchannel = len(pic.mode)
            img = img.view(pic.size[1], pic.size[0], nchannel)
Soumith Chintala's avatar
Soumith Chintala committed
48
            # put it from HWC to CHW format
49
            # yikes, this transpose takes 80% of the loading time/CPU
Soumith Chintala's avatar
Soumith Chintala committed
50
            img = img.transpose(0, 1).transpose(0, 2).contiguous()
51
52
        return img.float().div(255)

Adam Paszke's avatar
Adam Paszke committed
53

54
class ToPILImage(object):
Adam Paszke's avatar
Adam Paszke committed
55
    """Converts a torch.*Tensor of range [0, 1] and shape C x H x W
56
57
58
59
    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):
60
61
62
        npimg = pic
        mode = None
        if not isinstance(npimg, np.ndarray):
63
64
            npimg = pic.mul(255).byte().numpy()
            npimg = np.transpose(npimg, (1,2,0))
65
66
67
68
69
70
71

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

        return Image.fromarray(npimg, mode=mode)

soumith's avatar
soumith committed
72
73

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

    def __call__(self, tensor):
83
        # TODO: make efficient
soumith's avatar
soumith committed
84
85
86
87
88
89
        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
90
    """Rescales the input PIL.Image to the given 'size'.
91
92
93
94
95
96
    '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
97
98
99
100
101
102
103
104
105
    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:
106
107
108
            ow = self.size
            oh = int(self.size * h / w)
            return img.resize((ow, oh), self.interpolation)
soumith's avatar
soumith committed
109
        else:
110
111
112
            oh = self.size
            ow = int(self.size * w / h)
            return img.resize((ow, oh), self.interpolation)
soumith's avatar
soumith committed
113
114
115


class CenterCrop(object):
116
117
118
119
    """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
120
    def __init__(self, size):
121
122
123
124
        if isinstance(size, numbers.Number):
            self.size = (int(size), int(size))
        else:
            self.size = size
soumith's avatar
soumith committed
125
126
127

    def __call__(self, img):
        w, h = img.size
128
        th, tw = self.size
129
130
        x1 = int(round((w - tw) / 2.))
        y1 = int(round((h - th) / 2.))
131
        return img.crop((x1, y1, x1 + tw, y1 + th))
soumith's avatar
soumith committed
132
133


134
135
136
137
138
139
140
141
142
143
144
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
145
class Lambda(object):
Adam Paszke's avatar
Adam Paszke committed
146
    """Applies a lambda as a transform."""
Soumith Chintala's avatar
Soumith Chintala committed
147
148
149
150
151
152
153
    def __init__(self, lambd):
        assert type(lambd) is types.LambdaType
        self.lambd = lambd

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

154

soumith's avatar
soumith committed
155
class RandomCrop(object):
156
157
158
159
    """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
160
    def __init__(self, size, padding=0):
161
162
163
164
        if isinstance(size, numbers.Number):
            self.size = (int(size), int(size))
        else:
            self.size = size
soumith's avatar
soumith committed
165
166
167
168
        self.padding = padding

    def __call__(self, img):
        if self.padding > 0:
169
            img = ImageOps.expand(img, border=self.padding, fill=0)
soumith's avatar
soumith committed
170
171

        w, h = img.size
172
173
        th, tw = self.size
        if w == tw and h == th:
soumith's avatar
soumith committed
174
175
            return img

176
177
178
        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
179
180
181


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


class RandomSizedCrop(object):
191
192
193
194
195
196
    """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
197
198
199
200
201
202
203
204
    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
205
            aspect_ratio = random.uniform(3. / 4, 4. / 3)
soumith's avatar
soumith committed
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225

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