windy2612's picture
Update wpodnet/lib_detection.py
2095deb verified
# pylint: disable=invalid-name, redefined-outer-name, missing-docstring, non-parent-init-called, trailing-whitespace, line-too-long
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)
# Hàm normalize ảnh
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
# Hàm tái tạo từ predict value thành biến số, cắt từ ảnh chính ra biển số, nhãn...
def reconstruct(I, Iresized, Yr, lp_threshold):
bounds=[]
# 4 max-pooling layers, stride = 2
net_stride = 2**4
side = ((208 + 40)/2)/net_stride
# one line and two lines license plate size
one_line = (100, 23)
two_lines = (64, 46)
Probs = Yr[..., 0]
Affines = Yr[..., 2:]
xx, yy = np.where(Probs > lp_threshold)
# CNN input image size
WH = getWH(Iresized.shape)
# output feature map size
MN = WH/net_stride
vxx = vyy = 0.5 #alpha
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])
# affine transformation matrix
A = np.reshape(affine, (2, 3))
A[0, 0] = max(A[0, 0], 0)
A[1, 1] = max(A[1, 1], 0)
# identity transformation
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 size and type
#out_size, lp_type = (two_lines, 2) if ((final_labels_frontal[0].wh()[1] / final_labels_frontal[0].wh()[1]) >0.49) else (one_line, 1)
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
# Lấy tỷ lệ giữa W và H của ảnh và tìm ra chiều nhỏ nhất
ratio = float(max(I.shape[:2])) / min(I.shape[:2])
side = int(ratio * Dmin)
max_dim = min(side, Dmax)
I=im2single(I)
# Tính factor resize ảnh
min_dim_img = min(I.shape[:2])
factor = float(max_dim) / min_dim_img
# Tính W và H mới sau khi resize
w, h = (np.array(I.shape[1::-1], dtype=float) * factor).astype(int).tolist()
# Tiến hành resize ảnh
Iresized = cv2.resize(I, (w, h))
T = Iresized.copy()
# Chuyển thành Tensor
T = T.reshape((1, T.shape[0], T.shape[1], T.shape[2]))
# Tiến hành detect biển số bằng Wpod-net pretrain
Yr = model.predict(T,verbose=0)
# Remove các chiều =1 của Yr
Yr = np.squeeze(Yr)
# Tái tạo và trả về các biến gồm: Nhãn, Ảnh biến số, Loại biển số (1: dài: 2 vuông)
L, TLp, lp_type,bounds = reconstruct(I, Iresized, Yr, lp_threshold)
return L, TLp, lp_type,bounds