RivianG commited on
Commit
e338936
·
1 Parent(s): 9753465

Upload local_utils.py

Browse files
Files changed (1) hide show
  1. local_utils.py +221 -0
local_utils.py ADDED
@@ -0,0 +1,221 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # pylint: disable=invalid-name, redefined-outer-name, missing-docstring, non-parent-init-called, trailing-whitespace, line-too-long
2
+ import cv2
3
+ import numpy as np
4
+ import tensorflow as tf
5
+ from tensorflow.python.keras.backend import set_session
6
+
7
+
8
+ class Label:
9
+ def __init__(self, cl=-1, tl=np.array([0., 0.]), br=np.array([0., 0.]), prob=None):
10
+ self.__tl = tl
11
+ self.__br = br
12
+ self.__cl = cl
13
+ self.__prob = prob
14
+
15
+ def __str__(self):
16
+ return 'Class: %d, top left(x: %f, y: %f), bottom right(x: %f, y: %f)' % (
17
+ self.__cl, self.__tl[0], self.__tl[1], self.__br[0], self.__br[1])
18
+
19
+ def copy(self):
20
+ return Label(self.__cl, self.__tl, self.__br)
21
+
22
+ def wh(self): return self.__br - self.__tl
23
+
24
+ def cc(self): return self.__tl + self.wh() / 2
25
+
26
+ def tl(self): return self.__tl
27
+
28
+ def br(self): return self.__br
29
+
30
+ def tr(self): return np.array([self.__br[0], self.__tl[1]])
31
+
32
+ def bl(self): return np.array([self.__tl[0], self.__br[1]])
33
+
34
+ def cl(self): return self.__cl
35
+
36
+ def area(self): return np.prod(self.wh())
37
+
38
+ def prob(self): return self.__prob
39
+
40
+ def set_class(self, cl):
41
+ self.__cl = cl
42
+
43
+ def set_tl(self, tl):
44
+ self.__tl = tl
45
+
46
+ def set_br(self, br):
47
+ self.__br = br
48
+
49
+ def set_wh(self, wh):
50
+ cc = self.cc()
51
+ self.__tl = cc - .5 * wh
52
+ self.__br = cc + .5 * wh
53
+
54
+ def set_prob(self, prob):
55
+ self.__prob = prob
56
+
57
+
58
+ class DLabel(Label):
59
+ def __init__(self, cl, pts, prob):
60
+ self.pts = pts
61
+ tl = np.amin(pts, axis=1)
62
+ br = np.amax(pts, axis=1)
63
+ Label.__init__(self, cl, tl, br, prob)
64
+
65
+
66
+ def getWH(shape):
67
+ return np.array(shape[1::-1]).astype(float)
68
+
69
+
70
+ def IOU(tl1, br1, tl2, br2):
71
+ wh1, wh2 = br1 - tl1, br2 - tl2
72
+ assert ((wh1 >= 0).all() and (wh2 >= 0).all())
73
+
74
+ intersection_wh = np.maximum(np.minimum(br1, br2) - np.maximum(tl1, tl2), 0)
75
+ intersection_area = np.prod(intersection_wh)
76
+ area1, area2 = (np.prod(wh1), np.prod(wh2))
77
+ union_area = area1 + area2 - intersection_area
78
+ return intersection_area / union_area
79
+
80
+
81
+ def IOU_labels(l1, l2):
82
+ return IOU(l1.tl(), l1.br(), l2.tl(), l2.br())
83
+
84
+
85
+ def nms(Labels, iou_threshold=0.5):
86
+ SelectedLabels = []
87
+ Labels.sort(key=lambda l: l.prob(), reverse=True)
88
+
89
+ for label in Labels:
90
+ non_overlap = True
91
+ for sel_label in SelectedLabels:
92
+ if IOU_labels(label, sel_label) > iou_threshold:
93
+ non_overlap = False
94
+ break
95
+
96
+ if non_overlap:
97
+ SelectedLabels.append(label)
98
+ return SelectedLabels
99
+
100
+
101
+ def find_T_matrix(pts, t_pts):
102
+ A = np.zeros((8, 9))
103
+ for i in range(0, 4):
104
+ xi = pts[:, i]
105
+ xil = t_pts[:, i]
106
+ xi = xi.T
107
+
108
+ A[i * 2, 3:6] = -xil[2] * xi
109
+ A[i * 2, 6:] = xil[1] * xi
110
+ A[i * 2 + 1, :3] = xil[2] * xi
111
+ A[i * 2 + 1, 6:] = -xil[0] * xi
112
+
113
+ [U, S, V] = np.linalg.svd(A)
114
+ H = V[-1, :].reshape((3, 3))
115
+ return H
116
+
117
+
118
+ def getRectPts(tlx, tly, brx, bry):
119
+ return np.matrix([[tlx, brx, brx, tlx], [tly, tly, bry, bry], [1, 1, 1, 1]], dtype=float)
120
+
121
+
122
+ def normal(pts, side, mn, MN):
123
+ pts_MN_center_mn = pts * side
124
+ pts_MN = pts_MN_center_mn + mn.reshape((2, 1))
125
+ pts_prop = pts_MN / MN.reshape((2, 1))
126
+ return pts_prop
127
+
128
+
129
+ # Reconstruction function from predict value into plate crpoped from image
130
+ def reconstruct(I, Iresized, Yr, lp_threshold):
131
+ # 4 max-pooling layers, stride = 2
132
+ net_stride = 2 ** 4
133
+ side = ((208 + 40) / 2) / net_stride
134
+
135
+ # one line and two lines license plate size
136
+ one_line = (470, 110)
137
+ two_lines = (280, 200)
138
+
139
+ Probs = Yr[..., 0]
140
+ Affines = Yr[..., 2:]
141
+
142
+ xx, yy = np.where(Probs > lp_threshold)
143
+ # CNN input image size
144
+ WH = getWH(Iresized.shape)
145
+ # output feature map size
146
+ MN = WH / net_stride
147
+
148
+ vxx = vyy = 0.5 # alpha
149
+ base = lambda vx, vy: np.matrix([[-vx, -vy, 1], [vx, -vy, 1], [vx, vy, 1], [-vx, vy, 1]]).T
150
+ labels = []
151
+ labels_frontal = []
152
+
153
+ for i in range(len(xx)):
154
+ x, y = xx[i], yy[i]
155
+ affine = Affines[x, y]
156
+ prob = Probs[x, y]
157
+
158
+ mn = np.array([float(y) + 0.5, float(x) + 0.5])
159
+
160
+ # affine transformation matrix
161
+ A = np.reshape(affine, (2, 3))
162
+ A[0, 0] = max(A[0, 0], 0)
163
+ A[1, 1] = max(A[1, 1], 0)
164
+ # identity transformation
165
+ B = np.zeros((2, 3))
166
+ B[0, 0] = max(A[0, 0], 0)
167
+ B[1, 1] = max(A[1, 1], 0)
168
+
169
+ pts = np.array(A * base(vxx, vyy))
170
+ pts_frontal = np.array(B * base(vxx, vyy))
171
+
172
+ pts_prop = normal(pts, side, mn, MN)
173
+ frontal = normal(pts_frontal, side, mn, MN)
174
+
175
+ labels.append(DLabel(0, pts_prop, prob))
176
+ labels_frontal.append(DLabel(0, frontal, prob))
177
+
178
+ final_labels = nms(labels, 0.1)
179
+ final_labels_frontal = nms(labels_frontal, 0.1)
180
+
181
+ # print(final_labels_frontal)
182
+ assert final_labels_frontal, "" # "No License plate is founded!"
183
+
184
+ # LP size and type
185
+ out_size, lp_type = (two_lines, 2) if (
186
+ (final_labels_frontal[0].wh()[0] / final_labels_frontal[0].wh()[1]) < 1.7) else (one_line, 1)
187
+
188
+ TLp = []
189
+ Cor = []
190
+ if len(final_labels):
191
+ final_labels.sort(key=lambda x: x.prob(), reverse=True)
192
+ for _, label in enumerate(final_labels):
193
+ t_ptsh = getRectPts(0, 0, out_size[0], out_size[1])
194
+ ptsh = np.concatenate((label.pts * getWH(I.shape).reshape((2, 1)), np.ones((1, 4))))
195
+ H = find_T_matrix(ptsh, t_ptsh)
196
+ Ilp = cv2.warpPerspective(I, H, out_size, borderValue=0)
197
+ # cv2.imshow("plate", Ilp)
198
+ # cv2.waitKey(0)
199
+ TLp.append(Ilp)
200
+ Cor.append(ptsh)
201
+ return final_labels, TLp, lp_type, Cor
202
+
203
+
204
+ def detect_lp(graph, sess, model, I, max_dim, lp_threshold):
205
+ min_dim_img = min(I.shape[:2])
206
+ factor = float(max_dim) / min_dim_img
207
+ w, h = (np.array(I.shape[1::-1], dtype=float) * factor).astype(int).tolist()
208
+ Iresized = cv2.resize(I, (w, h))
209
+ T = Iresized.copy()
210
+ T = T.reshape((1, T.shape[0], T.shape[1], T.shape[2]))
211
+ with graph.as_default():
212
+ set_session(sess)
213
+
214
+
215
+ Yr = model.predict(T)
216
+ # print("1: ",Yr)
217
+ Yr = np.squeeze(Yr)
218
+ # print("2: ",Yr)
219
+ # print(Yr.shape)
220
+ L, TLp, lp_type, Cor = reconstruct(I, Iresized, Yr, lp_threshold)
221
+ return L, TLp, lp_type, Cor