更新CRNN,lightCRNN模型,修正pyqt的问号,增加识别框上方的车牌号,更新模型选择列表
This commit is contained in:
@@ -2,6 +2,7 @@ import cv2
|
||||
import numpy as np
|
||||
from ultralytics import YOLO
|
||||
import os
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
class LicensePlateYOLO:
|
||||
"""
|
||||
@@ -113,19 +114,38 @@ class LicensePlateYOLO:
|
||||
print(f"检测过程中出错: {e}")
|
||||
return []
|
||||
|
||||
def draw_detections(self, image, detections):
|
||||
def draw_detections(self, image, detections, plate_numbers=None):
|
||||
"""
|
||||
在图像上绘制检测结果
|
||||
|
||||
参数:
|
||||
image: 输入图像
|
||||
detections: 检测结果列表
|
||||
plate_numbers: 车牌号列表,与detections对应
|
||||
|
||||
返回:
|
||||
numpy.ndarray: 绘制了检测结果的图像
|
||||
"""
|
||||
draw_image = image.copy()
|
||||
|
||||
# 转换为PIL图像以支持中文字符
|
||||
pil_image = Image.fromarray(cv2.cvtColor(draw_image, cv2.COLOR_BGR2RGB))
|
||||
draw = ImageDraw.Draw(pil_image)
|
||||
|
||||
# 尝试加载中文字体
|
||||
try:
|
||||
# Windows系统常见的中文字体
|
||||
font_path = "C:/Windows/Fonts/simhei.ttf" # 黑体
|
||||
if not os.path.exists(font_path):
|
||||
font_path = "C:/Windows/Fonts/msyh.ttc" # 微软雅黑
|
||||
if not os.path.exists(font_path):
|
||||
font_path = "C:/Windows/Fonts/simsun.ttc" # 宋体
|
||||
|
||||
font = ImageFont.truetype(font_path, 20)
|
||||
except:
|
||||
# 如果无法加载字体,使用默认字体
|
||||
font = ImageFont.load_default()
|
||||
|
||||
for i, detection in enumerate(detections):
|
||||
box = detection['box']
|
||||
keypoints = detection['keypoints']
|
||||
@@ -133,6 +153,11 @@ class LicensePlateYOLO:
|
||||
confidence = detection['confidence']
|
||||
incomplete = detection.get('incomplete', False)
|
||||
|
||||
# 获取对应的车牌号
|
||||
plate_number = ""
|
||||
if plate_numbers and i < len(plate_numbers):
|
||||
plate_number = plate_numbers[i]
|
||||
|
||||
# 绘制边界框
|
||||
x1, y1, x2, y2 = map(int, box)
|
||||
|
||||
@@ -140,30 +165,53 @@ class LicensePlateYOLO:
|
||||
if class_name == '绿牌':
|
||||
box_color = (0, 255, 0) # 绿色
|
||||
elif class_name == '蓝牌':
|
||||
box_color = (255, 0, 0) # 蓝色
|
||||
box_color = (0, 0, 255) # 蓝色
|
||||
else:
|
||||
box_color = (128, 128, 128) # 灰色
|
||||
|
||||
cv2.rectangle(draw_image, (x1, y1), (x2, y2), box_color, 2)
|
||||
# 在PIL图像上绘制边界框
|
||||
draw.rectangle([(x1, y1), (x2, y2)], outline=box_color, width=2)
|
||||
|
||||
# 构建标签文本
|
||||
if plate_number:
|
||||
label = f"{class_name} {plate_number} {confidence:.2f}"
|
||||
else:
|
||||
label = f"{class_name} {confidence:.2f}"
|
||||
|
||||
# 绘制标签
|
||||
label = f"{class_name} {confidence:.2f}"
|
||||
if incomplete:
|
||||
label += " (不完整)"
|
||||
|
||||
# 计算文本大小和位置
|
||||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||||
font_scale = 0.6
|
||||
thickness = 2
|
||||
(text_width, text_height), _ = cv2.getTextSize(label, font, font_scale, thickness)
|
||||
# 计算文本大小
|
||||
bbox = draw.textbbox((0, 0), label, font=font)
|
||||
text_width = bbox[2] - bbox[0]
|
||||
text_height = bbox[3] - bbox[1]
|
||||
|
||||
# 绘制文本背景
|
||||
cv2.rectangle(draw_image, (x1, y1 - text_height - 10),
|
||||
(x1 + text_width, y1), box_color, -1)
|
||||
draw.rectangle([(x1, y1 - text_height - 10), (x1 + text_width, y1)],
|
||||
fill=box_color)
|
||||
|
||||
# 绘制文本
|
||||
cv2.putText(draw_image, label, (x1, y1 - 5),
|
||||
font, font_scale, (255, 255, 255), thickness)
|
||||
draw.text((x1, y1 - text_height - 5), label, fill=(255, 255, 255), font=font)
|
||||
|
||||
# 转换回OpenCV格式
|
||||
draw_image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
|
||||
|
||||
# 绘制关键点和连线(使用OpenCV)
|
||||
for i, detection in enumerate(detections):
|
||||
box = detection['box']
|
||||
keypoints = detection['keypoints']
|
||||
incomplete = detection.get('incomplete', False)
|
||||
|
||||
x1, y1, x2, y2 = map(int, box)
|
||||
|
||||
# 根据车牌类型选择颜色
|
||||
class_name = detection['class_name']
|
||||
if class_name == '绿牌':
|
||||
box_color = (0, 255, 0) # 绿色
|
||||
elif class_name == '蓝牌':
|
||||
box_color = (0, 0, 255) # 蓝色
|
||||
else:
|
||||
box_color = (128, 128, 128) # 灰色
|
||||
|
||||
# 绘制关键点和连线
|
||||
if len(keypoints) >= 4 and not incomplete:
|
||||
|
||||
Reference in New Issue
Block a user