From 9b0a5dc6fec77c5edd45c66e46b5b050d400b352 Mon Sep 17 00:00:00 2001 From: reshow Date: Wed, 4 Dec 2019 10:25:31 +0800 Subject: [PATCH] update augmentation --- augmentation.py | 142 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 142 insertions(+) diff --git a/augmentation.py b/augmentation.py index f62cf4b..d58edbe 100644 --- a/augmentation.py +++ b/augmentation.py @@ -1,4 +1,93 @@ import numpy as np +from skimage import io, transform +import math +import copy +from PIL import ImageEnhance, ImageOps, ImageFile, Image +import cv2 + + +# import numba + + +# sometimes = lambda aug: iaa.Sometimes(0.5, aug) + + +def randomColor(image): + """ + """ + PIL_image = Image.fromarray((image * 255.).astype(np.uint8)) + random_factor = np.random.randint(0, 31) / 10. + color_image = ImageEnhance.Color(PIL_image).enhance(random_factor) # 调整图像的饱和度 + random_factor = np.random.randint(10, 21) / 10. + brightness_image = ImageEnhance.Brightness(color_image).enhance(random_factor) # 调整图像的亮度 + random_factor = np.random.randint(10, 21) / 10. + contrast_image = ImageEnhance.Contrast(brightness_image).enhance(random_factor) # 调整图像对比度 + random_factor = np.random.randint(0, 31) / 10. + out = np.array(ImageEnhance.Sharpness(contrast_image).enhance(random_factor)) + out = out / 255. + return out + + +def getRotateMatrix(angle, image_shape): + [image_height, image_width, image_channel] = image_shape + t1 = np.array([[1, 0, -image_height / 2.], [0, 1, -image_width / 2.], [0, 0, 1]]) + r1 = np.array([[math.cos(angle), math.sin(angle), 0], [math.sin(-angle), math.cos(angle), 0], [0, 0, 1]]) + t2 = np.array([[1, 0, image_height / 2.], [0, 1, image_width / 2.], [0, 0, 1]]) + rt_mat = t2.dot(r1).dot(t1) + t1 = np.array([[1, 0, -image_height / 2.], [0, 1, -image_width / 2.], [0, 0, 1]]) + r1 = np.array([[math.cos(-angle), math.sin(-angle), 0], [math.sin(angle), math.cos(-angle), 0], [0, 0, 1]]) + t2 = np.array([[1, 0, image_height / 2.], [0, 1, image_width / 2.], [0, 0, 1]]) + rt_mat_inv = t2.dot(r1).dot(t1) + return rt_mat.astype(np.float32), rt_mat_inv.astype(np.float32) + + +def getRotateMatrix3D(angle, image_shape): + [image_height, image_width, image_channel] = image_shape + t1 = np.array([[1, 0, 0, -image_height / 2.], [0, 1, 0, -image_width / 2.], [0, 0, 1, 0], [0, 0, 0, 1]]) + r1 = np.array([[math.cos(angle), math.sin(angle), 0, 0], [math.sin(-angle), math.cos(angle), 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]) + t2 = np.array([[1, 0, 0, image_height / 2.], [0, 1, 0, image_width / 2.], [0, 0, 1, 0], [0, 0, 0, 1]]) + rt_mat = t2.dot(r1).dot(t1) + t1 = np.array([[1, 0, 0, -image_height / 2.], [0, 1, 0, -image_width / 2.], [0, 0, 1, 0], [0, 0, 0, 1]]) + r1 = np.array([[math.cos(-angle), math.sin(-angle), 0, 0], [math.sin(angle), math.cos(-angle), 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]) + t2 = np.array([[1, 0, 0, image_height / 2.], [0, 1, 0, image_width / 2.], [0, 0, 1, 0], [0, 0, 0, 1]]) + rt_mat_inv = t2.dot(r1).dot(t1) + return rt_mat.astype(np.float32), rt_mat_inv.astype(np.float32) + + +# @numba.jit(numba.float32(numba.float32,numba.float32)) +def myDot(a, b): + return np.dot(a, b) + + +def rotateData(x, y, angle_range=45, specify_angle=None): + if specify_angle is None: + angle = np.random.randint(-angle_range, angle_range) + angle = angle / 180. * np.pi + else: + angle = specify_angle + [image_height, image_width, image_channel] = x.shape + # move-rotate-move + [rform, rform_inv] = getRotateMatrix(angle, x.shape) + + # rotate_x = transform.warp(x, rform_inv, + # output_shape=(image_height, image_width)) + rotate_x = cv2.warpPerspective(x, rform, (image_height, image_width)) + rotate_y = y.copy() + rotate_y[:, :, 2] = 1. + rotate_y = rotate_y.reshape(image_width * image_height, image_channel) + # rotate_y = rotate_y.dot(rform.T) + rotate_y = myDot(rotate_y, rform.T) + rotate_y = rotate_y.reshape(image_height, image_width, image_channel) + rotate_y[:, :, 2] = y[:, :, 2] + # for i in range(image_height): + # for j in range(image_width): + # rotate_y[i][j][2] = 1. + # rotate_y[i][j] = rotate_y[i][j].dot(rform.T) + # rotate_y[i][j][2] = y[i][j][2] + # tex = np.ones((256, 256, 3)) + # from visualize import show + # show([rotate_y, tex, rotate_x.astype(np.float32)], mode='uvmap') + return rotate_x, rotate_y def gaussNoise(x, mean=0, var=0.001): @@ -9,6 +98,46 @@ def gaussNoise(x, mean=0, var=0.001): return out +def randomErase(x, max_num=4, s_l=0.02, s_h=0.3, r_1=0.3, r_2=1 / 0.3, v_l=0, v_h=1.0): + [img_h, img_w, img_c] = x.shape + out = x.copy() + num = np.random.randint(1, max_num) + + for i in range(num): + s = np.random.uniform(s_l, s_h) * img_h * img_w + r = np.random.uniform(r_1, r_2) + w = int(np.sqrt(s / r)) + h = int(np.sqrt(s * r)) + left = np.random.randint(0, img_w) + top = np.random.randint(0, img_h) + mask = np.zeros((img_h, img_w)) + mask[top:min(top + h, img_h), left:min(left + w, img_w)] = 1 + if np.random.rand() < 0.25: + c = np.random.uniform(v_l, v_h) + out[mask > 0] = c + elif np.random.rand() < 0.75: + c0 = np.random.uniform(v_l, v_h) + c1 = np.random.uniform(v_l, v_h) + c2 = np.random.uniform(v_l, v_h) + out0 = out[:, :, 0] + out0[mask > 0] = c0 + out1 = out[:, :, 1] + out1[mask > 0] = c1 + out2 = out[:, :, 2] + out2[mask > 0] = c2 + else: + c0 = np.random.uniform(v_l, v_h) + c1 = np.random.uniform(v_l, v_h) + c2 = np.random.uniform(v_l, v_h) + out0 = out[:, :, 0] + out0[mask > 0] *= c0 + out1 = out[:, :, 1] + out1[mask > 0] *= c1 + out2 = out[:, :, 2] + out2[mask > 0] *= c2 + return out + + def channelScale(x, min_rate=0.6, max_rate=1.4): out = x.copy() for i in range(3): @@ -18,11 +147,24 @@ def channelScale(x, min_rate=0.6, max_rate=1.4): def prnAugment_torch(x, y, is_rotate=True): + if is_rotate: + if np.random.rand() > 0.5: + x, y = rotateData(x, y, 90) + if np.random.rand() > 0.75: + x = randomErase(x) if np.random.rand() > 0.5: x = channelScale(x) + # if np.random.rand() > 0.75: + # x = gaussNoise(x) return x, y +# def prnAugment_torch(x, y, is_rotate=True): +# if np.random.rand() > 0.5: +# x = channelScale(x) +# return x, y + + if __name__ == '__main__': import time from skimage import io