import random
import numpy as np
import scipy
from scipy import ndimage
from PIL import Image, ImageEnhance, ImageOps
#See: https://github.com/4uiiurz1/pytorch-auto-augment
class AutoAugment(object):
def __init__(self):
self.policies = [
['Invert', 0.1, 7, 'Contrast', 0.2, 6],
['Rotate', 0.7, 2, 'TranslateX', 0.3, 9],
['Sharpness', 0.8, 1, 'Sharpness', 0.9, 3],
['ShearY', 0.5, 8, 'TranslateY', 0.7, 9],
['AutoContrast', 0.5, 8, 'Equalize', 0.9, 2],
['ShearY', 0.2, 7, 'Posterize', 0.3, 7],
['Color', 0.4, 3, 'Brightness', 0.6, 7],
['Sharpness', 0.3, 9, 'Brightness', 0.7, 9],
['Equalize', 0.6, 5, 'Equalize', 0.5, 1],
['Contrast', 0.6, 7, 'Sharpness', 0.6, 5],
['Color', 0.7, 7, 'TranslateX', 0.5, 8],
['Equalize', 0.3, 7, 'AutoContrast', 0.4, 8],
['TranslateY', 0.4, 3, 'Sharpness', 0.2, 6],
['Brightness', 0.9, 6, 'Color', 0.2, 8],
['Solarize', 0.5, 2, 'Invert', 0, 0.3],
['Equalize', 0.2, 0, 'AutoContrast', 0.6, 0],
['Equalize', 0.2, 8, 'Equalize', 0.6, 4],
['Color', 0.9, 9, 'Equalize', 0.6, 6],
['AutoContrast', 0.8, 4, 'Solarize', 0.2, 8],
['Brightness', 0.1, 3, 'Color', 0.7, 0],
['Solarize', 0.4, 5, 'AutoContrast', 0.9, 3],
['TranslateY', 0.9, 9, 'TranslateY', 0.7, 9],
['AutoContrast', 0.9, 2, 'Solarize', 0.8, 3],
['Equalize', 0.8, 8, 'Invert', 0.1, 3],
['TranslateY', 0.7, 9, 'AutoContrast', 0.9, 1],
]
def __call__(self, img):
img = apply_policy(img, self.policies[random.randrange(len(self.policies))])
return img
operations = {
'ShearX': lambda img, magnitude: shear_x(img, magnitude),
'ShearY': lambda img, magnitude: shear_y(img, magnitude),
'TranslateX': lambda img, magnitude: translate_x(img, magnitude),
'TranslateY': lambda img, magnitude: translate_y(img, magnitude),
'Rotate': lambda img, magnitude: rotate(img, magnitude),
'AutoContrast': lambda img, magnitude: auto_contrast(img, magnitude),
'Invert': lambda img, magnitude: invert(img, magnitude),
'Equalize': lambda img, magnitude: equalize(img, magnitude),
'Solarize': lambda img, magnitude: solarize(img, magnitude),
'Posterize': lambda img, magnitude: posterize(img, magnitude),
'Contrast': lambda img, magnitude: contrast(img, magnitude),
'Color': lambda img, magnitude: color(img, magnitude),
'Brightness': lambda img, magnitude: brightness(img, magnitude),
'Sharpness': lambda img, magnitude: sharpness(img, magnitude),
'Cutout': lambda img, magnitude: cutout(img, magnitude),
}
def apply_policy(img, policy):
if random.random() < policy[1]:
img = operations[policy[0]](img, policy[2])
if random.random() < policy[4]:
img = operations[policy[3]](img, policy[5])
return img
def transform_matrix_offset_center(matrix, x, y):
o_x = float(x) / 2 + 0.5
o_y = float(y) / 2 + 0.5
offset_matrix = np.array([[1, 0, o_x], [0, 1, o_y], [0, 0, 1]])
reset_matrix = np.array([[1, 0, -o_x], [0, 1, -o_y], [0, 0, 1]])
transform_matrix = offset_matrix @ matrix @ reset_matrix
return transform_matrix
def shear_x(img, magnitude):
img = np.array(img)
magnitudes = np.linspace(-0.3, 0.3, 11)
transform_matrix = np.array([[1, random.uniform(magnitudes[magnitude], magnitudes[magnitude+1]), 0],
[0, 1, 0],
[0, 0, 1]])
transform_matrix = transform_matrix_offset_center(transform_matrix, img.shape[0], img.shape[1])
affine_matrix = transform_matrix[:2, :2]
offset = transform_matrix[:2, 2]
img = np.stack([ndimage.interpolation.affine_transform(
img[:, :, c],
affine_matrix,
offset) for c in range(img.shape[2])], axis=2)
img = Image.fromarray(img)
return img
def shear_y(img, magnitude):
img = np.array(img)
magnitudes = np.linspace(-0.3, 0.3, 11)
transform_matrix = np.array([[1, 0, 0],
[random.uniform(magnitudes[magnitude], magnitudes[magnitude+1]), 1, 0],
[0, 0, 1]])
transform_matrix = transform_matrix_offset_center(transform_matrix, img.shape[0], img.shape[1])
affine_matrix = transform_matrix[:2, :2]
offset = transform_matrix[:2, 2]
img = np.stack([ndimage.interpolation.affine_transform(
img[:, :, c],
affine_matrix,
offset) for c in range(img.shape[2])], axis=2)
img = Image.fromarray(img)
return img
def translate_x(img, magnitude):
img = np.array(img)
magnitudes = np.linspace(-150/331, 150/331, 11)
transform_matrix = np.array([[1, 0, 0],
[0, 1, img.shape[1]*random.uniform(magnitudes[magnitude], magnitudes[magnitude+1])],
[0, 0, 1]])
transform_matrix = transform_matrix_offset_center(transform_matrix, img.shape[0], img.shape[1])
affine_matrix = transform_matrix[:2, :2]
offset = transform_matrix[:2, 2]
img = np.stack([ndimage.interpolation.affine_transform(
img[:, :, c],
affine_matrix,
offset) for c in range(img.shape[2])], axis=2)
img = Image.fromarray(img)
return img
def translate_y(img, magnitude):
img = np.array(img)
magnitudes = np.linspace(-150/331, 150/331, 11)
transform_matrix = np.array([[1, 0, img.shape[0]*random.uniform(magnitudes[magnitude], magnitudes[magnitude+1])],
[0, 1, 0],
[0, 0, 1]])
transform_matrix = transform_matrix_offset_center(transform_matrix, img.shape[0], img.shape[1])
affine_matrix = transform_matrix[:2, :2]
offset = transform_matrix[:2, 2]
img = np.stack([ndimage.interpolation.affine_transform(
img[:, :, c],
affine_matrix,
offset) for c in range(img.shape[2])], axis=2)
img = Image.fromarray(img)
return img
def rotate(img, magnitude):
img = np.array(img)
magnitudes = np.linspace(-30, 30, 11)
theta = np.deg2rad(random.uniform(magnitudes[magnitude], magnitudes[magnitude+1]))
transform_matrix = np.array([[np.cos(theta), -np.sin(theta), 0],
[np.sin(theta), np.cos(theta), 0],
[0, 0, 1]])
transform_matrix = transform_matrix_offset_center(transform_matrix, img.shape[0], img.shape[1])
affine_matrix = transform_matrix[:2, :2]
offset = transform_matrix[:2, 2]
img = np.stack([ndimage.interpolation.affine_transform(
img[:, :, c],
affine_matrix,
offset) for c in range(img.shape[2])], axis=2)
img = Image.fromarray(img)
return img
def auto_contrast(img, magnitude):
img = ImageOps.autocontrast(img)
return img
def invert(img, magnitude):
img = ImageOps.invert(img)
return img
def equalize(img, magnitude):
img = ImageOps.equalize(img)
return img
def solarize(img, magnitude):
magnitudes = np.linspace(0, 256, 11)
img = ImageOps.solarize(img, random.uniform(magnitudes[magnitude], magnitudes[magnitude+1]))
return img
def posterize(img, magnitude):
magnitudes = np.linspace(4, 8, 11)
img = ImageOps.posterize(img, int(round(random.uniform(magnitudes[magnitude], magnitudes[magnitude+1]))))
return img
def contrast(img, magnitude):
magnitudes = np.linspace(0.1, 1.9, 11)
img = ImageEnhance.Contrast(img).enhance(random.uniform(magnitudes[magnitude], magnitudes[magnitude+1]))
return img
def color(img, magnitude):
magnitudes = np.linspace(0.1, 1.9, 11)
img = ImageEnhance.Color(img).enhance(random.uniform(magnitudes[magnitude], magnitudes[magnitude+1]))
return img
def brightness(img, magnitude):
magnitudes = np.linspace(0.1, 1.9, 11)
img = ImageEnhance.Brightness(img).enhance(random.uniform(magnitudes[magnitude], magnitudes[magnitude+1]))
return img
def sharpness(img, magnitude):
magnitudes = np.linspace(0.1, 1.9, 11)
img = ImageEnhance.Sharpness(img).enhance(random.uniform(magnitudes[magnitude], magnitudes[magnitude+1]))
return img
def cutout(org_img, magnitude=None):
img = np.array(img)
magnitudes = np.linspace(0, 60/331, 11)
img = np.copy(org_img)
mask_val = img.mean()
if magnitude is None:
mask_size = 16
else:
mask_size = int(round(img.shape[0]*random.uniform(magnitudes[magnitude], magnitudes[magnitude+1])))
top = np.random.randint(0 - mask_size//2, img.shape[0] - mask_size)
left = np.random.randint(0 - mask_size//2, img.shape[1] - mask_size)
bottom = top + mask_size
right = left + mask_size
if top < 0:
top = 0
if left < 0:
left = 0
img[top:bottom, left:right, :].fill(mask_val)
img = Image.fromarray(img)
return img
class Cutout(object):
def __init__(self, length=16):
self.length = length
def __call__(self, img):
img = np.array(img)
mask_val = img.mean()
top = np.random.randint(0 - self.length//2, img.shape[0] - self.length)
left = np.random.randint(0 - self.length//2, img.shape[1] - self.length)
bottom = top + self.length
right = left + self.length
top = 0 if top < 0 else top
left = 0 if left < 0 else top
img[top:bottom, left:right, :] = mask_val
img = Image.fromarray(img)
return img
