Upload local_utils.py

local_utils.py

@@ -0,0 +1,230 @@
+
+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