ast / img_augm.py
jgwill's picture
add:ast-app
1b677c1
raw
history blame
7.8 kB
# Copyright (C) 2018 Artsiom Sanakoyeu and Dmytro Kotovenko
#
# This file is part of Adaptive Style Transfer
#
# Adaptive Style Transfer is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Adaptive Style Transfer is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
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 not inference stage apply the pipeline of augmentations.
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), :]
# Crop out patch of desired size.
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 the input image was too small to comprise patch of size crop_size,
# resize obtained patch to desired size.
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')
# scale the values to fit between 0 and 1
hsv /= 255.
# do the scalings & shiftings
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)
# cut off invalid values
hsv.clip(0.01, 0.99, hsv)
# round to full numbers
hsv = np.uint8(np.round(hsv * 254.))
# convert back to rgb image
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, :, :]