|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import numpy as np
|
|
import scipy.misc
|
|
import cv2
|
|
from PIL import Image
|
|
|
|
|
|
class Augmentor():
|
|
def __init__(self,
|
|
crop_size=(256, 256),
|
|
scale_augm_prb=0.5, scale_augm_range=0.2,
|
|
rotation_augm_prb=0.5, rotation_augm_range=0.15,
|
|
hsv_augm_prb=1.0,
|
|
hue_augm_shift=0.05,
|
|
saturation_augm_shift=0.05, saturation_augm_scale=0.05,
|
|
value_augm_shift=0.05, value_augm_scale=0.05,
|
|
affine_trnsfm_prb=0.5, affine_trnsfm_range=0.05,
|
|
horizontal_flip_prb=0.5,
|
|
vertical_flip_prb=0.5):
|
|
|
|
self.crop_size = crop_size
|
|
|
|
self.scale_augm_prb = scale_augm_prb
|
|
self.scale_augm_range = scale_augm_range
|
|
|
|
self.rotation_augm_prb = rotation_augm_prb
|
|
self.rotation_augm_range = rotation_augm_range
|
|
|
|
self.hsv_augm_prb = hsv_augm_prb
|
|
self.hue_augm_shift = hue_augm_shift
|
|
self.saturation_augm_scale = saturation_augm_scale
|
|
self.saturation_augm_shift = saturation_augm_shift
|
|
self.value_augm_scale = value_augm_scale
|
|
self.value_augm_shift = value_augm_shift
|
|
|
|
self.affine_trnsfm_prb = affine_trnsfm_prb
|
|
self.affine_trnsfm_range = affine_trnsfm_range
|
|
|
|
self.horizontal_flip_prb = horizontal_flip_prb
|
|
self.vertical_flip_prb = vertical_flip_prb
|
|
|
|
def __call__(self, image, is_inference=False):
|
|
if is_inference:
|
|
return cv2.resize(image, None, fx=self.crop_size[0], fy=self.crop_size[1], interpolation=cv2.INTER_CUBIC)
|
|
|
|
|
|
if self.scale_augm_prb > np.random.uniform():
|
|
image = self.scale(image=image,
|
|
scale_x=1. + np.random.uniform(low=-self.scale_augm_range, high=-self.scale_augm_range),
|
|
scale_y=1. + np.random.uniform(low=-self.scale_augm_range, high=-self.scale_augm_range)
|
|
)
|
|
|
|
|
|
rows, cols, ch = image.shape
|
|
image = np.pad(array=image, pad_width=[[rows // 4, rows // 4], [cols // 4, cols // 4], [0, 0]], mode='reflect')
|
|
if self.rotation_augm_prb > np.random.uniform():
|
|
image = self.rotate(image=image,
|
|
angle=np.random.uniform(low=-self.rotation_augm_range*90.,
|
|
high=self.rotation_augm_range*90.)
|
|
)
|
|
|
|
if self.affine_trnsfm_prb > np.random.uniform():
|
|
image = self.affine(image=image,
|
|
rng=self.affine_trnsfm_range
|
|
)
|
|
image = image[(rows // 4):-(rows // 4), (cols // 4):-(cols // 4), :]
|
|
|
|
|
|
image = self.crop(image=image,
|
|
crop_size=self.crop_size
|
|
)
|
|
|
|
if self.hsv_augm_prb > np.random.uniform():
|
|
image = self.hsv_transform(image=image,
|
|
hue_shift=self.hue_augm_shift,
|
|
saturation_shift=self.saturation_augm_shift,
|
|
saturation_scale=self.saturation_augm_scale,
|
|
value_shift=self.value_augm_shift,
|
|
value_scale=self.value_augm_scale)
|
|
|
|
if self.horizontal_flip_prb > np.random.uniform():
|
|
image = self.horizontal_flip(image)
|
|
|
|
if self.vertical_flip_prb > np.random.uniform():
|
|
image = self.vertical_flip(image)
|
|
|
|
return image
|
|
|
|
def scale(self, image, scale_x, scale_y):
|
|
"""
|
|
Args:
|
|
image:
|
|
scale_x: float positive value. New horizontal scale
|
|
scale_y: float positive value. New vertical scale
|
|
Returns:
|
|
"""
|
|
image = cv2.resize(image, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_CUBIC)
|
|
return image
|
|
|
|
def rotate(self, image, angle):
|
|
"""
|
|
Args:
|
|
image: input image
|
|
angle: angle of rotation in degrees
|
|
Returns:
|
|
"""
|
|
rows, cols, ch = image.shape
|
|
|
|
rot_M = cv2.getRotationMatrix2D((cols / 2, rows / 2), angle, 1)
|
|
image = cv2.warpAffine(image, rot_M, (cols, rows))
|
|
return image
|
|
|
|
def crop(self, image, crop_size=(256, 256)):
|
|
rows, cols, chs = image.shape
|
|
x = int(np.random.uniform(low=0, high=max(0, rows - crop_size[0])))
|
|
y = int(np.random.uniform(low=0, high=max(0, cols - crop_size[1])))
|
|
|
|
image = image[x:x+crop_size[0], y:y+crop_size[1], :]
|
|
|
|
|
|
if image.shape[0] < crop_size[0] or image.shape[1] < crop_size[1]:
|
|
image = scipy.misc.imresize(arr=image, size=crop_size)
|
|
return image
|
|
|
|
def hsv_transform(self, image,
|
|
hue_shift=0.2,
|
|
saturation_shift=0.2, saturation_scale=0.2,
|
|
value_shift=0.2, value_scale=0.2,
|
|
):
|
|
|
|
image = Image.fromarray(image)
|
|
hsv = np.array(image.convert("HSV"), 'float64')
|
|
|
|
|
|
hsv /= 255.
|
|
|
|
|
|
hsv[..., 0] += np.random.uniform(-hue_shift, hue_shift)
|
|
hsv[..., 1] *= np.random.uniform(1. / (1. + saturation_scale), 1. + saturation_scale)
|
|
hsv[..., 1] += np.random.uniform(-saturation_shift, saturation_shift)
|
|
hsv[..., 2] *= np.random.uniform(1. / (1. + value_scale), 1. + value_scale)
|
|
hsv[..., 2] += np.random.uniform(-value_shift, value_shift)
|
|
|
|
|
|
hsv.clip(0.01, 0.99, hsv)
|
|
|
|
|
|
hsv = np.uint8(np.round(hsv * 254.))
|
|
|
|
|
|
return np.asarray(Image.fromarray(hsv, "HSV").convert("RGB"))
|
|
|
|
|
|
def affine(self, image, rng):
|
|
rows, cols, ch = image.shape
|
|
pts1 = np.float32([[0., 0.], [0., 1.], [1., 0.]])
|
|
[x0, y0] = [0. + np.random.uniform(low=-rng, high=rng), 0. + np.random.uniform(low=-rng, high=rng)]
|
|
[x1, y1] = [0. + np.random.uniform(low=-rng, high=rng), 1. + np.random.uniform(low=-rng, high=rng)]
|
|
[x2, y2] = [1. + np.random.uniform(low=-rng, high=rng), 0. + np.random.uniform(low=-rng, high=rng)]
|
|
pts2 = np.float32([[x0, y0], [x1, y1], [x2, y2]])
|
|
affine_M = cv2.getAffineTransform(pts1, pts2)
|
|
image = cv2.warpAffine(image, affine_M, (cols, rows))
|
|
|
|
return image
|
|
|
|
def horizontal_flip(self, image):
|
|
return image[:, ::-1, :]
|
|
|
|
def vertical_flip(self, image):
|
|
return image[::-1, :, :]
|
|
|
|
|