584 lines
20 KiB
Python
584 lines
20 KiB
Python
import math
|
|
import random
|
|
|
|
import cv2
|
|
import lib.functional as F
|
|
import numpy as np
|
|
from PIL import Image
|
|
from scipy.ndimage.filters import gaussian_filter
|
|
from skimage import color
|
|
from skimage.draw import random_shapes
|
|
from skimage.filters import gaussian
|
|
from skimage.transform import resize
|
|
|
|
cv2.setNumThreads(0)
|
|
from numba import jit, u1, u2
|
|
|
|
|
|
class RGB2Lab(object):
|
|
def __init__(self):
|
|
pass
|
|
|
|
def __call__(self, inputs):
|
|
return color.rgb2lab(inputs)
|
|
|
|
|
|
class Guassian_noise(object):
|
|
"""Elastic distortion"""
|
|
|
|
def __init__(self, noise_sigma=0.1):
|
|
self.noise_sigma = noise_sigma
|
|
|
|
def __call__(self, inputs):
|
|
h = inputs.shape[0]
|
|
w = inputs.shape[1]
|
|
noisy_image = inputs
|
|
noise = np.random.randn(h, w) * self.noise_sigma
|
|
noisy_image[:, :, 0] = inputs[:, :, 0] + noise
|
|
|
|
return noisy_image
|
|
|
|
|
|
class Distortion(object):
|
|
"""Elastic distortion"""
|
|
|
|
def __init__(self, distortion_level=3, flip_probability=0):
|
|
self.alpha_max = distortion_level
|
|
self.flip_probability = flip_probability
|
|
|
|
def __call__(self, inputs):
|
|
if np.random.rand() < self.flip_probability:
|
|
inputs = inputs.transpose(Image.FLIP_LEFT_RIGHT)
|
|
|
|
inputs = np.array(inputs)
|
|
alpha = np.random.rand() * self.alpha_max
|
|
sigma = 50
|
|
random_state = np.random.RandomState(None)
|
|
shape = inputs.shape[0], inputs.shape[1]
|
|
|
|
dx = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha * 1000
|
|
dy = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha * 1000
|
|
|
|
x, y = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing="ij")
|
|
remap_image = cv2.remap(
|
|
inputs, (dy + y).astype(np.float32), (dx + x).astype(np.float32), interpolation=cv2.INTER_LINEAR
|
|
)
|
|
|
|
return Image.fromarray(remap_image)
|
|
|
|
|
|
class Distortion_with_flow(object):
|
|
"""Elastic distortion"""
|
|
|
|
def __init__(self):
|
|
return
|
|
|
|
def __call__(self, inputs, dx, dy):
|
|
inputs = np.array(inputs)
|
|
shape = inputs.shape[0], inputs.shape[1]
|
|
inputs = np.array(inputs)
|
|
remap_image = forward_mapping(inputs, dy, dx, maxIter=3, precision=1e-3)
|
|
|
|
return Image.fromarray(remap_image)
|
|
|
|
|
|
class Normalize(object):
|
|
def __init__(self):
|
|
pass
|
|
|
|
def __call__(self, inputs):
|
|
inputs[0:1, :, :] = F.normalize(inputs[0:1, :, :], 50, 1)
|
|
inputs[1:3, :, :] = F.normalize(inputs[1:3, :, :], (0, 0), (1, 1))
|
|
return inputs
|
|
|
|
|
|
class ToTensor(object):
|
|
def __init__(self):
|
|
pass
|
|
|
|
def __call__(self, inputs):
|
|
return F.to_mytensor(inputs)
|
|
|
|
|
|
class RandomErasing(object):
|
|
"""
|
|
Class that performs Random Erasing in Random Erasing Data Augmentation by Zhong et al.
|
|
-------------------------------------------------------------------------------------
|
|
probability: The probability that the operation will be performed.
|
|
sl: min erasing area
|
|
sh: max erasing area
|
|
r1: min aspect ratio
|
|
mean: erasing value
|
|
-------------------------------------------------------------------------------------
|
|
"""
|
|
|
|
def __init__(self, probability=0.6, sl=0.05, sh=0.6):
|
|
self.probability = probability
|
|
self.sl = sl
|
|
self.sh = sh
|
|
|
|
def __call__(self, img):
|
|
img = np.array(img)
|
|
if random.uniform(0, 1) > self.probability:
|
|
return Image.fromarray(img)
|
|
|
|
area = img.shape[0] * img.shape[1]
|
|
h0 = img.shape[0]
|
|
w0 = img.shape[1]
|
|
channel = img.shape[2]
|
|
|
|
h = int(round(random.uniform(self.sl, self.sh) * h0))
|
|
w = int(round(random.uniform(self.sl, self.sh) * w0))
|
|
|
|
if w < img.shape[1] and h < img.shape[0]:
|
|
x1 = random.randint(0, img.shape[0] - h)
|
|
y1 = random.randint(0, img.shape[1] - w)
|
|
img[x1 : x1 + h, y1 : y1 + w, :] = np.random.rand(h, w, channel) * 255
|
|
|
|
return Image.fromarray(img)
|
|
|
|
return Image.fromarray(img)
|
|
|
|
|
|
class CenteredPad(object):
|
|
"""
|
|
pad the frame with black border,
|
|
make square image for processing
|
|
"""
|
|
|
|
def __init__(self, image_size):
|
|
self.image_size = image_size
|
|
|
|
def __call__(self, img):
|
|
img = np.array(img)
|
|
width = np.size(img, 1)
|
|
height = np.size(img, 0)
|
|
old_size = [height, width]
|
|
|
|
ratio = float(self.image_size) / max(height, width)
|
|
new_size = [int(x * ratio) for x in old_size]
|
|
I_resize = resize(img, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
|
width_new = np.size(I_resize, 1)
|
|
height_new = np.size(I_resize, 0)
|
|
|
|
I_pad = np.zeros((self.image_size, self.image_size, 3))
|
|
start_height = (self.image_size - new_size[0]) // 2
|
|
start_width = (self.image_size - new_size[1]) // 2
|
|
I_pad[start_height : (start_height + height_new), start_width : (start_width + width_new), :] = I_resize
|
|
|
|
return Image.fromarray(I_pad.astype(np.uint8))
|
|
|
|
|
|
class centeredPad_with_height(object):
|
|
"""
|
|
pad the image according to the height
|
|
"""
|
|
|
|
def __init__(self, image_size):
|
|
self.height = image_size[0]
|
|
self.width = image_size[1]
|
|
|
|
def __call__(self, image):
|
|
# pad the image to 16:9
|
|
# pad height
|
|
I = np.array(image)
|
|
|
|
# for padded input
|
|
if I.shape[0] == I.shape[1] and I.shape[1] == self.width:
|
|
return Image.fromarray(I.astype(np.uint8))
|
|
|
|
width0 = np.size(I, 1)
|
|
height0 = np.size(I, 0)
|
|
old_size = [height0, width0]
|
|
height = self.height
|
|
width = self.width
|
|
|
|
ratio = height / height0
|
|
new_size = [int(x * ratio) for x in old_size]
|
|
I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
|
width_new = np.size(I_resize, 1)
|
|
height_new = np.size(I_resize, 0)
|
|
|
|
# if exceed the expected width
|
|
if width_new > width:
|
|
I_resize = I_resize[:, math.floor(width_new - width) // 2 : (math.floor(width_new - width) // 2 + width), :]
|
|
width_new = np.size(I_resize, 1)
|
|
height_new = np.size(I_resize, 0)
|
|
|
|
# lines: 56~200
|
|
I_pad = np.zeros((width, width, 3))
|
|
start_height = (width - height_new) // 2
|
|
start_width = (width - width_new) // 2
|
|
I_pad[start_height : (start_height + height_new), start_width : (start_width + width_new), :] = I_resize
|
|
|
|
return Image.fromarray(I_pad.astype(np.uint8))
|
|
|
|
|
|
class CenterPad(object):
|
|
def __init__(self, image_size):
|
|
self.height = image_size[0]
|
|
self.width = image_size[1]
|
|
|
|
def __call__(self, image):
|
|
# pad the image to 16:9
|
|
# pad height
|
|
I = np.array(image)
|
|
|
|
# for padded input
|
|
height_old = np.size(I, 0)
|
|
width_old = np.size(I, 1)
|
|
old_size = [height_old, width_old]
|
|
height = self.height
|
|
width = self.width
|
|
I_pad = np.zeros((height, width, np.size(I, 2)))
|
|
|
|
ratio = height / width
|
|
if height_old / width_old == ratio:
|
|
if height_old == height:
|
|
return Image.fromarray(I.astype(np.uint8))
|
|
new_size = [int(x * height / height_old) for x in old_size]
|
|
I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
|
return Image.fromarray(I_resize.astype(np.uint8))
|
|
|
|
if height_old / width_old > ratio: # pad the width and crop
|
|
new_size = [int(x * width / width_old) for x in old_size]
|
|
I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
|
width_resize = np.size(I_resize, 1)
|
|
height_resize = np.size(I_resize, 0)
|
|
start_height = (height_resize - height) // 2
|
|
I_pad[:, :, :] = I_resize[start_height : (start_height + height), :, :]
|
|
else: # pad the height and crop
|
|
new_size = [int(x * height / height_old) for x in old_size]
|
|
I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
|
width_resize = np.size(I_resize, 1)
|
|
height_resize = np.size(I_resize, 0)
|
|
start_width = (width_resize - width) // 2
|
|
I_pad[:, :, :] = I_resize[:, start_width : (start_width + width), :]
|
|
|
|
return Image.fromarray(I_pad.astype(np.uint8))
|
|
|
|
|
|
class CenterPad_threshold(object):
|
|
def __init__(self, image_size, threshold=3 / 4):
|
|
self.height = image_size[0]
|
|
self.width = image_size[1]
|
|
self.threshold = threshold
|
|
|
|
def __call__(self, image):
|
|
# pad the image to 16:9
|
|
# pad height
|
|
I = np.array(image)
|
|
|
|
# for padded input
|
|
height_old = np.size(I, 0)
|
|
width_old = np.size(I, 1)
|
|
old_size = [height_old, width_old]
|
|
height = self.height
|
|
width = self.width
|
|
I_pad = np.zeros((height, width, np.size(I, 2)))
|
|
|
|
ratio = height / width
|
|
|
|
if height_old / width_old == ratio:
|
|
if height_old == height:
|
|
return Image.fromarray(I.astype(np.uint8))
|
|
new_size = [int(x * height / height_old) for x in old_size]
|
|
I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
|
return Image.fromarray(I_resize.astype(np.uint8))
|
|
|
|
if height_old / width_old > self.threshold:
|
|
width_new, height_new = width_old, int(width_old * self.threshold)
|
|
height_margin = height_old - height_new
|
|
height_crop_start = height_margin // 2
|
|
I_crop = I[height_crop_start : (height_crop_start + height_new), :, :]
|
|
I_resize = resize(
|
|
I_crop, [height, width], mode="reflect", preserve_range=True, clip=False, anti_aliasing=True
|
|
)
|
|
|
|
return Image.fromarray(I_resize.astype(np.uint8))
|
|
|
|
if height_old / width_old > ratio: # pad the width and crop
|
|
new_size = [int(x * width / width_old) for x in old_size]
|
|
I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
|
width_resize = np.size(I_resize, 1)
|
|
height_resize = np.size(I_resize, 0)
|
|
start_height = (height_resize - height) // 2
|
|
I_pad[:, :, :] = I_resize[start_height : (start_height + height), :, :]
|
|
else: # pad the height and crop
|
|
new_size = [int(x * height / height_old) for x in old_size]
|
|
I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
|
width_resize = np.size(I_resize, 1)
|
|
height_resize = np.size(I_resize, 0)
|
|
start_width = (width_resize - width) // 2
|
|
I_pad[:, :, :] = I_resize[:, start_width : (start_width + width), :]
|
|
|
|
return Image.fromarray(I_pad.astype(np.uint8))
|
|
|
|
|
|
class CenterPadCrop_numpy(object):
|
|
"""
|
|
pad the image according to the height
|
|
"""
|
|
|
|
def __init__(self, image_size):
|
|
self.height = image_size[0]
|
|
self.width = image_size[1]
|
|
|
|
def __call__(self, image, threshold=3 / 4):
|
|
# pad the image to 16:9
|
|
# pad height
|
|
I = np.array(image)
|
|
# for padded input
|
|
height_old = np.size(I, 0)
|
|
width_old = np.size(I, 1)
|
|
old_size = [height_old, width_old]
|
|
height = self.height
|
|
width = self.width
|
|
padding_size = width
|
|
if image.ndim == 2:
|
|
I_pad = np.zeros((width, width))
|
|
else:
|
|
I_pad = np.zeros((width, width, I.shape[2]))
|
|
|
|
ratio = height / width
|
|
if height_old / width_old == ratio:
|
|
return I
|
|
|
|
if height_old / width_old > threshold:
|
|
width_new, height_new = width_old, int(width_old * threshold)
|
|
height_margin = height_old - height_new
|
|
height_crop_start = height_margin // 2
|
|
I_crop = I[height_start : (height_start + height_new), :]
|
|
I_resize = resize(
|
|
I_crop, [height, width], mode="reflect", preserve_range=True, clip=False, anti_aliasing=True
|
|
)
|
|
return I_resize
|
|
|
|
if height_old / width_old > ratio: # pad the width and crop
|
|
new_size = [int(x * width / width_old) for x in old_size]
|
|
I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
|
width_resize = np.size(I_resize, 1)
|
|
height_resize = np.size(I_resize, 0)
|
|
start_height = (height_resize - height) // 2
|
|
start_height_block = (padding_size - height) // 2
|
|
if image.ndim == 2:
|
|
I_pad[start_height_block : (start_height_block + height), :] = I_resize[
|
|
start_height : (start_height + height), :
|
|
]
|
|
else:
|
|
I_pad[start_height_block : (start_height_block + height), :, :] = I_resize[
|
|
start_height : (start_height + height), :, :
|
|
]
|
|
else: # pad the height and crop
|
|
new_size = [int(x * height / height_old) for x in old_size]
|
|
I_resize = resize(I, new_size, mode="reflect", preserve_range=True, clip=False, anti_aliasing=True)
|
|
width_resize = np.size(I_resize, 1)
|
|
height_resize = np.size(I_resize, 0)
|
|
start_width = (width_resize - width) // 2
|
|
start_width_block = (padding_size - width) // 2
|
|
if image.ndim == 2:
|
|
I_pad[:, start_width_block : (start_width_block + width)] = I_resize[
|
|
:, start_width : (start_width + width)
|
|
]
|
|
|
|
else:
|
|
I_pad[:, start_width_block : (start_width_block + width), :] = I_resize[
|
|
:, start_width : (start_width + width), :
|
|
]
|
|
|
|
crop_start_height = (I_pad.shape[0] - height) // 2
|
|
crop_start_width = (I_pad.shape[1] - width) // 2
|
|
|
|
if image.ndim == 2:
|
|
return I_pad[
|
|
crop_start_height : (crop_start_height + height), crop_start_width : (crop_start_width + width)
|
|
]
|
|
else:
|
|
return I_pad[
|
|
crop_start_height : (crop_start_height + height), crop_start_width : (crop_start_width + width), :
|
|
]
|
|
|
|
|
|
@jit(nopython=True)
|
|
def iterSearchShader(padu, padv, xr, yr, W, H, maxIter, precision):
|
|
# print('processing location', (xr, yr))
|
|
#
|
|
if abs(padu[yr, xr]) < precision and abs(padv[yr, xr]) < precision:
|
|
return xr, yr
|
|
|
|
# Our initialize method in this paper, can see the overleaf for detail
|
|
if (xr + 1) <= (W - 1):
|
|
dif = padu[yr, xr + 1] - padu[yr, xr]
|
|
else:
|
|
dif = padu[yr, xr] - padu[yr, xr - 1]
|
|
u_next = padu[yr, xr] / (1 + dif)
|
|
if (yr + 1) <= (H - 1):
|
|
dif = padv[yr + 1, xr] - padv[yr, xr]
|
|
else:
|
|
dif = padv[yr, xr] - padv[yr - 1, xr]
|
|
v_next = padv[yr, xr] / (1 + dif)
|
|
i = xr - u_next
|
|
j = yr - v_next
|
|
i_int = int(i)
|
|
j_int = int(j)
|
|
|
|
# The same as traditional iterative search method
|
|
for _ in range(maxIter):
|
|
if not 0 <= i <= (W - 1) or not 0 <= j <= (H - 1):
|
|
return i, j
|
|
|
|
u11 = padu[j_int, i_int]
|
|
v11 = padv[j_int, i_int]
|
|
|
|
u12 = padu[j_int, i_int + 1]
|
|
v12 = padv[j_int, i_int + 1]
|
|
|
|
int1 = padu[j_int + 1, i_int]
|
|
v21 = padv[j_int + 1, i_int]
|
|
|
|
int2 = padu[j_int + 1, i_int + 1]
|
|
v22 = padv[j_int + 1, i_int + 1]
|
|
|
|
u = (
|
|
u11 * (i_int + 1 - i) * (j_int + 1 - j)
|
|
+ u12 * (i - i_int) * (j_int + 1 - j)
|
|
+ int1 * (i_int + 1 - i) * (j - j_int)
|
|
+ int2 * (i - i_int) * (j - j_int)
|
|
)
|
|
|
|
v = (
|
|
v11 * (i_int + 1 - i) * (j_int + 1 - j)
|
|
+ v12 * (i - i_int) * (j_int + 1 - j)
|
|
+ v21 * (i_int + 1 - i) * (j - j_int)
|
|
+ v22 * (i - i_int) * (j - j_int)
|
|
)
|
|
|
|
i_next = xr - u
|
|
j_next = yr - v
|
|
|
|
if abs(i - i_next) < precision and abs(j - j_next) < precision:
|
|
return i, j
|
|
|
|
i = i_next
|
|
j = j_next
|
|
|
|
# if the search doesn't converge within max iter, it will return the last iter result
|
|
return i_next, j_next
|
|
|
|
|
|
# Bilinear interpolation
|
|
@jit(nopython=True)
|
|
def biInterpolation(distorted, i, j):
|
|
i = u2(i)
|
|
j = u2(j)
|
|
Q11 = distorted[j, i]
|
|
Q12 = distorted[j, i + 1]
|
|
Q21 = distorted[j + 1, i]
|
|
Q22 = distorted[j + 1, i + 1]
|
|
|
|
return u1(
|
|
Q11 * (i + 1 - i) * (j + 1 - j)
|
|
+ Q12 * (i - i) * (j + 1 - j)
|
|
+ Q21 * (i + 1 - i) * (j - j)
|
|
+ Q22 * (i - i) * (j - j)
|
|
)
|
|
|
|
|
|
@jit(nopython=True)
|
|
def iterSearch(distortImg, resultImg, padu, padv, W, H, maxIter=5, precision=1e-2):
|
|
for xr in range(W):
|
|
for yr in range(H):
|
|
# (xr, yr) is the point in result image, (i, j) is the search result in distorted image
|
|
i, j = iterSearchShader(padu, padv, xr, yr, W, H, maxIter, precision)
|
|
|
|
# reflect the pixels outside the border
|
|
if i > W - 1:
|
|
i = 2 * W - 1 - i
|
|
if i < 0:
|
|
i = -i
|
|
if j > H - 1:
|
|
j = 2 * H - 1 - j
|
|
if j < 0:
|
|
j = -j
|
|
|
|
# Bilinear interpolation to get the pixel at (i, j) in distorted image
|
|
resultImg[yr, xr, 0] = biInterpolation(
|
|
distortImg[:, :, 0],
|
|
i,
|
|
j,
|
|
)
|
|
resultImg[yr, xr, 1] = biInterpolation(
|
|
distortImg[:, :, 1],
|
|
i,
|
|
j,
|
|
)
|
|
resultImg[yr, xr, 2] = biInterpolation(
|
|
distortImg[:, :, 2],
|
|
i,
|
|
j,
|
|
)
|
|
return None
|
|
|
|
|
|
def forward_mapping(source_image, u, v, maxIter=5, precision=1e-2):
|
|
"""
|
|
warp the image according to the forward flow
|
|
u: horizontal
|
|
v: vertical
|
|
"""
|
|
H = source_image.shape[0]
|
|
W = source_image.shape[1]
|
|
|
|
distortImg = np.array(np.zeros((H + 1, W + 1, 3)), dtype=np.uint8)
|
|
distortImg[0:H, 0:W] = source_image[0:H, 0:W]
|
|
distortImg[H, 0:W] = source_image[H - 1, 0:W]
|
|
distortImg[0:H, W] = source_image[0:H, W - 1]
|
|
distortImg[H, W] = source_image[H - 1, W - 1]
|
|
|
|
padu = np.array(np.zeros((H + 1, W + 1)), dtype=np.float32)
|
|
padu[0:H, 0:W] = u[0:H, 0:W]
|
|
padu[H, 0:W] = u[H - 1, 0:W]
|
|
padu[0:H, W] = u[0:H, W - 1]
|
|
padu[H, W] = u[H - 1, W - 1]
|
|
|
|
padv = np.array(np.zeros((H + 1, W + 1)), dtype=np.float32)
|
|
padv[0:H, 0:W] = v[0:H, 0:W]
|
|
padv[H, 0:W] = v[H - 1, 0:W]
|
|
padv[0:H, W] = v[0:H, W - 1]
|
|
padv[H, W] = v[H - 1, W - 1]
|
|
|
|
resultImg = np.array(np.zeros((H, W, 3)), dtype=np.uint8)
|
|
iterSearch(distortImg, resultImg, padu, padv, W, H, maxIter, precision)
|
|
return resultImg
|
|
|
|
|
|
def random_mask(H, W, mask_size=200):
|
|
"""
|
|
mask: ranges in [0,1]
|
|
"""
|
|
masked_image = np.zeros([H, W, 3])
|
|
mask = random_shapes(
|
|
[H, W],
|
|
max_shapes=1,
|
|
min_shapes=1,
|
|
max_size=mask_size,
|
|
min_size=mask_size / 2,
|
|
multichannel=False,
|
|
intensity_range=[0, 0],
|
|
)[0]
|
|
mask = np.stack((mask, mask, mask), axis=-1)
|
|
random_state = np.random.RandomState(None)
|
|
distortion_range = 50
|
|
alpha = np.random.rand() * 6
|
|
forward_dx = (
|
|
gaussian_filter((random_state.rand(H, W) * 2 - 1), distortion_range, mode="constant", cval=0) * alpha * 1000
|
|
)
|
|
forward_dy = (
|
|
gaussian_filter((random_state.rand(H, W) * 2 - 1), distortion_range, mode="constant", cval=0) * alpha * 1000
|
|
)
|
|
mask = forward_mapping(mask, forward_dy, forward_dx, maxIter=3, precision=1e-3) / 255
|
|
mask = 1 - gaussian(mask, sigma=0, preserve_range=True, multichannel=False, anti_aliasing=True)
|
|
mask = mask[:, :, 0]
|
|
|
|
return mask
|