wehub-resource-sync
234 lines
5.8 KiB
Python
234 lines
5.8 KiB
Python
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image, ImageOps
|
|
|
|
'''
|
|
PIL resize (W,H)
|
|
Torch resize is (H,W)
|
|
'''
|
|
class Shrink:
|
|
def __init__(self):
|
|
self.tps = cv2.createThinPlateSplineShapeTransformer()
|
|
self.translateXAbs = TranslateXAbs()
|
|
self.translateYAbs = TranslateYAbs()
|
|
|
|
def __call__(self, img, mag=-1, prob=1.):
|
|
if np.random.uniform(0,1) > prob:
|
|
return img
|
|
|
|
W, H = img.size
|
|
img = np.array(img)
|
|
srcpt = list()
|
|
dstpt = list()
|
|
|
|
W_33 = 0.33 * W
|
|
W_50 = 0.50 * W
|
|
W_66 = 0.66 * W
|
|
|
|
H_50 = 0.50 * H
|
|
|
|
P = 0
|
|
|
|
#frac = 0.4
|
|
|
|
b = [.2, .3, .4]
|
|
if mag<0 or mag>=len(b):
|
|
index = 0
|
|
else:
|
|
index = mag
|
|
frac = b[index]
|
|
|
|
# left-most
|
|
srcpt.append([P, P])
|
|
srcpt.append([P, H-P])
|
|
x = np.random.uniform(frac-.1, frac)*W_33
|
|
y = np.random.uniform(frac-.1, frac)*H_50
|
|
dstpt.append([P+x, P+y])
|
|
dstpt.append([P+x, H-P-y])
|
|
|
|
# 2nd left-most
|
|
srcpt.append([P+W_33, P])
|
|
srcpt.append([P+W_33, H-P])
|
|
dstpt.append([P+W_33, P+y])
|
|
dstpt.append([P+W_33, H-P-y])
|
|
|
|
# 3rd left-most
|
|
srcpt.append([P+W_66, P])
|
|
srcpt.append([P+W_66, H-P])
|
|
dstpt.append([P+W_66, P+y])
|
|
dstpt.append([P+W_66, H-P-y])
|
|
|
|
# right-most
|
|
srcpt.append([W-P, P])
|
|
srcpt.append([W-P, H-P])
|
|
dstpt.append([W-P-x, P+y])
|
|
dstpt.append([W-P-x, H-P-y])
|
|
|
|
N = len(dstpt)
|
|
matches = [cv2.DMatch(i, i, 0) for i in range(N)]
|
|
dst_shape = np.array(dstpt).reshape((-1, N, 2))
|
|
src_shape = np.array(srcpt).reshape((-1, N, 2))
|
|
self.tps.estimateTransformation(dst_shape, src_shape, matches)
|
|
img = self.tps.warpImage(img)
|
|
img = Image.fromarray(img)
|
|
|
|
if np.random.uniform(0, 1) < 0.5:
|
|
img = self.translateXAbs(img, val=x)
|
|
else:
|
|
img = self.translateYAbs(img, val=y)
|
|
|
|
return img
|
|
|
|
|
|
class Rotate:
|
|
def __init__(self, square_side=224):
|
|
self.side = square_side
|
|
|
|
def __call__(self, img, iscurve=False, mag=-1, prob=1.):
|
|
if np.random.uniform(0,1) > prob:
|
|
return img
|
|
|
|
W, H = img.size
|
|
|
|
if H!=self.side or W!=self.side:
|
|
img = img.resize((self.side, self.side), Image.BICUBIC)
|
|
|
|
b = [20., 40, 60]
|
|
if mag<0 or mag>=len(b):
|
|
index = 1
|
|
else:
|
|
index = mag
|
|
rotate_angle = b[index]
|
|
|
|
angle = np.random.uniform(rotate_angle-20, rotate_angle)
|
|
if np.random.uniform(0, 1) < 0.5:
|
|
angle = -angle
|
|
|
|
#angle = np.random.normal(loc=0., scale=rotate_angle)
|
|
#angle = min(angle, 2*rotate_angle)
|
|
#angle = max(angle, -2*rotate_angle)
|
|
|
|
expand = False if iscurve else True
|
|
img = img.rotate(angle=angle, resample=Image.BICUBIC, expand=expand)
|
|
img = img.resize((W, H), Image.BICUBIC)
|
|
|
|
return img
|
|
|
|
class Perspective:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def __call__(self, img, mag=-1, prob=1.):
|
|
if np.random.uniform(0,1) > prob:
|
|
return img
|
|
|
|
W, H = img.size
|
|
|
|
# upper-left, upper-right, lower-left, lower-right
|
|
src = np.float32([[0, 0], [W, 0], [0, H], [W, H]])
|
|
#low = 0.3
|
|
|
|
b = [.1, .2, .3]
|
|
if mag<0 or mag>=len(b):
|
|
index = 2
|
|
else:
|
|
index = mag
|
|
low = b[index]
|
|
|
|
high = 1 - low
|
|
if np.random.uniform(0, 1) > 0.5:
|
|
toprightY = np.random.uniform(low, low+.1)*H
|
|
bottomrightY = np.random.uniform(high-.1, high)*H
|
|
dest = np.float32([[0, 0], [W, toprightY], [0, H], [W, bottomrightY]])
|
|
else:
|
|
topleftY = np.random.uniform(low, low+.1)*H
|
|
bottomleftY = np.random.uniform(high-.1, high)*H
|
|
dest = np.float32([[0, topleftY], [W, 0], [0, bottomleftY], [W, H]])
|
|
M = cv2.getPerspectiveTransform(src, dest)
|
|
img = np.array(img)
|
|
img = cv2.warpPerspective(img, M, (W, H) )
|
|
img = Image.fromarray(img)
|
|
|
|
return img
|
|
|
|
|
|
class TranslateX:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def __call__(self, img, mag=-1, prob=1.):
|
|
if np.random.uniform(0,1) > prob:
|
|
return img
|
|
|
|
b = [.03, .06, .09]
|
|
if mag<0 or mag>=len(b):
|
|
index = 2
|
|
else:
|
|
index = mag
|
|
v = b[index]
|
|
v = np.random.uniform(v-0.03, v)
|
|
|
|
v = v * img.size[0]
|
|
if np.random.uniform(0,1) > 0.5:
|
|
v = -v
|
|
return img.transform(img.size, Image.AFFINE, (1, 0, v, 0, 1, 0))
|
|
|
|
|
|
class TranslateY:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def __call__(self, img, mag=-1, prob=1.):
|
|
if np.random.uniform(0,1) > prob:
|
|
return img
|
|
|
|
b = [.07, .14, .21]
|
|
if mag<0 or mag>=len(b):
|
|
index = 2
|
|
else:
|
|
index = mag
|
|
v = b[index]
|
|
v = np.random.uniform(v-0.07, v)
|
|
|
|
v = v * img.size[1]
|
|
if np.random.uniform(0,1) > 0.5:
|
|
v = -v
|
|
return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, v))
|
|
|
|
|
|
class TranslateXAbs:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def __call__(self, img, val=0, prob=1.):
|
|
if np.random.uniform(0,1) > prob:
|
|
return img
|
|
|
|
v = np.random.uniform(0, val)
|
|
|
|
if np.random.uniform(0,1) > 0.5:
|
|
v = -v
|
|
return img.transform(img.size, Image.AFFINE, (1, 0, v, 0, 1, 0))
|
|
|
|
|
|
class TranslateYAbs:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def __call__(self, img, val=0, prob=1.):
|
|
if np.random.uniform(0,1) > prob:
|
|
return img
|
|
|
|
v = np.random.uniform(0, val)
|
|
|
|
if np.random.uniform(0,1) > 0.5:
|
|
v = -v
|
|
return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, v))
|
|
|
|
|
|
|
|
|
|
|
|
|