更新 main.py
This commit is contained in:
parent
2b443b7373
commit
7d80c90b47
384
main.py
384
main.py
@ -2,6 +2,7 @@ import sys
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
from collections import defaultdict, deque
|
||||
from PyQt5.QtWidgets import QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, \
|
||||
QFileDialog, QFrame, QScrollArea, QComboBox
|
||||
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 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):
|
||||
"""摄像头线程类"""
|
||||
frame_ready = pyqtSignal(np.ndarray)
|
||||
@ -118,6 +303,7 @@ class LicensePlateWidget(QWidget):
|
||||
def init_ui(self, class_name, corrected_image, plate_number):
|
||||
layout = QHBoxLayout()
|
||||
layout.setContentsMargins(10, 5, 10, 5)
|
||||
layout.setSpacing(8) # 设置组件间距
|
||||
|
||||
# 车牌类型标签
|
||||
type_label = QLabel(class_name)
|
||||
@ -155,7 +341,6 @@ class LicensePlateWidget(QWidget):
|
||||
|
||||
# 矫正后的车牌图像
|
||||
image_label = QLabel()
|
||||
image_label.setFixedSize(120, 40)
|
||||
image_label.setStyleSheet("border: 1px solid #ddd; background-color: white;")
|
||||
|
||||
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)
|
||||
|
||||
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.setAlignment(Qt.AlignCenter)
|
||||
else:
|
||||
# 当没有图像时,设置固定尺寸显示提示信息
|
||||
image_label.setFixedSize(120, 40)
|
||||
image_label.setText("车牌未完全\n进入摄像头")
|
||||
image_label.setAlignment(Qt.AlignCenter)
|
||||
image_label.setStyleSheet("border: 1px solid #ddd; background-color: #f5f5f5; color: #666;")
|
||||
|
||||
# 车牌号标签
|
||||
# 车牌号标签 - 使用自适应宽度
|
||||
number_label = QLabel(plate_number)
|
||||
number_label.setFixedWidth(150)
|
||||
number_label.setMinimumWidth(120) # 设置最小宽度
|
||||
number_label.setMaximumWidth(200) # 设置最大宽度
|
||||
number_label.setAlignment(Qt.AlignCenter)
|
||||
number_label.setStyleSheet(
|
||||
"QLabel { "
|
||||
@ -190,6 +403,11 @@ class LicensePlateWidget(QWidget):
|
||||
"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(image_label)
|
||||
@ -197,6 +415,9 @@ class LicensePlateWidget(QWidget):
|
||||
layout.addStretch()
|
||||
|
||||
self.setLayout(layout)
|
||||
# 调整整体组件的最小高度以适应动态图像尺寸
|
||||
min_height = max(60, image_label.height() + 20) # 至少60像素高度
|
||||
self.setMinimumHeight(min_height)
|
||||
self.setStyleSheet(
|
||||
"QWidget { "
|
||||
"background-color: white; "
|
||||
@ -221,6 +442,13 @@ class MainWindow(QMainWindow):
|
||||
self.last_plate_results = [] # 存储上一次的车牌识别结果
|
||||
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_detector()
|
||||
self.init_camera()
|
||||
@ -292,7 +520,7 @@ class MainWindow(QMainWindow):
|
||||
# 右侧结果显示区域
|
||||
right_frame = QFrame()
|
||||
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_layout = QVBoxLayout(right_frame)
|
||||
|
||||
@ -388,6 +616,32 @@ class MainWindow(QMainWindow):
|
||||
model_path = os.path.join(os.path.dirname(__file__), "yolopart", "yolo11s-pose42.pt")
|
||||
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):
|
||||
"""初始化摄像头线程"""
|
||||
self.camera_thread = CameraThread()
|
||||
@ -401,7 +655,11 @@ class MainWindow(QMainWindow):
|
||||
|
||||
def start_camera(self):
|
||||
"""启动摄像头"""
|
||||
# 重置处理状态和清理界面
|
||||
self.reset_processing_state()
|
||||
|
||||
if self.camera_thread.start_camera():
|
||||
self.current_mode = "camera"
|
||||
self.start_button.setEnabled(False)
|
||||
self.stop_button.setEnabled(True)
|
||||
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:
|
||||
self.stop_video()
|
||||
|
||||
# 重置处理状态和清理界面
|
||||
self.reset_processing_state()
|
||||
|
||||
# 选择视频文件
|
||||
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:
|
||||
self.stop_video()
|
||||
|
||||
# 重置处理状态和清理界面
|
||||
self.reset_processing_state()
|
||||
|
||||
# 选择图片文件
|
||||
image_path, _ = QFileDialog.getOpenFileName(self, "选择图片文件", "", "图片文件 (*.jpg *.jpeg *.png *.bmp)")
|
||||
|
||||
@ -538,6 +802,9 @@ class MainWindow(QMainWindow):
|
||||
|
||||
def process_frame(self, frame):
|
||||
"""处理摄像头帧"""
|
||||
if frame is None:
|
||||
return
|
||||
|
||||
self.current_frame = frame.copy()
|
||||
|
||||
# 先显示原始帧,保证视频流畅播放
|
||||
@ -686,48 +953,47 @@ class MainWindow(QMainWindow):
|
||||
self.camera_label.setText(f"显示错误: {str(e)}")
|
||||
|
||||
def update_results_display(self):
|
||||
"""更新右侧结果显示"""
|
||||
# 更新车牌数量
|
||||
count = len(self.detections)
|
||||
self.count_label.setText(f"识别到的车牌数量: {count}")
|
||||
"""更新右侧结果显示(使用稳定化结果)"""
|
||||
print(f"开始更新结果显示,当前模式: {self.current_mode}, 检测数量: {len(self.detections) if self.detections else 0}")
|
||||
|
||||
# 准备新的车牌结果列表
|
||||
new_plate_results = []
|
||||
for i, detection in enumerate(self.detections):
|
||||
# 矫正车牌图像
|
||||
if not self.detections:
|
||||
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("无检测结果,已清空界面")
|
||||
return
|
||||
|
||||
# 获取矫正图像和识别文本
|
||||
corrected_images = []
|
||||
plate_texts = []
|
||||
|
||||
for detection in self.detections:
|
||||
corrected_image = self.correct_license_plate(detection)
|
||||
corrected_images.append(corrected_image)
|
||||
|
||||
# 获取车牌号,传入车牌类型信息
|
||||
plate_number = 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
|
||||
})
|
||||
if corrected_image is not None:
|
||||
plate_text = self.recognize_plate_number(corrected_image, detection['class_name'])
|
||||
else:
|
||||
plate_text = "识别失败"
|
||||
plate_texts.append(plate_text)
|
||||
|
||||
# 比较新旧结果是否相同
|
||||
results_changed = False
|
||||
if len(self.last_plate_results) != len(new_plate_results):
|
||||
results_changed = True
|
||||
else:
|
||||
for i in range(len(new_plate_results)):
|
||||
if i >= len(self.last_plate_results):
|
||||
results_changed = True
|
||||
break
|
||||
|
||||
last_result = self.last_plate_results[i]
|
||||
new_result = new_plate_results[i]
|
||||
|
||||
# 比较车牌类型和车牌号
|
||||
if (last_result['class_name'] != new_result['class_name'] or
|
||||
last_result['plate_number'] != new_result['plate_number']):
|
||||
results_changed = True
|
||||
break
|
||||
# 使用稳定器获取稳定的识别结果
|
||||
stable_results = self.plate_stabilizer.update_and_get_stable_result(
|
||||
self.detections, corrected_images, plate_texts
|
||||
)
|
||||
|
||||
# 更新车牌数量显示
|
||||
self.count_label.setText(f"识别到的车牌数量: {len(stable_results)}")
|
||||
print(f"稳定结果数量: {len(stable_results)}")
|
||||
|
||||
# 检查结果是否发生变化
|
||||
results_changed = self.check_results_changed(stable_results)
|
||||
print(f"结果是否变化: {results_changed}")
|
||||
|
||||
# 只有当结果发生变化时才更新显示
|
||||
if results_changed:
|
||||
# 清除之前的结果
|
||||
for i in reversed(range(self.results_layout.count())):
|
||||
@ -735,18 +1001,42 @@ class MainWindow(QMainWindow):
|
||||
if child:
|
||||
child.setParent(None)
|
||||
|
||||
# 添加新的结果
|
||||
for result in new_plate_results:
|
||||
# 添加新的稳定结果
|
||||
for i, result in enumerate(stable_results):
|
||||
plate_widget = LicensePlateWidget(
|
||||
result['id'],
|
||||
i + 1, # 显示序号
|
||||
result['class_name'],
|
||||
result['corrected_image'],
|
||||
result['plate_number']
|
||||
)
|
||||
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):
|
||||
"""矫正车牌图像"""
|
||||
|
Loading…
x
Reference in New Issue
Block a user