import argparse import copy import os import time from pprint import pprint import cv2 import numpy as np import torch from models.experimental import attempt_load from plate_recognition.double_plate_split_merge import get_split_merge from plate_recognition.plate_rec import ( allFilePath, cv_imread, get_plate_result, init_model, ) from utils.cv_puttext import cv2ImgAddText from utils.datasets import letterbox from utils.general import check_img_size, non_max_suppression_face, scale_coords clors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (0, 255, 255)] danger = ["危", "险"] 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 = 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 load_model(weights, device): # 加载检测模型 model = attempt_load(weights, map_location=device) # FP32 model return model def scale_coords_landmarks( img1_shape, coords, img0_shape, ratio_pad=None ): # 返回到原图坐标 if ratio_pad is None: # calculate from img0_shape gain = min( img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1] ) # gain = old / new pad = ( (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2, ) # wh padding else: gain = ratio_pad[0][0] pad = ratio_pad[1] coords[:, [0, 2, 4, 6]] -= pad[0] # x padding coords[:, [1, 3, 5, 7]] -= pad[1] # y padding coords[:, :8] /= gain coords[:, 0].clamp_(0, img0_shape[1]) # x1 coords[:, 1].clamp_(0, img0_shape[0]) # y1 coords[:, 2].clamp_(0, img0_shape[1]) # x2 coords[:, 3].clamp_(0, img0_shape[0]) # y2 coords[:, 4].clamp_(0, img0_shape[1]) # x3 coords[:, 5].clamp_(0, img0_shape[0]) # y3 coords[:, 6].clamp_(0, img0_shape[1]) # x4 coords[:, 7].clamp_(0, img0_shape[0]) # y4 # coords[:, 8].clamp_(0, img0_shape[1]) # x5 # coords[:, 9].clamp_(0, img0_shape[0]) # y5 return coords def get_plate_rec_landmark( img, xyxy, conf, landmarks, class_num, device, plate_rec_model, is_color=False ): # 获取车牌坐标以及四个角点坐标并识别车牌号 h, w, c = img.shape result_dict = {} 1 or round(0.002 * (h + w) / 2) + 1 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) # 透视变换得到车牌小图 if class_label: # 判断是否是双层车牌，是双牌的话进行分割后然后拼接 roi_img = get_split_merge(roi_img) if not is_color: plate_number, rec_prob = get_plate_result( roi_img, device, plate_rec_model, is_color=is_color ) # 对车牌小图进行识别 else: plate_number, rec_prob, plate_color, color_conf = get_plate_result( roi_img, device, plate_rec_model, is_color=is_color ) # cv2.imwrite("roi.jpg",roi_img) result_dict["rect"] = rect # 车牌roi区域 result_dict["detect_conf"] = conf # 检测区域置信度 result_dict["landmarks"] = landmarks_np.tolist() # 车牌角点坐标 result_dict["plate_no"] = plate_number # 车牌号 result_dict["rec_conf"] = rec_prob # 每个字符的概率 result_dict["roi_height"] = roi_img.shape[0] # 车牌高度 result_dict["plate_color"] = "" if is_color: result_dict["plate_color"] = plate_color # 车牌颜色 result_dict["color_conf"] = color_conf # 颜色置信度 result_dict["plate_type"] = class_label # 单双层 0单层 1双层 return result_dict def detect_Recognition_plate( model, orgimg, device, plate_rec_model, img_size, is_color=False ): # 获取车牌信息 # img_size = opt_img_size conf_thres = 0.3 ##### 置信度阈值 ##### iou_thres = 0.5 # nms的iou值 dict_list = [] img0 = copy.deepcopy(orgimg) assert orgimg is not None, "Image Not Found " h0, w0 = orgimg.shape[:2] r = img_size / max(h0, w0) if r != 1: interp = cv2.INTER_AREA if r < 1 else cv2.INTER_LINEAR img0 = cv2.resize(img0, (int(w0 * r), int(h0 * r)), interpolation=interp) imgsz = check_img_size(img_size, s=model.stride.max()) # 检查 img_size img = letterbox(img0, new_shape=imgsz)[ 0 ] # 检测前处理，图片长宽变为32倍数 img = ( img[:, :, ::-1].transpose(2, 0, 1).copy() ) # 图片的BGR排列转为RGB,然后将图片的H,W,C排列变为C,H,W排列 t0 = time.time() img = torch.from_numpy(img).to(device) img = img.float() 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_face(pred, conf_thres, iou_thres) # 检测进程 for i, det in enumerate(pred): # 对每张图片遍历 if len(det): det[:, :4] = scale_coords(img.shape[2:], det[:, :4], orgimg.shape).round() # 打印结果 for c in det[:, -1].unique(): n = (det[:, -1] == c).sum() # 检测每个种类 det[:, 5:13] = scale_coords_landmarks( img.shape[2:], det[:, 5:13], orgimg.shape ).round() for j in range(det.size()[0]): xyxy = det[j, :4].view(-1).tolist() conf = det[j, 4].cpu().numpy() landmarks = det[j, 5:13].view(-1).tolist() class_num = det[j, 13].cpu().numpy() result_dict = get_plate_rec_landmark( orgimg, xyxy, conf, landmarks, class_num, device, plate_rec_model, is_color=is_color, ) dict_list.append(result_dict) return dict_list # cv2.imwrite('result.jpg', orgimg) def draw_result(orgimg, dict_list, is_color=True): # 将车牌结果画出 result_str = "" if dict_list: pprint(dict_list) 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)) height_area = result["roi_height"] landmarks = result["landmarks"] result_p = result["plate_no"] if result["plate_type"] == 0: # 单层 result_p += " " + result["plate_color"] else: # 双层 result_p += " " + result["plate_color"] + "双层" result_str += result_p + " " for i in range(4): # 关键点 cv2.circle( orgimg, (int(landmarks[i][0]), int(landmarks[i][1])), 5, clors[i], -1 ) 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])-30)), ( int(rect_area[0] + round(1.2 * labelSize[0][0]))+50, rect_area[1] + labelSize[1], ), (255, 255, 255), cv2.FILLED , ) # 画文字框,背景白色 if len(result) >= 1: orgimg = cv2ImgAddText( orgimg, result_p, rect_area[0], int(rect_area[1] - round(1.6 * labelSize[0][1]))-30, (0, 0, 0), 30, ) if result_str: print(result_str) return orgimg def get_second(capture): if capture.isOpened(): rate = capture.get(5) # 帧速率 FrameNumber = capture.get(7) # 视频文件的帧数 duration = FrameNumber / rate return int(rate), int(FrameNumber), int(duration) def start(image_path="imgs"): # 测试图片路径 parser = argparse.ArgumentParser() parser.add_argument( "--detect_model", nargs="+", type=str, default="weights/plate_detect.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( "--is_color", type=bool, default=True, help="plate color" ) # 识别颜色 parser.add_argument( "--image_path", type=str, default=image_path, help="source" ) # 图片路径 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("--video", type=str, default="", help="source") # 视频的路径 parser.add_argument( "--stream", type=str, default="", help="RTSP/RTMP video stream URL" ) # 视频流地址 device = torch.device( "cuda" if torch.cuda.is_available() else "cpu" ) # device =torch.device("cpu") opt = parser.parse_args() print(opt) save_path = opt.output count = 0 if not os.path.exists(save_path): os.mkdir(save_path) detect_model = load_model( opt.detect_model, device ) # 初始化检测模型 plate_rec_model = init_model( device, opt.rec_model, is_color=opt.is_color ) # 初始化识别模型 # 计算参数量 total = sum(p.numel() for p in detect_model.parameters()) total_1 = sum(p.numel() for p in plate_rec_model.parameters()) print("detect params: %.2fM,rec params: %.2fM" % (total / 1e6, total_1 / 1e6)) # plate_color_model =init_color_model(opt.color_model,device) time_all = 0 time_begin = time.time() # 处理视频流 if opt.stream: cap = cv2.VideoCapture(opt.stream) if not cap.isOpened(): print(f"无法打开视频流: {opt.stream}") return fps = cap.get(cv2.CAP_PROP_FPS) or 25 width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) out = cv2.VideoWriter("stream_result.mp4", cv2.VideoWriter_fourcc(*"MP4V"), fps, (width, height)) frame_count = 0 fps_all = 0 print(f"开始处理视频流: {opt.stream}") while True: t1 = cv2.getTickCount() frame_count += 1 ret, img = cap.read() if not ret: print("视频流读取结束或出错") break dict_list = detect_Recognition_plate( detect_model, img, device, plate_rec_model, opt.img_size, is_color=opt.is_color, ) ori_img = draw_result(img, dict_list) t2 = cv2.getTickCount() infer_time = (t2 - t1) / cv2.getTickFrequency() fps = 1.0 / infer_time fps_all += fps str_fps = f"fps:{fps:.2f}" cv2.putText(ori_img, str_fps, (20, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2) cv2.imshow("Stream Result", ori_img) out.write(ori_img) if cv2.waitKey(1) & 0xFF == ord('q'): break cap.release() out.release() cv2.destroyAllWindows() print(f"总帧数: {frame_count}, 平均FPS: {fps_all / frame_count:.2f}") # 处理本地视频 elif opt.video: video_name = opt.video 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("result.mp4", fourcc, fps, (width, height)) # 写入视频 frame_count = 0 fps_all = 0 rate, FrameNumber, duration = get_second(capture) if capture.isOpened(): while True: t1 = cv2.getTickCount() frame_count += 1 print(f"第{frame_count} 帧", end=" ") ret, img = capture.read() if not ret: break # if frame_count%rate==0: img0 = copy.deepcopy(img) dict_list = detect_Recognition_plate( detect_model, img, device, plate_rec_model, opt.img_size, is_color=opt.is_color, ) ori_img = draw_result(img, dict_list) t2 = cv2.getTickCount() infer_time = (t2 - t1) / cv2.getTickFrequency() fps = 1.0 / infer_time fps_all += fps str_fps = f"fps:{fps:.4f}" cv2.putText( ori_img, str_fps, (20, 20), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, ) cv2.imshow("haha", ori_img) cv2.waitKey(0) out.write(ori_img) # 处理图片 else: if not os.path.isfile(opt.image_path): # 目录 file_list = [] allFilePath( opt.image_path, file_list ) # 将目录下的所有图片文件路径读取到file_list里面 for img_path in file_list: # 遍历图片文件 print(count, img_path, end=" ") time_b = time.time() # 开始时间 img = cv_imread(img_path) # opencv 读取图片 if img is None: continue if img.shape[-1] == 4: # 图片如果是4个通道的，将其转为3个通道 img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR) # detect_one(model,img_path,device) dict_list = detect_Recognition_plate( detect_model, img, device, plate_rec_model, opt.img_size, is_color=opt.is_color, ) # 检测以及识别车牌 pprint(dict_list) ori_img = draw_result(img, dict_list) # 将结果画在图上 img_name = os.path.basename(img_path) save_img_path = os.path.join(save_path, img_name) # 图片保存的路径 time_e = time.time() time_gap = time_e - time_b # 计算单个图片识别耗时 if count: time_all += time_gap if isinstance(ori_img, cv2.UMat): ori_img = cv2.UMat.get(ori_img) cv2.imwrite(save_img_path, ori_img) # opencv将识别的图片保存 count += 1 # cv2.namedWindow("result", cv2.WINDOW_NORMAL) # cv2.resizeWindow("result", 800, 600) cv2.imshow("result", ori_img) cv2.waitKey(0) cv2.destroyAllWindows() print( f"sumTime time is {time.time() - time_begin} s, average pic time is {time_all / (len(file_list) - 1)}" ) else: # 单个图片 print(count, opt.image_path, end=" ") img = cv_imread(opt.image_path) if img.shape[-1] == 4: img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR) # detect_one(model,img_path,device) dict_list = detect_Recognition_plate( detect_model, img, device, plate_rec_model, opt.img_size, is_color=opt.is_color, ) ori_img = draw_result(img, dict_list) img_name = os.path.basename(opt.image_path) save_img_path = os.path.join(save_path, img_name) cv2.imwrite(save_img_path, ori_img) if __name__ == '__main__': start()