functional_pil.py 12.5 KB
Newer Older
1
import numbers
2
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
3

vfdev's avatar
vfdev committed
4
import numpy as np
5
import torch
6
from PIL import Image, ImageOps, ImageEnhance
7
from typing_extensions import Literal
vfdev's avatar
vfdev committed
8

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


@torch.jit.unused
vfdev's avatar
vfdev committed
16
def _is_pil_image(img: Any) -> bool:
17
18
19
20
21
22
    if accimage is not None:
        return isinstance(img, (Image.Image, accimage.Image))
    else:
        return isinstance(img, Image.Image)


23
24
25
26
27
28
29
30
31
@torch.jit.unused
def get_dimensions(img: Any) -> List[int]:
    if _is_pil_image(img):
        channels = len(img.getbands())
        width, height = img.size
        return [channels, height, width]
    raise TypeError(f"Unexpected type {type(img)}")


vfdev's avatar
vfdev committed
32
@torch.jit.unused
33
def get_image_size(img: Any) -> List[int]:
vfdev's avatar
vfdev committed
34
    if _is_pil_image(img):
35
        return list(img.size)
36
    raise TypeError(f"Unexpected type {type(img)}")
vfdev's avatar
vfdev committed
37
38


39
@torch.jit.unused
40
def get_image_num_channels(img: Any) -> int:
41
    if _is_pil_image(img):
42
        return len(img.getbands())
43
    raise TypeError(f"Unexpected type {type(img)}")
44
45


46
@torch.jit.unused
47
def hflip(img: Image.Image) -> Image.Image:
48
    if not _is_pil_image(img):
49
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
50
51
52
53
54

    return img.transpose(Image.FLIP_LEFT_RIGHT)


@torch.jit.unused
55
def vflip(img: Image.Image) -> Image.Image:
56
    if not _is_pil_image(img):
57
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
58
59

    return img.transpose(Image.FLIP_TOP_BOTTOM)
60
61
62


@torch.jit.unused
63
def adjust_brightness(img: Image.Image, brightness_factor: float) -> Image.Image:
64
    if not _is_pil_image(img):
65
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
66
67
68
69
70
71
72

    enhancer = ImageEnhance.Brightness(img)
    img = enhancer.enhance(brightness_factor)
    return img


@torch.jit.unused
73
def adjust_contrast(img: Image.Image, contrast_factor: float) -> Image.Image:
74
    if not _is_pil_image(img):
75
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
76
77
78
79
80
81
82

    enhancer = ImageEnhance.Contrast(img)
    img = enhancer.enhance(contrast_factor)
    return img


@torch.jit.unused
83
def adjust_saturation(img: Image.Image, saturation_factor: float) -> Image.Image:
84
    if not _is_pil_image(img):
85
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
86
87
88
89
90
91
92

    enhancer = ImageEnhance.Color(img)
    img = enhancer.enhance(saturation_factor)
    return img


@torch.jit.unused
93
def adjust_hue(img: Image.Image, hue_factor: float) -> Image.Image:
94
    if not (-0.5 <= hue_factor <= 0.5):
95
        raise ValueError(f"hue_factor ({hue_factor}) is not in [-0.5, 0.5].")
96
97

    if not _is_pil_image(img):
98
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
99
100

    input_mode = img.mode
101
    if input_mode in {"L", "1", "I", "F"}:
102
103
        return img

104
    h, s, v = img.convert("HSV").split()
105
106
107

    np_h = np.array(h, dtype=np.uint8)
    # uint8 addition take cares of rotation across boundaries
108
    with np.errstate(over="ignore"):
109
        np_h += np.uint8(hue_factor * 255)
110
    h = Image.fromarray(np_h, "L")
111

112
    img = Image.merge("HSV", (h, s, v)).convert(input_mode)
113
    return img
114
115


116
@torch.jit.unused
117
118
119
120
121
122
def adjust_gamma(
    img: Image.Image,
    gamma: float,
    gain: float = 1.0,
) -> Image.Image:

123
    if not _is_pil_image(img):
124
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
125
126

    if gamma < 0:
127
        raise ValueError("Gamma should be a non-negative real number")
128
129

    input_mode = img.mode
130
    img = img.convert("RGB")
131
    gamma_map = [int((255 + 1 - 1e-3) * gain * pow(ele / 255.0, gamma)) for ele in range(256)] * 3
132
133
134
135
136
137
    img = img.point(gamma_map)  # use PIL's point-function to accelerate this part

    img = img.convert(input_mode)
    return img


138
@torch.jit.unused
139
140
141
142
def pad(
    img: Image.Image,
    padding: Union[int, List[int], Tuple[int, ...]],
    fill: Optional[Union[float, List[float], Tuple[float, ...]]] = 0,
143
    padding_mode: Literal["constant", "edge", "reflect", "symmetric"] = "constant",
144
145
) -> Image.Image:

146
    if not _is_pil_image(img):
147
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
148
149
150
151
152
153
154
155
156
157
158
159

    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 not isinstance(padding_mode, str):
        raise TypeError("Got inappropriate padding_mode arg")

    if isinstance(padding, list):
        padding = tuple(padding)

    if isinstance(padding, tuple) and len(padding) not in [1, 2, 4]:
160
        raise ValueError(f"Padding must be an int or a 1, 2, or 4 element tuple, not a {len(padding)} element tuple")
161
162
163
164
165
166
167
168
169

    if isinstance(padding, tuple) and len(padding) == 1:
        # Compatibility with `functional_tensor.pad`
        padding = padding[0]

    if padding_mode not in ["constant", "edge", "reflect", "symmetric"]:
        raise ValueError("Padding mode should be either constant, edge, reflect or symmetric")

    if padding_mode == "constant":
170
        opts = _parse_fill(fill, img, name="fill")
171
172
        if img.mode == "P":
            palette = img.getpalette()
173
            image = ImageOps.expand(img, border=padding, **opts)
174
175
176
            image.putpalette(palette)
            return image

177
        return ImageOps.expand(img, border=padding, **opts)
178
179
180
181
182
183
184
185
186
187
188
189
    else:
        if isinstance(padding, int):
            pad_left = pad_right = pad_top = pad_bottom = padding
        if isinstance(padding, tuple) and len(padding) == 2:
            pad_left = pad_right = padding[0]
            pad_top = pad_bottom = padding[1]
        if isinstance(padding, tuple) and len(padding) == 4:
            pad_left = padding[0]
            pad_top = padding[1]
            pad_right = padding[2]
            pad_bottom = padding[3]

190
191
192
193
194
195
196
197
198
        p = [pad_left, pad_top, pad_right, pad_bottom]
        cropping = -np.minimum(p, 0)

        if cropping.any():
            crop_left, crop_top, crop_right, crop_bottom = cropping
            img = img.crop((crop_left, crop_top, img.width - crop_right, img.height - crop_bottom))

        pad_left, pad_top, pad_right, pad_bottom = np.maximum(p, 0)

199
        if img.mode == "P":
200
201
            palette = img.getpalette()
            img = np.asarray(img)
202
            img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right)), mode=padding_mode)
203
204
205
206
207
208
209
210
211
212
213
214
215
            img = Image.fromarray(img)
            img.putpalette(palette)
            return img

        img = np.asarray(img)
        # RGB image
        if len(img.shape) == 3:
            img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right), (0, 0)), padding_mode)
        # Grayscale image
        if len(img.shape) == 2:
            img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right)), padding_mode)

        return Image.fromarray(img)
vfdev's avatar
vfdev committed
216
217
218


@torch.jit.unused
219
220
221
222
223
224
225
226
def crop(
    img: Image.Image,
    top: int,
    left: int,
    height: int,
    width: int,
) -> Image.Image:

vfdev's avatar
vfdev committed
227
    if not _is_pil_image(img):
228
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
vfdev's avatar
vfdev committed
229
230

    return img.crop((left, top, left + width, top + height))
vfdev's avatar
vfdev committed
231
232
233


@torch.jit.unused
234
235
236
237
238
239
240
def resize(
    img: Image.Image,
    size: Union[Sequence[int], int],
    interpolation: int = Image.BILINEAR,
    max_size: Optional[int] = None,
) -> Image.Image:

vfdev's avatar
vfdev committed
241
    if not _is_pil_image(img):
242
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
vfdev's avatar
vfdev committed
243
    if not (isinstance(size, int) or (isinstance(size, Sequence) and len(size) in (1, 2))):
244
        raise TypeError(f"Got inappropriate size arg: {size}")
vfdev's avatar
vfdev committed
245

246
247
248
    if isinstance(size, Sequence) and len(size) == 1:
        size = size[0]
    if isinstance(size, int):
vfdev's avatar
vfdev committed
249
        w, h = img.size
250
251
252
253
254
255
256
257
258
259
260
261
262
263

        short, long = (w, h) if w <= h else (h, w)
        new_short, new_long = size, int(size * long / short)

        if max_size is not None:
            if max_size <= size:
                raise ValueError(
                    f"max_size = {max_size} must be strictly greater than the requested "
                    f"size for the smaller edge size = {size}"
                )
            if new_long > max_size:
                new_short, new_long = int(max_size * new_short / new_long), max_size

        new_w, new_h = (new_short, new_long) if w <= h else (new_long, new_short)
264
265
266
267
268

        if (w, h) == (new_w, new_h):
            return img
        else:
            return img.resize((new_w, new_h), interpolation)
vfdev's avatar
vfdev committed
269
    else:
270
271
272
273
274
        if max_size is not None:
            raise ValueError(
                "max_size should only be passed if size specifies the length of the smaller edge, "
                "i.e. size should be an int or a sequence of length 1 in torchscript mode."
            )
vfdev's avatar
vfdev committed
275
        return img.resize(size[::-1], interpolation)
vfdev's avatar
vfdev committed
276
277
278


@torch.jit.unused
279
280
281
282
283
284
def _parse_fill(
    fill: Optional[Union[float, List[float], Tuple[float, ...]]],
    img: Image.Image,
    name: str = "fillcolor",
) -> Dict[str, Optional[Union[float, List[float], Tuple[float, ...]]]]:

285
    # Process fill color for affine transforms
vfdev's avatar
vfdev committed
286
287
288
289
290
    num_bands = len(img.getbands())
    if fill is None:
        fill = 0
    if isinstance(fill, (int, float)) and num_bands > 1:
        fill = tuple([fill] * num_bands)
291
292
    if isinstance(fill, (list, tuple)):
        if len(fill) != num_bands:
293
            msg = "The number of elements in 'fill' does not match the number of bands of the image ({} != {})"
294
295
296
            raise ValueError(msg.format(len(fill), num_bands))

        fill = tuple(fill)
vfdev's avatar
vfdev committed
297

298
    return {name: fill}
vfdev's avatar
vfdev committed
299
300
301


@torch.jit.unused
302
303
304
305
306
307
308
def affine(
    img: Image.Image,
    matrix: List[float],
    interpolation: int = Image.NEAREST,
    fill: Optional[Union[float, List[float], Tuple[float, ...]]] = 0,
) -> Image.Image:

vfdev's avatar
vfdev committed
309
    if not _is_pil_image(img):
310
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
vfdev's avatar
vfdev committed
311
312

    output_size = img.size
313
    opts = _parse_fill(fill, img)
314
    return img.transform(output_size, Image.AFFINE, matrix, interpolation, **opts)
vfdev's avatar
vfdev committed
315
316
317


@torch.jit.unused
318
319
320
321
322
323
324
325
326
def rotate(
    img: Image.Image,
    angle: float,
    interpolation: int = Image.NEAREST,
    expand: bool = False,
    center: Optional[Tuple[int, int]] = None,
    fill: Optional[Union[float, List[float], Tuple[float, ...]]] = 0,
) -> Image.Image:

vfdev's avatar
vfdev committed
327
    if not _is_pil_image(img):
328
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
vfdev's avatar
vfdev committed
329

330
    opts = _parse_fill(fill, img)
331
    return img.rotate(angle, interpolation, expand, center, **opts)
332
333
334


@torch.jit.unused
335
336
337
338
339
340
341
def perspective(
    img: Image.Image,
    perspective_coeffs: float,
    interpolation: int = Image.BICUBIC,
    fill: Optional[Union[float, List[float], Tuple[float, ...]]] = 0,
) -> Image.Image:

342
    if not _is_pil_image(img):
343
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
344

345
    opts = _parse_fill(fill, img)
346
347

    return img.transform(img.size, Image.PERSPECTIVE, perspective_coeffs, interpolation, **opts)
348
349
350


@torch.jit.unused
351
def to_grayscale(img: Image.Image, num_output_channels: int) -> Image.Image:
352
    if not _is_pil_image(img):
353
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
354
355

    if num_output_channels == 1:
356
        img = img.convert("L")
357
    elif num_output_channels == 3:
358
        img = img.convert("L")
359
360
        np_img = np.array(img, dtype=np.uint8)
        np_img = np.dstack([np_img, np_img, np_img])
361
        img = Image.fromarray(np_img, "RGB")
362
    else:
363
        raise ValueError("num_output_channels should be either 1 or 3")
364
365

    return img
366
367
368


@torch.jit.unused
369
def invert(img: Image.Image) -> Image.Image:
370
    if not _is_pil_image(img):
371
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
372
373
374
375
    return ImageOps.invert(img)


@torch.jit.unused
376
def posterize(img: Image.Image, bits: int) -> Image.Image:
377
    if not _is_pil_image(img):
378
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
379
380
381
382
    return ImageOps.posterize(img, bits)


@torch.jit.unused
383
def solarize(img: Image.Image, threshold: int) -> Image.Image:
384
    if not _is_pil_image(img):
385
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
386
387
388
389
    return ImageOps.solarize(img, threshold)


@torch.jit.unused
390
def adjust_sharpness(img: Image.Image, sharpness_factor: float) -> Image.Image:
391
    if not _is_pil_image(img):
392
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
393
394
395
396
397
398
399

    enhancer = ImageEnhance.Sharpness(img)
    img = enhancer.enhance(sharpness_factor)
    return img


@torch.jit.unused
400
def autocontrast(img: Image.Image) -> Image.Image:
401
    if not _is_pil_image(img):
402
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
403
404
405
406
    return ImageOps.autocontrast(img)


@torch.jit.unused
407
def equalize(img: Image.Image) -> Image.Image:
408
    if not _is_pil_image(img):
409
        raise TypeError(f"img should be PIL Image. Got {type(img)}")
410
    return ImageOps.equalize(img)