import cv2 import matplotlib.pyplot as plt import numpy as np from openvino.runtime import Core import os import time import copy from PIL import Image, ImageDraw, ImageFont import argparse def cv_imread(path): img=cv2.imdecode(np.fromfile(path,dtype=np.uint8),-1) return img def allFilePath(rootPath,allFIleList): fileList = os.listdir(rootPath) for temp in fileList: if os.path.isfile(os.path.join(rootPath,temp)): # if temp.endswith("jpg"): allFIleList.append(os.path.join(rootPath,temp)) else: allFilePath(os.path.join(rootPath,temp),allFIleList) mean_value,std_value=((0.588,0.193))#识别模型均值标准差 plateName=r"#京沪津渝冀晋蒙辽吉黑苏浙皖闽赣鲁豫鄂湘粤桂琼川贵云藏陕甘青宁新学警港澳挂使领民航危0123456789ABCDEFGHJKLMNPQRSTUVWXYZ险品" def rec_pre_precessing(img,size=(48,168)): #识别前处理 img =cv2.resize(img,(168,48)) img = img.astype(np.float32) img = (img/255-mean_value)/std_value img = img.transpose(2,0,1) img = img.reshape(1,*img.shape) return img def decodePlate(preds): #识别后处理 pre=0 newPreds=[] preds=preds.astype(np.int8)[0] for i in range(len(preds)): if preds[i]!=0 and preds[i]!=pre: newPreds.append(preds[i]) pre=preds[i] plate="" for i in newPreds: plate+=plateName[int(i)] return plate def load_model(onnx_path): ie = Core() model_onnx = ie.read_model(model=onnx_path) compiled_model_onnx = ie.compile_model(model=model_onnx, device_name="CPU") output_layer_onnx = compiled_model_onnx.output(0) return compiled_model_onnx,output_layer_onnx def get_plate_result(img,rec_model,rec_output): img =rec_pre_precessing(img) # time_b = time.time() res_onnx = rec_model([img])[rec_output] # time_e= time.time() index =np.argmax(res_onnx,axis=-1) #找出最大概率的那个字符的序号 plate_no = decodePlate(index) # print(f'{plate_no},time is {time_e-time_b}') return plate_no def get_split_merge(img): #双层车牌进行分割后识别 h,w,c = img.shape img_upper = img[0:int(5/12*h),:] img_lower = img[int(1/3*h):,:] img_upper = cv2.resize(img_upper,(img_lower.shape[1],img_lower.shape[0])) new_img = np.hstack((img_upper,img_lower)) return new_img 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) (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 the warped image return warped def my_letter_box(img,size=(640,640)): h,w,c = img.shape r = min(size[0]/h,size[1]/w) new_h,new_w = int(h*r),int(w*r) top = int((size[0]-new_h)/2) left = int((size[1]-new_w)/2) bottom = size[0]-new_h-top right = size[1]-new_w-left img_resize = cv2.resize(img,(new_w,new_h)) img = cv2.copyMakeBorder(img_resize,top,bottom,left,right,borderType=cv2.BORDER_CONSTANT,value=(114,114,114)) return img,r,left,top def xywh2xyxy(boxes): xywh =copy.deepcopy(boxes) xywh[:,0]=boxes[:,0]-boxes[:,2]/2 xywh[:,1]=boxes[:,1]-boxes[:,3]/2 xywh[:,2]=boxes[:,0]+boxes[:,2]/2 xywh[:,3]=boxes[:,1]+boxes[:,3]/2 return xywh def my_nms(boxes,iou_thresh): index = np.argsort(boxes[:,4])[::-1] keep = [] while index.size >0: i = index[0] keep.append(i) x1=np.maximum(boxes[i,0],boxes[index[1:],0]) y1=np.maximum(boxes[i,1],boxes[index[1:],1]) x2=np.minimum(boxes[i,2],boxes[index[1:],2]) y2=np.minimum(boxes[i,3],boxes[index[1:],3]) w = np.maximum(0,x2-x1) h = np.maximum(0,y2-y1) inter_area = w*h union_area = (boxes[i,2]-boxes[i,0])*(boxes[i,3]-boxes[i,1])+(boxes[index[1:],2]-boxes[index[1:],0])*(boxes[index[1:],3]-boxes[index[1:],1]) iou = inter_area/(union_area-inter_area) idx = np.where(iou<=iou_thresh)[0] index = index[idx+1] return keep def restore_box(boxes,r,left,top): boxes[:,[0,2,5,7,9,11]]-=left boxes[:,[1,3,6,8,10,12]]-=top boxes[:,[0,2,5,7,9,11]]/=r boxes[:,[1,3,6,8,10,12]]/=r return boxes def detect_pre_precessing(img,img_size): img,r,left,top=my_letter_box(img,img_size) # cv2.imwrite("1.jpg",img) img =img[:,:,::-1].transpose(2,0,1).copy().astype(np.float32) img=img/255 img=img.reshape(1,*img.shape) return img,r,left,top def post_precessing(dets,r,left,top,conf_thresh=0.3,iou_thresh=0.5):#检测后处理 choice = dets[:,:,4]>conf_thresh dets=dets[choice] dets[:,13:15]*=dets[:,4:5] box = dets[:,:4] boxes = xywh2xyxy(box) score= np.max(dets[:,13:15],axis=-1,keepdims=True) index = np.argmax(dets[:,13:15],axis=-1).reshape(-1,1) output = np.concatenate((boxes,score,dets[:,5:13],index),axis=1) reserve_=my_nms(output,iou_thresh) output=output[reserve_] output = restore_box(output,r,left,top) return output def rec_plate(outputs,img0,rec_model,rec_output): dict_list=[] for output in outputs: result_dict={} rect=output[:4].tolist() land_marks = output[5:13].reshape(4,2) roi_img = four_point_transform(img0,land_marks) label = int(output[-1]) if label==1: #代表是双层车牌 roi_img = get_split_merge(roi_img) plate_no = get_plate_result(roi_img,rec_model,rec_output) #得到车牌识别结果 result_dict['rect']=rect result_dict['landmarks']=land_marks.tolist() result_dict['plate_no']=plate_no result_dict['roi_height']=roi_img.shape[0] dict_list.append(result_dict) return dict_list def cv2ImgAddText(img, text, left, top, textColor=(0, 255, 0), textSize=20): if (isinstance(img, np.ndarray)): #判断是否OpenCV图片类型 img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) draw = ImageDraw.Draw(img) fontText = ImageFont.truetype( "fonts/platech.ttf", textSize, encoding="utf-8") draw.text((left, top), text, textColor, font=fontText) return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR) 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]=min(orgimg.shape[1],int(y-padding_h)) rect_area[2]=max(0,int(rect_area[2]+padding_w)) rect_area[3]=min(orgimg.shape[0],int(rect_area[3]+padding_h)) height_area = result['roi_height'] landmarks=result['landmarks'] result = result['plate_no'] result_str+=result+" " # for i in range(4): #关键点 # cv2.circle(orgimg, (int(landmarks[i][0]), int(landmarks[i][1])), 5, clors[i], -1) if len(result)>=6: cv2.rectangle(orgimg,(rect_area[0],rect_area[1]),(rect_area[2],rect_area[3]),(0,0,255),2) #画框 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 def get_second(capture): if capture.isOpened(): rate = capture.get(5) # 帧速率 FrameNumber = capture.get(7) # 视频文件的帧数 duration = FrameNumber/rate # 帧速率/视频总帧数 是时间，除以60之后单位是分钟 return int(rate),int(FrameNumber),int(duration) if __name__=="__main__": parser = argparse.ArgumentParser() parser.add_argument('--detect_model',type=str, default=r'weights/plate_detect.onnx', help='model.pt path(s)') #检测模型 parser.add_argument('--rec_model', type=str, default='weights/plate_rec.onnx', help='model.pt path(s)')#识别模型 parser.add_argument('--image_path', type=str, default='imgs', help='source') parser.add_argument('--img_size', type=int, default=640, help='inference size (pixels)') parser.add_argument('--output', type=str, default='result1', help='source') opt = parser.parse_args() file_list=[] file_folder=opt.image_path allFilePath(file_folder,file_list) rec_onnx_path =opt.rec_model detect_onnx_path=opt.detect_model rec_model,rec_output=load_model(rec_onnx_path) detect_model,detect_output=load_model(detect_onnx_path) count=0 img_size=(opt.img_size,opt.img_size) begin=time.time() save_path=opt.output if not os.path.exists(save_path): os.mkdir(save_path) for pic_ in file_list: count+=1 print(count,pic_,end=" ") img=cv2.imread(pic_) time_b = time.time() if img.shape[-1]==4: img = cv2.cvtColor(img,cv2.COLOR_BGRA2BGR) img0 = copy.deepcopy(img) img,r,left,top = detect_pre_precessing(img,img_size) #检测前处理 # print(img.shape) det_result = detect_model([img])[detect_output] outputs = post_precessing(det_result,r,left,top) #检测后处理 time_1 = time.time() result_list=rec_plate(outputs,img0,rec_model,rec_output) time_e= time.time() print(f'耗时 {time_e-time_b} s') ori_img = draw_result(img0,result_list) img_name = os.path.basename(pic_) save_img_path = os.path.join(save_path,img_name) cv2.imwrite(save_img_path,ori_img) print(f"总共耗时{time.time()-begin} s") # video_name = r"plate.mp4" # capture=cv2.VideoCapture(video_name) # fourcc = cv2.VideoWriter_fourcc(*'MP4V') # fps = capture.get(cv2.CAP_PROP_FPS) # 帧数 # width, height = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) # 宽高 # out = cv2.VideoWriter('2result.mp4', fourcc, fps, (width, height)) # 写入视频 # frame_count = 0 # fps_all=0 # rate,FrameNumber,duration=get_second(capture) # # with open("example.csv",mode='w',newline='') as example_file: # # fieldnames = ['车牌', '时间'] # # writer = csv.DictWriter(example_file, fieldnames=fieldnames, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL) # # writer.writeheader() # if capture.isOpened(): # while True: # t1 = cv2.getTickCount() # frame_count+=1 # ret,img=capture.read() # if not ret: # break # # if frame_count%rate==0: # img0 = copy.deepcopy(img) # img,r,left,top = detect_pre_precessing(img,img_size) #检测前处理 # # print(img.shape) # det_result = detect_model([img])[detect_output] # outputs = post_precessing(det_result,r,left,top) #检测后处理 # result_list=rec_plate(outputs,img0,rec_model,rec_output) # ori_img = draw_result(img0,result_list) # t2 =cv2.getTickCount() # infer_time =(t2-t1)/cv2.getTickFrequency() # fps=1.0/infer_time # fps_all+=fps # str_fps = f'fps:{fps:.4f}' # out.write(ori_img) # cv2.putText(ori_img,str_fps,(20,20),cv2.FONT_HERSHEY_SIMPLEX,1,(0,255,0),2) # cv2.imshow("haha",ori_img) # cv2.waitKey(1) # # current_time = int(frame_count/FrameNumber*duration) # # sec = current_time%60 # # minute = current_time//60 # # for result_ in result_list: # # plate_no = result_['plate_no'] # # if not is_car_number(pattern_str,plate_no): # # continue # # print(f'车牌号:{plate_no},时间:{minute}分{sec}秒') # # time_str =f'{minute}分{sec}秒' # # writer.writerow({"车牌":plate_no,"时间":time_str}) # # out.write(ori_img) # else: # print("失败") # capture.release() # out.release() # cv2.destroyAllWindows() # print(f"all frame is {frame_count},average fps is {fps_all/frame_count}")