-
Notifications
You must be signed in to change notification settings - Fork 0
/
single.py
103 lines (92 loc) · 4.79 KB
/
single.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
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image, ImageEnhance, ImageOps
import numpy as np
import random
class Rand_Augment():
def __init__(self, Numbers=None, max_Magnitude=None):
self.transforms = ['autocontrast', 'equalize', 'rotate', 'solarize', 'color', 'posterize',
'contrast', 'brightness', 'sharpness', 'shearX', 'shearY', 'translateX', 'translateY']
if Numbers is None:
self.Numbers = len(self.transforms) // 2
else:
self.Numbers = Numbers
if max_Magnitude is None:
self.max_Magnitude = 10
else:
self.max_Magnitude = max_Magnitude
fillcolor = 128
self.ranges = {
# these Magnitude range , you must test it yourself , see what will happen after these operation ,
# it is no need to obey the value in autoaugment.py
"shearX": np.linspace(0, 0.3, 10),
"shearY": np.linspace(0, 0.3, 10),
"translateX": np.linspace(0, 0.2, 10),
"translateY": np.linspace(0, 0.2, 10),
"rotate": np.linspace(0, 360, 10),
"color": np.linspace(0.0, 0.9, 10),
"posterize": np.round(np.linspace(8, 4, 10), 0).astype(np.int),
"solarize": np.linspace(256, 231, 10),
"contrast": np.linspace(0.0, 0.5, 10),
"sharpness": np.linspace(0.0, 0.9, 10),
"brightness": np.linspace(0.0, 0.3, 10),
"autocontrast": [0] * 10,
"equalize": [0] * 10,
"invert": [0] * 10
}
self.func = {
"shearX": lambda img, magnitude: img.transform(
img.size, Image.AFFINE, (1, magnitude * random.choice([-1, 1]), 0, 0, 1, 0),
Image.BICUBIC, fill=fillcolor),
"shearY": lambda img, magnitude: img.transform(
img.size, Image.AFFINE, (1, 0, 0, magnitude * random.choice([-1, 1]), 1, 0),
Image.BICUBIC, fill=fillcolor),
"translateX": lambda img, magnitude: img.transform(
img.size, Image.AFFINE, (1, 0, magnitude * img.size[0] * random.choice([-1, 1]), 0, 1, 0),
fill=fillcolor),
"translateY": lambda img, magnitude: img.transform(
img.size, Image.AFFINE, (1, 0, 0, 0, 1, magnitude * img.size[1] * random.choice([-1, 1])),
fill=fillcolor),
"rotate": lambda img, magnitude: self.rotate_with_fill(img, magnitude),
# "rotate": lambda img, magnitude: img.rotate(magnitude * random.choice([-1, 1])),
"color": lambda img, magnitude: ImageEnhance.Color(img).enhance(1 + magnitude * random.choice([-1, 1])),
"posterize": lambda img, magnitude: ImageOps.posterize(img, magnitude),
"solarize": lambda img, magnitude: ImageOps.solarize(img, magnitude),
"contrast": lambda img, magnitude: ImageEnhance.Contrast(img).enhance(
1 + magnitude * random.choice([-1, 1])),
"sharpness": lambda img, magnitude: ImageEnhance.Sharpness(img).enhance(
1 + magnitude * random.choice([-1, 1])),
"brightness": lambda img, magnitude: ImageEnhance.Brightness(img).enhance(
1 + magnitude * random.choice([-1, 1])),
"autocontrast": lambda img, magnitude: ImageOps.autocontrast(img),
"equalize": lambda img, magnitude: img,
"invert": lambda img, magnitude: ImageOps.invert(img)
}
def rand_augment(self):
M = np.random.randint(0, self.max_Magnitude, self.Numbers)
sampled_ops = np.random.choice(self.transforms, self.Numbers)
return [(op, Magnitude) for (op, Magnitude) in zip(sampled_ops, M)]
def __call__(self, image):
operations = self.rand_augment()
for (op_name, M) in operations:
operation = self.func[op_name]
mag = self.ranges[op_name][M]
image = operation(image, mag)
return image
def rotate_with_fill(self, img, magnitude):
# I don't know why rotate must change to RGBA , it is copy from Autoaugment - pytorch
rot = img.convert("RGBA").rotate(magnitude)
return Image.composite(rot, Image.new("RGBA", rot.size, (128,) * 4), rot).convert(img.mode)
def test_single_operation(self, image, op_name, M=-1):
operation = self.func[op_name]
mag = self.ranges[op_name][M]
image = operation(image, mag)
return image
if __name__ == '__main__':
img_augment = Rand_Augment()
img_origal = Image.open(r'./example-test-pics/dog4.png')
img_final = img_augment(img_origal)
plt.imshow(img_final)
plt.savefig('dog-after-4.png')
plt.show()