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 | |