更新 main.py

This commit is contained in:
spdis 2025-09-17 22:12:45 +08:00
parent 2b443b7373
commit 7d80c90b47

380
main.py
View File

@ -2,6 +2,7 @@ import sys
import os import os
import cv2 import cv2
import numpy as np import numpy as np
from collections import defaultdict, deque
from PyQt5.QtWidgets import QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, \ from PyQt5.QtWidgets import QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, \
QFileDialog, QFrame, QScrollArea, QComboBox QFileDialog, QFrame, QScrollArea, QComboBox
from PyQt5.QtCore import QTimer, Qt, pyqtSignal, QThread from PyQt5.QtCore import QTimer, Qt, pyqtSignal, QThread
@ -18,6 +19,190 @@ from yolopart.detector import LicensePlateYOLO
# from CRNN_part.crnn_interface import LPRNmodel_predict # from CRNN_part.crnn_interface import LPRNmodel_predict
# from CRNN_part.crnn_interface import LPRNinitialize_model # from CRNN_part.crnn_interface import LPRNinitialize_model
class PlateStabilizer:
"""车牌识别结果稳定器"""
def __init__(self, history_size=10, confidence_threshold=0.6, stability_frames=5):
self.history_size = history_size # 历史帧数量
self.confidence_threshold = confidence_threshold # 置信度阈值
self.stability_frames = stability_frames # 稳定帧数要求
# 存储每个车牌的历史识别结果
self.plate_histories = defaultdict(lambda: deque(maxlen=history_size))
# 存储当前稳定的车牌结果
self.stable_results = {}
# 车牌ID计数器
self.plate_id_counter = 0
# 车牌位置追踪
self.plate_positions = {}
def calculate_plate_distance(self, pos1, pos2):
"""计算两个车牌位置的距离"""
if pos1 is None or pos2 is None:
return float('inf')
# 计算中心点距离
center1 = ((pos1[0] + pos1[2]) / 2, (pos1[1] + pos1[3]) / 2)
center2 = ((pos2[0] + pos2[2]) / 2, (pos2[1] + pos2[3]) / 2)
return np.sqrt((center1[0] - center2[0])**2 + (center1[1] - center2[1])**2)
def match_plates_to_history(self, current_detections):
"""将当前检测结果匹配到历史记录"""
matched_plates = {}
used_ids = set()
for detection in current_detections:
bbox = detection.get('bbox', [0, 0, 0, 0])
best_match_id = None
min_distance = float('inf')
# 寻找最佳匹配的历史车牌
for plate_id, last_pos in self.plate_positions.items():
if plate_id in used_ids:
continue
distance = self.calculate_plate_distance(bbox, last_pos)
if distance < min_distance and distance < 100: # 距离阈值
min_distance = distance
best_match_id = plate_id
if best_match_id is not None:
matched_plates[best_match_id] = detection
used_ids.add(best_match_id)
self.plate_positions[best_match_id] = bbox
else:
# 创建新的车牌ID
new_id = f"plate_{self.plate_id_counter}"
self.plate_id_counter += 1
matched_plates[new_id] = detection
self.plate_positions[new_id] = bbox
return matched_plates
def calculate_confidence(self, plate_text, detection_quality=1.0):
"""计算识别结果的置信度"""
if not plate_text or plate_text == "识别失败":
return 0.0
# 基础置信度基于文本长度和字符类型
base_confidence = 0.5
# 长度合理性检查
if 7 <= len(plate_text) <= 8:
base_confidence += 0.2
# 字符类型检查(中文+字母+数字的组合)
has_chinese = any('\u4e00' <= char <= '\u9fff' for char in plate_text)
has_letter = any(char.isalpha() for char in plate_text)
has_digit = any(char.isdigit() for char in plate_text)
if has_chinese and has_letter and has_digit:
base_confidence += 0.2
# 检测质量影响
confidence = base_confidence * detection_quality
return min(confidence, 1.0)
def update_and_get_stable_result(self, current_detections, corrected_images, plate_texts):
"""更新历史记录并返回稳定的识别结果"""
if not current_detections:
return []
# 匹配当前检测到历史记录
matched_plates = self.match_plates_to_history(current_detections)
stable_results = []
for plate_id, detection in matched_plates.items():
# 获取对应的矫正图像和识别文本
detection_idx = current_detections.index(detection)
corrected_image = corrected_images[detection_idx] if detection_idx < len(corrected_images) else None
plate_text = plate_texts[detection_idx] if detection_idx < len(plate_texts) else "识别失败"
# 计算置信度
confidence = self.calculate_confidence(plate_text)
# 添加到历史记录
history_entry = {
'text': plate_text,
'confidence': confidence,
'detection': detection,
'corrected_image': corrected_image
}
self.plate_histories[plate_id].append(history_entry)
# 计算稳定结果
stable_text = self.get_stable_text(plate_id)
if stable_text and stable_text != "识别失败":
stable_results.append({
'id': plate_id,
'class_name': detection['class_name'],
'corrected_image': corrected_image,
'plate_number': stable_text,
'detection': detection
})
return stable_results
def get_stable_text(self, plate_id):
"""获取指定车牌的稳定识别结果"""
history = self.plate_histories[plate_id]
if len(history) < 3: # 历史记录太少,返回最新结果
return history[-1]['text'] if history else "识别失败"
# 统计各种识别结果的加权投票
text_votes = defaultdict(float)
total_confidence = 0
for entry in history:
text = entry['text']
confidence = entry['confidence']
if text != "识别失败" and confidence > 0.3:
text_votes[text] += confidence
total_confidence += confidence
if not text_votes:
return "识别失败"
# 找到得票最高的结果
best_text = max(text_votes.items(), key=lambda x: x[1])
# 检查是否足够稳定(得票率超过阈值)
vote_ratio = best_text[1] / total_confidence if total_confidence > 0 else 0
if vote_ratio >= self.confidence_threshold:
return best_text[0]
else:
# 不够稳定,返回最近的高置信度结果
recent_high_conf = [entry for entry in list(history)[-5:]
if entry['confidence'] > 0.5 and entry['text'] != "识别失败"]
if recent_high_conf:
return recent_high_conf[-1]['text']
else:
return history[-1]['text']
def clear_old_plates(self, current_plate_ids):
"""清理不再出现的车牌历史记录"""
# 移除超过一定时间未更新的车牌
plates_to_remove = []
for plate_id in self.plate_histories.keys():
if plate_id not in current_plate_ids:
plates_to_remove.append(plate_id)
for plate_id in plates_to_remove:
if plate_id in self.plate_histories:
del self.plate_histories[plate_id]
if plate_id in self.plate_positions:
del self.plate_positions[plate_id]
if plate_id in self.stable_results:
del self.stable_results[plate_id]
class CameraThread(QThread): class CameraThread(QThread):
"""摄像头线程类""" """摄像头线程类"""
frame_ready = pyqtSignal(np.ndarray) frame_ready = pyqtSignal(np.ndarray)
@ -118,6 +303,7 @@ class LicensePlateWidget(QWidget):
def init_ui(self, class_name, corrected_image, plate_number): def init_ui(self, class_name, corrected_image, plate_number):
layout = QHBoxLayout() layout = QHBoxLayout()
layout.setContentsMargins(10, 5, 10, 5) layout.setContentsMargins(10, 5, 10, 5)
layout.setSpacing(8) # 设置组件间距
# 车牌类型标签 # 车牌类型标签
type_label = QLabel(class_name) type_label = QLabel(class_name)
@ -155,7 +341,6 @@ class LicensePlateWidget(QWidget):
# 矫正后的车牌图像 # 矫正后的车牌图像
image_label = QLabel() image_label = QLabel()
image_label.setFixedSize(120, 40)
image_label.setStyleSheet("border: 1px solid #ddd; background-color: white;") image_label.setStyleSheet("border: 1px solid #ddd; background-color: white;")
if corrected_image is not None: if corrected_image is not None:
@ -169,16 +354,44 @@ class LicensePlateWidget(QWidget):
q_image = QImage(corrected_image.data, w, h, bytes_per_line, QImage.Format_Grayscale8) q_image = QImage(corrected_image.data, w, h, bytes_per_line, QImage.Format_Grayscale8)
pixmap = QPixmap.fromImage(q_image) pixmap = QPixmap.fromImage(q_image)
scaled_pixmap = pixmap.scaled(120, 40, Qt.KeepAspectRatio, Qt.SmoothTransformation)
# 动态计算显示尺寸,保持车牌的宽高比
original_width = pixmap.width()
original_height = pixmap.height()
# 设置最大显示尺寸限制
max_width = 150
max_height = 60
# 计算缩放比例,确保图像完整显示
width_ratio = max_width / original_width if original_width > 0 else 1
height_ratio = max_height / original_height if original_height > 0 else 1
scale_ratio = min(width_ratio, height_ratio, 1.0) # 不放大,只缩小
# 计算实际显示尺寸
display_width = int(original_width * scale_ratio)
display_height = int(original_height * scale_ratio)
# 确保最小显示尺寸
display_width = max(display_width, 80)
display_height = max(display_height, 25)
# 设置标签尺寸并缩放图像
image_label.setFixedSize(display_width, display_height)
scaled_pixmap = pixmap.scaled(display_width, display_height, Qt.KeepAspectRatio, Qt.SmoothTransformation)
image_label.setPixmap(scaled_pixmap) image_label.setPixmap(scaled_pixmap)
image_label.setAlignment(Qt.AlignCenter)
else: else:
# 当没有图像时,设置固定尺寸显示提示信息
image_label.setFixedSize(120, 40)
image_label.setText("车牌未完全\n进入摄像头") image_label.setText("车牌未完全\n进入摄像头")
image_label.setAlignment(Qt.AlignCenter) image_label.setAlignment(Qt.AlignCenter)
image_label.setStyleSheet("border: 1px solid #ddd; background-color: #f5f5f5; color: #666;") image_label.setStyleSheet("border: 1px solid #ddd; background-color: #f5f5f5; color: #666;")
# 车牌号标签 # 车牌号标签 - 使用自适应宽度
number_label = QLabel(plate_number) number_label = QLabel(plate_number)
number_label.setFixedWidth(150) number_label.setMinimumWidth(120) # 设置最小宽度
number_label.setMaximumWidth(200) # 设置最大宽度
number_label.setAlignment(Qt.AlignCenter) number_label.setAlignment(Qt.AlignCenter)
number_label.setStyleSheet( number_label.setStyleSheet(
"QLabel { " "QLabel { "
@ -190,6 +403,11 @@ class LicensePlateWidget(QWidget):
"font-weight: bold; " "font-weight: bold; "
"}" "}"
) )
# 根据文本长度调整宽度
font_metrics = number_label.fontMetrics()
text_width = font_metrics.boundingRect(plate_number).width()
optimal_width = max(120, min(200, text_width + 20)) # 加20像素的边距
number_label.setFixedWidth(optimal_width)
layout.addWidget(type_label) layout.addWidget(type_label)
layout.addWidget(image_label) layout.addWidget(image_label)
@ -197,6 +415,9 @@ class LicensePlateWidget(QWidget):
layout.addStretch() layout.addStretch()
self.setLayout(layout) self.setLayout(layout)
# 调整整体组件的最小高度以适应动态图像尺寸
min_height = max(60, image_label.height() + 20) # 至少60像素高度
self.setMinimumHeight(min_height)
self.setStyleSheet( self.setStyleSheet(
"QWidget { " "QWidget { "
"background-color: white; " "background-color: white; "
@ -221,6 +442,13 @@ class MainWindow(QMainWindow):
self.last_plate_results = [] # 存储上一次的车牌识别结果 self.last_plate_results = [] # 存储上一次的车牌识别结果
self.current_recognition_method = "CRNN" # 当前识别方法 self.current_recognition_method = "CRNN" # 当前识别方法
# 添加车牌稳定器
self.plate_stabilizer = PlateStabilizer(
history_size=15, # 保存15帧历史
confidence_threshold=0.7, # 70%置信度阈值
stability_frames=5 # 需要5帧稳定
)
self.init_ui() self.init_ui()
self.init_detector() self.init_detector()
self.init_camera() self.init_camera()
@ -292,7 +520,7 @@ class MainWindow(QMainWindow):
# 右侧结果显示区域 # 右侧结果显示区域
right_frame = QFrame() right_frame = QFrame()
right_frame.setFrameStyle(QFrame.StyledPanel) right_frame.setFrameStyle(QFrame.StyledPanel)
right_frame.setFixedWidth(400) right_frame.setFixedWidth(460)
right_frame.setStyleSheet("QFrame { background-color: #fafafa; border: 2px solid #ddd; }") right_frame.setStyleSheet("QFrame { background-color: #fafafa; border: 2px solid #ddd; }")
right_layout = QVBoxLayout(right_frame) right_layout = QVBoxLayout(right_frame)
@ -388,6 +616,32 @@ class MainWindow(QMainWindow):
model_path = os.path.join(os.path.dirname(__file__), "yolopart", "yolo11s-pose42.pt") model_path = os.path.join(os.path.dirname(__file__), "yolopart", "yolo11s-pose42.pt")
self.detector = LicensePlateYOLO(model_path) self.detector = LicensePlateYOLO(model_path)
def reset_processing_state(self):
"""重置处理状态和清理界面"""
# 重置处理标志
self.is_processing = False
# 清空当前帧和检测结果
self.current_frame = None
self.detections = []
# 重置车牌稳定器
self.plate_stabilizer = PlateStabilizer(
history_size=15,
confidence_threshold=0.7,
stability_frames=5
)
# 清空右侧结果显示
self.count_label.setText("识别到的车牌数量: 0")
for i in reversed(range(self.results_layout.count())):
child = self.results_layout.itemAt(i).widget()
if child:
child.setParent(None)
self.last_plate_results = []
print("处理状态已重置,界面已清理")
def init_camera(self): def init_camera(self):
"""初始化摄像头线程""" """初始化摄像头线程"""
self.camera_thread = CameraThread() self.camera_thread = CameraThread()
@ -401,7 +655,11 @@ class MainWindow(QMainWindow):
def start_camera(self): def start_camera(self):
"""启动摄像头""" """启动摄像头"""
# 重置处理状态和清理界面
self.reset_processing_state()
if self.camera_thread.start_camera(): if self.camera_thread.start_camera():
self.current_mode = "camera"
self.start_button.setEnabled(False) self.start_button.setEnabled(False)
self.stop_button.setEnabled(True) self.stop_button.setEnabled(True)
self.camera_label.setText("摄像头启动中...") self.camera_label.setText("摄像头启动中...")
@ -435,6 +693,9 @@ class MainWindow(QMainWindow):
elif self.current_mode == "video" and self.video_thread and self.video_thread.running: elif self.current_mode == "video" and self.video_thread and self.video_thread.running:
self.stop_video() self.stop_video()
# 重置处理状态和清理界面
self.reset_processing_state()
# 选择视频文件 # 选择视频文件
video_path, _ = QFileDialog.getOpenFileName(self, "选择视频文件", "", "视频文件 (*.mp4 *.avi *.mov *.mkv)") video_path, _ = QFileDialog.getOpenFileName(self, "选择视频文件", "", "视频文件 (*.mp4 *.avi *.mov *.mkv)")
@ -483,6 +744,9 @@ class MainWindow(QMainWindow):
elif self.current_mode == "video" and self.video_thread and self.video_thread.running: elif self.current_mode == "video" and self.video_thread and self.video_thread.running:
self.stop_video() self.stop_video()
# 重置处理状态和清理界面
self.reset_processing_state()
# 选择图片文件 # 选择图片文件
image_path, _ = QFileDialog.getOpenFileName(self, "选择图片文件", "", "图片文件 (*.jpg *.jpeg *.png *.bmp)") image_path, _ = QFileDialog.getOpenFileName(self, "选择图片文件", "", "图片文件 (*.jpg *.jpeg *.png *.bmp)")
@ -538,6 +802,9 @@ class MainWindow(QMainWindow):
def process_frame(self, frame): def process_frame(self, frame):
"""处理摄像头帧""" """处理摄像头帧"""
if frame is None:
return
self.current_frame = frame.copy() self.current_frame = frame.copy()
# 先显示原始帧,保证视频流畅播放 # 先显示原始帧,保证视频流畅播放
@ -686,48 +953,47 @@ class MainWindow(QMainWindow):
self.camera_label.setText(f"显示错误: {str(e)}") self.camera_label.setText(f"显示错误: {str(e)}")
def update_results_display(self): def update_results_display(self):
"""更新右侧结果显示""" """更新右侧结果显示(使用稳定化结果)"""
# 更新车牌数量 print(f"开始更新结果显示,当前模式: {self.current_mode}, 检测数量: {len(self.detections) if self.detections else 0}")
count = len(self.detections)
self.count_label.setText(f"识别到的车牌数量: {count}")
# 准备新的车牌结果列表 if not self.detections:
new_plate_results = [] self.count_label.setText("识别到的车牌数量: 0")
for i, detection in enumerate(self.detections): # 清除显示
# 矫正车牌图像 for i in reversed(range(self.results_layout.count())):
child = self.results_layout.itemAt(i).widget()
if child:
child.setParent(None)
self.last_plate_results = []
print("无检测结果,已清空界面")
return
# 获取矫正图像和识别文本
corrected_images = []
plate_texts = []
for detection in self.detections:
corrected_image = self.correct_license_plate(detection) corrected_image = self.correct_license_plate(detection)
corrected_images.append(corrected_image)
# 获取车牌号,传入车牌类型信息 if corrected_image is not None:
plate_number = self.recognize_plate_number(corrected_image, detection['class_name']) plate_text = self.recognize_plate_number(corrected_image, detection['class_name'])
# 添加到新结果列表
new_plate_results.append({
'id': i + 1,
'class_name': detection['class_name'],
'corrected_image': corrected_image,
'plate_number': plate_number
})
# 比较新旧结果是否相同
results_changed = False
if len(self.last_plate_results) != len(new_plate_results):
results_changed = True
else: else:
for i in range(len(new_plate_results)): plate_text = "识别失败"
if i >= len(self.last_plate_results): plate_texts.append(plate_text)
results_changed = True
break
last_result = self.last_plate_results[i] # 使用稳定器获取稳定的识别结果
new_result = new_plate_results[i] stable_results = self.plate_stabilizer.update_and_get_stable_result(
self.detections, corrected_images, plate_texts
)
# 比较车牌类型和车牌号 # 更新车牌数量显示
if (last_result['class_name'] != new_result['class_name'] or self.count_label.setText(f"识别到的车牌数量: {len(stable_results)}")
last_result['plate_number'] != new_result['plate_number']): print(f"稳定结果数量: {len(stable_results)}")
results_changed = True
break # 检查结果是否发生变化
results_changed = self.check_results_changed(stable_results)
print(f"结果是否变化: {results_changed}")
# 只有当结果发生变化时才更新显示
if results_changed: if results_changed:
# 清除之前的结果 # 清除之前的结果
for i in reversed(range(self.results_layout.count())): for i in reversed(range(self.results_layout.count())):
@ -735,18 +1001,42 @@ class MainWindow(QMainWindow):
if child: if child:
child.setParent(None) child.setParent(None)
# 添加新的结果 # 添加新的稳定结果
for result in new_plate_results: for i, result in enumerate(stable_results):
plate_widget = LicensePlateWidget( plate_widget = LicensePlateWidget(
result['id'], i + 1, # 显示序号
result['class_name'], result['class_name'],
result['corrected_image'], result['corrected_image'],
result['plate_number'] result['plate_number']
) )
self.results_layout.addWidget(plate_widget) self.results_layout.addWidget(plate_widget)
print(f"添加车牌widget: {result['plate_number']}")
# 更新存储的上一次结果 # 更新存储的结果
self.last_plate_results = new_plate_results self.last_plate_results = stable_results
# 清理旧的车牌记录
current_plate_ids = [result['id'] for result in stable_results]
self.plate_stabilizer.clear_old_plates(current_plate_ids)
print("结果显示更新完成")
def check_results_changed(self, new_results):
"""检查识别结果是否发生变化"""
if len(self.last_plate_results) != len(new_results):
return True
for i, new_result in enumerate(new_results):
if i >= len(self.last_plate_results):
return True
old_result = self.last_plate_results[i]
# 比较关键字段
if (old_result.get('class_name') != new_result.get('class_name') or
old_result.get('plate_number') != new_result.get('plate_number')):
return True
return False
def correct_license_plate(self, detection): def correct_license_plate(self, detection):
"""矫正车牌图像""" """矫正车牌图像"""