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