Spaces:
Runtime error
Runtime error
import cv2 | |
import numpy as np | |
import os | |
def getPath(test_image_root = "./TDCN_IMG/"): | |
path_list={} | |
for i,angle_path in enumerate(os.listdir(test_image_root)): | |
angle = os.path.join(test_image_root,angle_path) | |
path_list[angle_path]=[] | |
for img in os.listdir(angle): | |
path_= os.path.join(angle,img) | |
path_list[angle_path].append([img.strip('.jpeg'),path_]) | |
return path_list | |
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 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 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(tlx, tly, brx, bry): | |
return np.matrix([[tlx, brx, brx, tlx], [tly, tly, bry, bry], [1, 1, 1, 1]], dtype=float) | |
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 | |
# Reconstruction function from predict value into plate crpoped from image | |
def reconstruct(I, Iresized, Yr, lp_threshold): | |
# 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 = (470, 110) | |
two_lines = (280, 200) | |
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) | |
#print(final_labels_frontal) | |
assert final_labels_frontal, "No License plate is founded!" | |
# LP size and type | |
out_size, lp_type = (two_lines, 2) if ((final_labels_frontal[0].wh()[0] / final_labels_frontal[0].wh()[1]) < 1.7) else (one_line, 1) | |
TLp = [] | |
Cor = [] | |
if len(final_labels): | |
final_labels.sort(key=lambda x: x.prob(), reverse=True) | |
for _, label in enumerate(final_labels): | |
t_ptsh = getRectPts(0, 0, out_size[0], out_size[1]) | |
ptsh = np.concatenate((label.pts * getWH(I.shape).reshape((2, 1)), np.ones((1, 4)))) | |
H = find_T_matrix(ptsh, t_ptsh) | |
Ilp = cv2.warpPerspective(I, H, out_size, borderValue=0) | |
TLp.append(Ilp) | |
Cor.append(ptsh) | |
return final_labels, TLp, lp_type, Cor | |
def detect_lp(model, I, max_dim, lp_threshold): | |
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) | |
Yr = np.squeeze(Yr) | |
#print(Yr.shape) | |
L, TLp, lp_type, Cor = reconstruct(I, Iresized, Yr, lp_threshold) | |
return L, TLp, lp_type, Cor | |