-
Notifications
You must be signed in to change notification settings - Fork 75
/
detect_rec_plate.py
190 lines (171 loc) · 8.22 KB
/
detect_rec_plate.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
import argparse
import time
import os
import copy
import cv2
import torch
import numpy as np
import torch.backends.cudnn as cudnn
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from plate_recognition.plate_rec import get_plate_result,allFilePath,init_model,cv_imread
from plate_recognition.double_plate_split_merge import get_split_merge
from utils.datasets import letterbox
from utils.cv_puttext import cv2ImgAddText
def cv_imread(path):
img=cv2.imdecode(np.fromfile(path,dtype=np.uint8),-1)
return img
clors = [(255,0,0),(0,255,0),(0,0,255),(255,255,0),(0,255,255)]
def order_points(pts): #关键点按照（左上，右上，右下，左下）排列
rect = np.zeros((4, 2), dtype = "float32")
s = pts.sum(axis = 1)
rect[0] = pts[np.argmin(s)]
rect[2] = pts[np.argmax(s)]
diff = np.diff(pts, axis = 1)
rect[1] = pts[np.argmin(diff)]
rect[3] = pts[np.argmax(diff)]
return rect
def four_point_transform(image, pts): #透视变换
# rect = order_points(pts)
rect=pts.astype("float32")
(tl, tr, br, bl) = rect
widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
maxWidth = max(int(widthA), int(widthB))
heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
maxHeight = max(int(heightA), int(heightB))
dst = np.array([
[0, 0],
[maxWidth - 1, 0],
[maxWidth - 1, maxHeight - 1],
[0, maxHeight - 1]], dtype = "float32")
M = cv2.getPerspectiveTransform(rect, dst)
warped = cv2.warpPerspective(image, M, (maxWidth, maxHeight))
return warped
def get_plate_rec_landmark(img, xyxy, conf, landmarks, class_num,device,plate_rec_model):
h,w,c = img.shape
result_dict={}
tl = 1 or round(0.002 * (h + w) / 2) + 1 # line/font thickness
x1 = int(xyxy[0])
y1 = int(xyxy[1])
x2 = int(xyxy[2])
y2 = int(xyxy[3])
height=y2-y1
landmarks_np=np.zeros((4,2))
rect=[x1,y1,x2,y2]
for i in range(4):
point_x = int(landmarks[2 * i])
point_y = int(landmarks[2 * i + 1])
landmarks_np[i]=np.array([point_x,point_y])
class_label= int(class_num) #车牌的的类型0代表单牌，1代表双层车牌
roi_img = four_point_transform(img,landmarks_np) #透视变换得到车牌小图
# cv2.imwrite("roi.jpg",roi_img)
# roi_img_h = roi_img.shape[0]
# roi_img_w = roi_img.shape[1]
# if roi_img_w/roi_img_h<3:
# class_label=
# h_w_r = roi_img_w/roi_img_h
if class_label : #判断是否是双层车牌，是双牌的话进行分割后然后拼接
roi_img=get_split_merge(roi_img)
plate_number,rec_prob,plate_color,color_conf = get_plate_result(roi_img,device,plate_rec_model) #对车牌小图进行识别
result_dict['rect']=rect
result_dict['landmarks']=landmarks_np.tolist()
result_dict['plate_no']=plate_number
result_dict['rec_conf']=rec_prob #每个字符的概率
result_dict['plate_color']=plate_color
result_dict['color_conf']=color_conf
result_dict['roi_height']=roi_img.shape[0]
result_dict['score']=conf
result_dict['label']=class_label
return result_dict
def detect_Recognition_plate(model, orgimg, device,plate_rec_model,img_size):
conf_thres = 0.3
iou_thres = 0.5
dict_list=[]
im0 = copy.deepcopy(orgimg)
imgsz=(img_size,img_size)
img = letterbox(im0, new_shape=imgsz)[0]
img = img[:, :, ::-1].transpose(2, 0, 1).copy() # BGR to RGB, to 3x640X640
img = torch.from_numpy(img).to(device)
img = img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
pred = model(img)[0]
pred = non_max_suppression(pred, conf_thres=conf_thres, iou_thres=iou_thres, kpt_label=4,agnostic=True)
for i, det in enumerate(pred):
if len(det):
# Rescale boxes from img_size to im0 size
scale_coords(img.shape[2:], det[:, :4], im0.shape, kpt_label=False)
scale_coords(img.shape[2:], det[:, 6:], im0.shape, kpt_label=4, step=3)
for j in range(det.size()[0]):
xyxy = det[j, :4].view(-1).tolist()
conf = det[j, 4].cpu().numpy()
landmarks = det[j, 6:].view(-1).tolist()
landmarks = [landmarks[0],landmarks[1],landmarks[3],landmarks[4],landmarks[6],landmarks[7],landmarks[9],landmarks[10]]
class_num = det[j, 5].cpu().numpy()
result_dict = get_plate_rec_landmark(orgimg, xyxy, conf, landmarks, class_num,device,plate_rec_model)
dict_list.append(result_dict)
return dict_list
def draw_result(orgimg,dict_list):
result_str =""
for result in dict_list:
rect_area = result['rect']
x,y,w,h = rect_area[0],rect_area[1],rect_area[2]-rect_area[0],rect_area[3]-rect_area[1]
padding_w = 0.05*w
padding_h = 0.11*h
rect_area[0]=max(0,int(x-padding_w))
rect_area[1]=max(0,int(y-padding_h))
rect_area[2]=min(orgimg.shape[1],int(rect_area[2]+padding_w))
rect_area[3]=min(orgimg.shape[0],int(rect_area[3]+padding_h))
rect_area = [int(x) for x in rect_area]
height_area = result['roi_height']
landmarks=result['landmarks']
result_p = result['plate_no']+" "+result['plate_color']
result_str+=result_p+" "
cv2.rectangle(orgimg,(rect_area[0],rect_area[1]),(rect_area[2],rect_area[3]),(0,0,255),2) #画框
labelSize = cv2.getTextSize(result_p,cv2.FONT_HERSHEY_SIMPLEX,0.5,1)
if rect_area[0]+labelSize[0][0]>orgimg.shape[1]: #防止显示的文字越界
rect_area[0]=int(orgimg.shape[1]-labelSize[0][0])
orgimg=cv2.rectangle(orgimg,(rect_area[0],int(rect_area[1]-round(1.6*labelSize[0][1]))),(int(rect_area[0]+round(1.2*labelSize[0][0])),rect_area[1]+labelSize[1]),(255,255,255),cv2.FILLED)
if len(result)>1:
for i in range(4): #关键点
cv2.circle(orgimg, (int(landmarks[i][0]), int(landmarks[i][1])), 5, clors[i], -1)
orgimg=cv2ImgAddText(orgimg,result_p,rect_area[0],int(rect_area[1]-round(1.6*labelSize[0][1])),(0,0,0),21)
# orgimg=cv2ImgAddText(orgimg,result,rect_area[0]-height_area,rect_area[1]-height_area-10,(0,255,0),height_area)
print(result_str)
return orgimg
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--detect_model', nargs='+', type=str, default='weights/yolov7-lite-s.pt', help='model.pt path(s)') #检测模型
parser.add_argument('--rec_model', type=str, default='weights/plate_rec_color.pth', help='model.pt path(s)') #车牌识别 +颜色识别
parser.add_argument('--source', type=str, default='imgs', help='source') # file/folder, 0 for webcam
# parser.add_argument('--img-size', nargs= '+', type=int, default=640, help='inference size (pixels)')
parser.add_argument('--img_size', type=int, default=640, help='inference size (pixels)')
parser.add_argument('--output', type=str, default='result', help='source')
parser.add_argument('--kpt-label', type=int, default=4, help='number of keypoints')
device =torch.device("cuda" if torch.cuda.is_available() else "cpu")
# device = torch.device("cpu")
opt = parser.parse_args()
print(opt)
model = attempt_load(opt.detect_model, map_location=device)
# torch.save()
plate_rec_model=init_model(device,opt.rec_model)
if not os.path.exists(opt.output):
os.mkdir(opt.output)
file_list=[]
allFilePath(opt.source,file_list)
time_b = time.time()
for pic_ in file_list:
print(pic_,end=" ")
img = cv_imread(pic_)
if img.shape[-1]==4:
img=cv2.cvtColor(img,cv2.COLOR_BGRA2BGR)
# img = my_letter_box(img)
dict_list=detect_Recognition_plate(model, img, device,plate_rec_model,opt.img_size)
ori_img=draw_result(img,dict_list)
img_name = os.path.basename(pic_)
save_img_path = os.path.join(opt.output,img_name)
cv2.imwrite(save_img_path,ori_img)
print(f"elasted time is {time.time()-time_b} s")