|
|
|
from os.path import splitext |
|
import cv2 |
|
import numpy as np |
|
from keras.models import load_model |
|
import os |
|
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' |
|
|
|
class Label: |
|
def __init__(self, cl=-1, tl=np.array([0., 0.]), br=np.array([0., 0.]), prob=None): |
|
self.__tl = tl |
|
self.__br = br |
|
self.__cl = cl |
|
self.__prob = prob |
|
|
|
def __str__(self): |
|
return 'Class: %d, top left(x: %f, y: %f), bottom right(x: %f, y: %f)' % ( |
|
self.__cl, self.__tl[0], self.__tl[1], self.__br[0], self.__br[1]) |
|
|
|
def copy(self): |
|
return Label(self.__cl, self.__tl, self.__br) |
|
|
|
def wh(self): return self.__br - self.__tl |
|
|
|
def cc(self): return self.__tl + self.wh() / 2 |
|
|
|
def tl(self): return self.__tl |
|
|
|
def br(self): return self.__br |
|
|
|
def tr(self): return np.array([self.__br[0], self.__tl[1]]) |
|
|
|
def bl(self): return np.array([self.__tl[0], self.__br[1]]) |
|
|
|
def cl(self): return self.__cl |
|
|
|
def area(self): return np.prod(self.wh()) |
|
|
|
def prob(self): return self.__prob |
|
|
|
def set_class(self, cl): |
|
self.__cl = cl |
|
|
|
def set_tl(self, tl): |
|
self.__tl = tl |
|
|
|
def set_br(self, br): |
|
self.__br = br |
|
|
|
def set_wh(self, wh): |
|
cc = self.cc() |
|
self.__tl = cc - .5 * wh |
|
self.__br = cc + .5 * wh |
|
|
|
def set_prob(self, prob): |
|
self.__prob = prob |
|
|
|
class DLabel(Label): |
|
def __init__(self, cl, pts, prob): |
|
self.pts = pts |
|
tl = np.amin(pts, axis=1) |
|
br = np.amax(pts, axis=1) |
|
Label.__init__(self, cl, tl, br, prob) |
|
|
|
|
|
def im2single(Image): |
|
return Image.astype('float32') / 255 |
|
|
|
def getWH(shape): |
|
return np.array(shape[1::-1]).astype(float) |
|
|
|
def IOU(tl1, br1, tl2, br2): |
|
wh1, wh2 = br1-tl1, br2-tl2 |
|
assert((wh1 >= 0).all() and (wh2 >= 0).all()) |
|
|
|
intersection_wh = np.maximum(np.minimum(br1, br2) - np.maximum(tl1, tl2), 0) |
|
intersection_area = np.prod(intersection_wh) |
|
area1, area2 = (np.prod(wh1), np.prod(wh2)) |
|
union_area = area1 + area2 - intersection_area |
|
return intersection_area/union_area |
|
|
|
def IOU_labels(l1, l2): |
|
return IOU(l1.tl(), l1.br(), l2.tl(), l2.br()) |
|
|
|
def nms(Labels, iou_threshold=0.5): |
|
SelectedLabels = [] |
|
Labels.sort(key=lambda l: l.prob(), reverse=True) |
|
|
|
for label in Labels: |
|
non_overlap = True |
|
for sel_label in SelectedLabels: |
|
if IOU_labels(label, sel_label) > iou_threshold: |
|
non_overlap = False |
|
break |
|
|
|
if non_overlap: |
|
SelectedLabels.append(label) |
|
return SelectedLabels |
|
|
|
def load_model_wpod(path): |
|
model = load_model(path) |
|
return model |
|
|
|
def find_T_matrix(pts, t_pts): |
|
A = np.zeros((8, 9)) |
|
for i in range(0, 4): |
|
xi = pts[:, i] |
|
xil = t_pts[:, i] |
|
xi = xi.T |
|
|
|
A[i*2, 3:6] = -xil[2]*xi |
|
A[i*2, 6:] = xil[1]*xi |
|
A[i*2+1, :3] = xil[2]*xi |
|
A[i*2+1, 6:] = -xil[0]*xi |
|
|
|
[U, S, V] = np.linalg.svd(A) |
|
H = V[-1, :].reshape((3, 3)) |
|
return H |
|
|
|
def getRectPts(a, b): |
|
return np.array([[0,0], [a, 0], [a, b],[0,b]],np.float32) |
|
|
|
def normal(pts, side, mn, MN): |
|
pts_MN_center_mn = pts * side |
|
pts_MN = pts_MN_center_mn + mn.reshape((2, 1)) |
|
pts_prop = pts_MN / MN.reshape((2, 1)) |
|
return pts_prop |
|
def get_bound(x,y): |
|
bound =[] |
|
for i in range(0,len(x)): |
|
point =[x[i],y[i]] |
|
bound.append(point) |
|
return bound |
|
def calculate_ratio(bound): |
|
def distance(point1,point2): |
|
x1=point1[0] |
|
y1=point1[1] |
|
x2=point2[0] |
|
y2=point2[1] |
|
distance = np.sqrt((x2 - x1)**2 + (y2 - y1)**2) |
|
return distance |
|
box = bound |
|
dis1= distance(box[0],box[1]) |
|
dis2 = distance(box[1],box[2]) |
|
dis3 = distance(box[2],box[3]) |
|
dis4 = distance(box[3],box[0]) |
|
width = (dis1+dis3)/2 |
|
height= (dis2+dis4)/2 |
|
ratio = height/width |
|
if ratio>0.55: |
|
return 2 |
|
return 1 |
|
|
|
def reconstruct(I, Iresized, Yr, lp_threshold): |
|
bounds=[] |
|
|
|
net_stride = 2**4 |
|
side = ((208 + 40)/2)/net_stride |
|
|
|
|
|
one_line = (100, 23) |
|
two_lines = (64, 46) |
|
|
|
Probs = Yr[..., 0] |
|
Affines = Yr[..., 2:] |
|
|
|
xx, yy = np.where(Probs > lp_threshold) |
|
|
|
WH = getWH(Iresized.shape) |
|
|
|
MN = WH/net_stride |
|
|
|
vxx = vyy = 0.5 |
|
base = lambda vx, vy: np.matrix([[-vx, -vy, 1], [vx, -vy, 1], [vx, vy, 1], [-vx, vy, 1]]).T |
|
labels = [] |
|
labels_frontal = [] |
|
|
|
for i in range(len(xx)): |
|
x, y = xx[i], yy[i] |
|
affine = Affines[x, y] |
|
prob = Probs[x, y] |
|
|
|
mn = np.array([float(y) + 0.5, float(x) + 0.5]) |
|
|
|
|
|
A = np.reshape(affine, (2, 3)) |
|
A[0, 0] = max(A[0, 0], 0) |
|
A[1, 1] = max(A[1, 1], 0) |
|
|
|
B = np.zeros((2, 3)) |
|
B[0, 0] = max(A[0, 0], 0) |
|
B[1, 1] = max(A[1, 1], 0) |
|
|
|
pts = np.array(A*base(vxx, vyy)) |
|
pts_frontal = np.array(B*base(vxx, vyy)) |
|
|
|
pts_prop = normal(pts, side, mn, MN) |
|
frontal = normal(pts_frontal, side, mn, MN) |
|
|
|
labels.append(DLabel(0, pts_prop, prob)) |
|
labels_frontal.append(DLabel(0, frontal, prob)) |
|
|
|
final_labels = nms(labels, 0.1) |
|
final_labels_frontal = nms(labels_frontal, 0.1) |
|
if (len(final_labels_frontal)>0): |
|
|
|
|
|
|
|
|
|
lp_type=0 |
|
TLp = [] |
|
if len(final_labels): |
|
final_labels.sort(key=lambda x: x.prob(), reverse=True) |
|
for _, label in enumerate(final_labels): |
|
ptsh = np.concatenate((label.pts * getWH(I.shape).reshape((2, 1)), np.ones((1, 4)))) |
|
bound = get_bound(ptsh[0],ptsh[1]) |
|
pts=np.array(bound,dtype=np.float32) |
|
bounds.append(bound) |
|
lp_type = calculate_ratio(bound) |
|
if lp_type==2: |
|
out_size=two_lines |
|
else: out_size=one_line |
|
t_ptsh = getRectPts(out_size[0], out_size[1]) |
|
H=cv2.getPerspectiveTransform(pts,t_ptsh) |
|
Ilp = cv2.warpPerspective(I,H, (int(out_size[0]),int(out_size[1]))) |
|
TLp.append(Ilp) |
|
|
|
return final_labels, TLp, lp_type,bounds |
|
else: |
|
return None,[], None,None |
|
def detect_lp(model, I, lp_threshold): |
|
Dmax = 608 |
|
Dmin = 288 |
|
|
|
|
|
ratio = float(max(I.shape[:2])) / min(I.shape[:2]) |
|
side = int(ratio * Dmin) |
|
max_dim = min(side, Dmax) |
|
I=im2single(I) |
|
|
|
min_dim_img = min(I.shape[:2]) |
|
factor = float(max_dim) / min_dim_img |
|
|
|
|
|
w, h = (np.array(I.shape[1::-1], dtype=float) * factor).astype(int).tolist() |
|
|
|
|
|
Iresized = cv2.resize(I, (w, h)) |
|
|
|
T = Iresized.copy() |
|
|
|
|
|
T = T.reshape((1, T.shape[0], T.shape[1], T.shape[2])) |
|
|
|
|
|
Yr = model.predict(T,verbose=0) |
|
|
|
|
|
Yr = np.squeeze(Yr) |
|
|
|
|
|
|
|
|
|
L, TLp, lp_type,bounds = reconstruct(I, Iresized, Yr, lp_threshold) |
|
return L, TLp, lp_type,bounds |
|
|