-
-
Notifications
You must be signed in to change notification settings - Fork 297
/
transforms.py
155 lines (119 loc) · 3.91 KB
/
transforms.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
import numpy as np
try:
import Image
except ImportError:
from PIL import Image
from dezero.utils import pair
class Compose:
"""Compose several transforms.
Args:
transforms (list): list of transforms
"""
def __init__(self, transforms=[]):
self.transforms = transforms
def __call__(self, img):
if not self.transforms:
return img
for t in self.transforms:
img = t(img)
return img
# =============================================================================
# Transforms for PIL Image
# =============================================================================
class Convert:
def __init__(self, mode='RGB'):
self.mode = mode
def __call__(self, img):
if self.mode == 'BGR':
img = img.convert('RGB')
r, g, b = img.split()
img = Image.merge('RGB', (b, g, r))
return img
else:
return img.convert(self.mode)
class Resize:
"""Resize the input PIL image to the given size.
Args:
size (int or (int, int)): Desired output size
mode (int): Desired interpolation.
"""
def __init__(self, size, mode=Image.BILINEAR):
self.size = pair(size)
self.mode = mode
def __call__(self, img):
return img.resize(self.size, self.mode)
class CenterCrop:
"""Resize the input PIL image to the given size.
Args:
size (int or (int, int)): Desired output size.
mode (int): Desired interpolation.
"""
def __init__(self, size):
self.size = pair(size)
def __call__(self, img):
W, H = img.size
OW, OH = self.size
left = (W - OW) // 2
right = W - ((W - OW) // 2 + (W - OW) % 2)
up = (H - OH) // 2
bottom = H - ((H - OH) // 2 + (H - OH) % 2)
return img.crop((left, up, right, bottom))
class ToArray:
"""Convert PIL Image to NumPy array."""
def __init__(self, dtype=np.float32):
self.dtype = dtype
def __call__(self, img):
if isinstance(img, np.ndarray):
return img
if isinstance(img, Image.Image):
img = np.asarray(img)
img = img.transpose(2, 0, 1)
img = img.astype(self.dtype)
return img
else:
raise TypeError
class ToPIL:
"""Convert NumPy array to PIL Image."""
def __call__(self, array):
data = array.transpose(1, 2, 0)
return Image.fromarray(data)
class RandomHorizontalFlip:
pass
# =============================================================================
# Transforms for NumPy ndarray
# =============================================================================
class Normalize:
"""Normalize a NumPy array with mean and standard deviation.
Args:
mean (float or sequence): mean for all values or sequence of means for
each channel.
std (float or sequence):
"""
def __init__(self, mean=0, std=1):
self.mean = mean
self.std = std
def __call__(self, array):
mean, std = self.mean, self.std
if not np.isscalar(mean):
mshape = [1] * array.ndim
mshape[0] = len(array) if len(self.mean) == 1 else len(self.mean)
mean = np.array(self.mean, dtype=array.dtype).reshape(*mshape)
if not np.isscalar(std):
rshape = [1] * array.ndim
rshape[0] = len(array) if len(self.std) == 1 else len(self.std)
std = np.array(self.std, dtype=array.dtype).reshape(*rshape)
return (array - mean) / std
class Flatten:
"""Flatten a NumPy array.
"""
def __call__(self, array):
return array.flatten()
class AsType:
def __init__(self, dtype=np.float32):
self.dtype = dtype
def __call__(self, array):
return array.astype(self.dtype)
ToFloat = AsType
class ToInt(AsType):
def __init__(self, dtype=int):
self.dtype = dtype