wehub-resource-sync
242 lines
6.5 KiB
Python
242 lines
6.5 KiB
Python
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image, ImageOps
|
|
|
|
'''
|
|
PIL resize (W,H)
|
|
Torch resize is (H,W)
|
|
'''
|
|
class Stretch:
|
|
def __init__(self):
|
|
self.tps = cv2.createThinPlateSplineShapeTransformer()
|
|
|
|
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 = len(b)-1
|
|
else:
|
|
index = mag
|
|
frac = b[index]
|
|
|
|
# left-most
|
|
srcpt.append([P, P])
|
|
srcpt.append([P, H-P])
|
|
srcpt.append([P, H_50])
|
|
x = np.random.uniform(0, frac)*W_33 #if np.random.uniform(0,1) > 0.5 else 0
|
|
dstpt.append([P+x, P])
|
|
dstpt.append([P+x, H-P])
|
|
dstpt.append([P+x, H_50])
|
|
|
|
# 2nd left-most
|
|
srcpt.append([P+W_33, P])
|
|
srcpt.append([P+W_33, H-P])
|
|
x = np.random.uniform(-frac, frac)*W_33
|
|
dstpt.append([P+W_33+x, P])
|
|
dstpt.append([P+W_33+x, H-P])
|
|
|
|
# 3rd left-most
|
|
srcpt.append([P+W_66, P])
|
|
srcpt.append([P+W_66, H-P])
|
|
x = np.random.uniform(-frac, frac)*W_33
|
|
dstpt.append([P+W_66+x, P])
|
|
dstpt.append([P+W_66+x, H-P])
|
|
|
|
# right-most
|
|
srcpt.append([W-P, P])
|
|
srcpt.append([W-P, H-P])
|
|
srcpt.append([W-P, H_50])
|
|
x = np.random.uniform(-frac, 0)*W_33 #if np.random.uniform(0,1) > 0.5 else 0
|
|
dstpt.append([W-P+x, P])
|
|
dstpt.append([W-P+x, H-P])
|
|
dstpt.append([W-P+x, H_50])
|
|
|
|
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)
|
|
|
|
return img
|
|
|
|
|
|
class Distort:
|
|
def __init__(self):
|
|
self.tps = cv2.createThinPlateSplineShapeTransformer()
|
|
|
|
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 = len(b)-1
|
|
else:
|
|
index = mag
|
|
frac = b[index]
|
|
|
|
# top pts
|
|
srcpt.append([P, P])
|
|
x = np.random.uniform(0, frac)*W_33
|
|
y = np.random.uniform(0, frac)*H_50
|
|
dstpt.append([P+x, P+y])
|
|
|
|
srcpt.append([P+W_33, P])
|
|
x = np.random.uniform(-frac, frac)*W_33
|
|
y = np.random.uniform(0, frac)*H_50
|
|
dstpt.append([P+W_33+x, P+y])
|
|
|
|
srcpt.append([P+W_66, P])
|
|
x = np.random.uniform(-frac, frac)*W_33
|
|
y = np.random.uniform(0, frac)*H_50
|
|
dstpt.append([P+W_66+x, P+y])
|
|
|
|
srcpt.append([W-P, P])
|
|
x = np.random.uniform(-frac, 0)*W_33
|
|
y = np.random.uniform(0, frac)*H_50
|
|
dstpt.append([W-P+x, P+y])
|
|
|
|
# bottom pts
|
|
srcpt.append([P, H-P])
|
|
x = np.random.uniform(0, frac)*W_33
|
|
y = np.random.uniform(-frac, 0)*H_50
|
|
dstpt.append([P+x, H-P+y])
|
|
|
|
srcpt.append([P+W_33, H-P])
|
|
x = np.random.uniform(-frac, frac)*W_33
|
|
y = np.random.uniform(-frac, 0)*H_50
|
|
dstpt.append([P+W_33+x, H-P+y])
|
|
|
|
srcpt.append([P+W_66, H-P])
|
|
x = np.random.uniform(-frac, frac)*W_33
|
|
y = np.random.uniform(-frac, 0)*H_50
|
|
dstpt.append([P+W_66+x, H-P+y])
|
|
|
|
srcpt.append([W-P, H-P])
|
|
x = np.random.uniform(-frac, 0)*W_33
|
|
y = np.random.uniform(-frac, 0)*H_50
|
|
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)
|
|
|
|
return img
|
|
|
|
|
|
class Curve:
|
|
def __init__(self, square_side=224):
|
|
self.tps = cv2.createThinPlateSplineShapeTransformer()
|
|
self.side = square_side
|
|
|
|
def __call__(self, img, 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)
|
|
|
|
isflip = np.random.uniform(0,1) > 0.5
|
|
if isflip:
|
|
img = ImageOps.flip(img)
|
|
#img = TF.vflip(img)
|
|
|
|
img = np.array(img)
|
|
w = self.side
|
|
h = self.side
|
|
w_25 = 0.25 * w
|
|
w_50 = 0.50 * w
|
|
w_75 = 0.75 * w
|
|
|
|
b = [1.1, .95, .8]
|
|
if mag<0 or mag>=len(b):
|
|
index = 0
|
|
else:
|
|
index = mag
|
|
rmin = b[index]
|
|
|
|
r = np.random.uniform(rmin, rmin+.1)*h
|
|
x1 = (r**2 - w_50**2)**0.5
|
|
h1 = r - x1
|
|
|
|
t = np.random.uniform(0.4, 0.5)*h
|
|
|
|
w2 = w_50*t/r
|
|
hi = x1*t/r
|
|
h2 = h1 + hi
|
|
|
|
sinb_2 = ((1 - x1/r)/2)**0.5
|
|
cosb_2 = ((1 + x1/r)/2)**0.5
|
|
w3 = w_50 - r*sinb_2
|
|
h3 = r - r*cosb_2
|
|
|
|
w4 = w_50 - (r-t)*sinb_2
|
|
h4 = r - (r-t)*cosb_2
|
|
|
|
w5 = 0.5*w2
|
|
h5 = h1 + 0.5*hi
|
|
h_50 = 0.50*h
|
|
|
|
srcpt = [(0,0 ), (w,0 ), (w_50,0), (0,h ), (w,h ), (w_25,0), (w_75,0 ), (w_50,h), (w_25,h), (w_75,h ), (0,h_50), (w,h_50 )]
|
|
dstpt = [(0,h1), (w,h1), (w_50,0), (w2,h2), (w-w2,h2), (w3, h3), (w-w3,h3), (w_50,t), (w4,h4 ), (w-w4,h4), (w5,h5 ), (w-w5,h5)]
|
|
|
|
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 isflip:
|
|
#img = TF.vflip(img)
|
|
img = ImageOps.flip(img)
|
|
rect = (0, self.side//2, self.side, self.side)
|
|
else:
|
|
rect = (0, 0, self.side, self.side//2)
|
|
|
|
img = img.crop(rect)
|
|
img = img.resize((W, H), Image.BICUBIC)
|
|
return img
|
|
|
|
|