import sys,os
from PyQt5.QtCore import *
from PyQt5.QtWidgets import *
from PyQt5.QtGui import *
import cv2
from ultralytics import YOLO
from collections import Counter
import logging
class MainWindow(QMainWindow):
def __init__(self):
super().__init__()
self.xiaolian_ui()
self.con=0.25
self.current_annotated_image = None
self.detection_type = None
self.result = None
def xiaolian_ui(self):
self.setWindowTitle("@author:笑脸惹桃花")
self.setWindowIcon(QIcon("icon.png"))
self.setFixedSize(1380, 590)
# 创建主左右布局
main_layout = QHBoxLayout()
left_panel = QVBoxLayout()
left_panel.setContentsMargins(10, 20, 10, 20) # 增加上下边距
left_panel.setSpacing(15)
# 垂直排列的功能按钮
button_layout = QVBoxLayout()
button_layout.setSpacing(12) # 设置按钮间距
# 创建并添加所有按钮
self.load_model_button = QPushButton("📁 模型选择")
self.image_detect_button = QPushButton("🖼️ 图片检测")
self.folder_detect_button = QPushButton("📂 文件夹检测")
self.save_button = QPushButton("💾 保存结果")
self.exit_button = QPushButton("❌ 退出程序")
# 统一设置按钮属性
for btn in [self.load_model_button, self.image_detect_button,
self.folder_detect_button, self.save_button, self.exit_button]:
btn.setFixedSize(145, 35)
btn.setCursor(Qt.PointingHandCursor)
button_layout.addWidget(btn)
# 添加按钮组与信息框之间的伸缩空间
button_layout.addStretch(1)
# 信息输出框
status_group = QGroupBox("检测信息")
status_group.setStyleSheet("""
QGroupBox {
border: none;
background: rgba(245, 245, 245, 0.9);
border-radius: 8px;
padding: 10px;
}
QGroupBox::title {
subcontrol-origin: margin;
left: 10px;
padding: 0 5px;
color: #2980b9;
font: bold 14px "Segoe UI";
background: transparent;
}
""")
status_layout = QVBoxLayout()
self.output_text = QTextEdit(self)
self.output_text.setReadOnly(True)
self.output_text.setStyleSheet("""
background: white;
border: 1px solid #dcdcdc;
border-radius: 5px;
padding: 8px;
font: 13px "Segoe UI";
color: #2c3e50;
min-height: 150px;
""")
status_layout.addWidget(self.output_text)
status_group.setLayout(status_layout)
status_layout.addStretch()
# 组装左侧面板
left_panel.addLayout(button_layout)
left_panel.addWidget(status_group)
# 右侧视频显示区域 (宽度占比3/4)
right_panel = QHBoxLayout()
right_panel.setContentsMargins(10, 0, 10, 0)
# 创建图像显示区域
self.label1 = QLabel()
self.label1.setAlignment(Qt.AlignCenter)
self.label1.setMinimumSize(580, 560)
self.label1.setMaximumSize(580, 560)
self.label1.setStyleSheet('''
border:3px solid #6950a1;
background-color: black;
border-radius: 8px;
''')
self.label2 = QLabel()
self.label2.setAlignment(Qt.AlignCenter)
self.label2.setMinimumSize(580, 560)
self.label2.setMaximumSize(580, 560)
self.label2.setStyleSheet('''
border:3px solid #6950a1;
background-color: black;
border-radius: 8px;
''')
# 添加图像标签并设置间距
right_panel.addWidget(self.label1)
right_panel.addSpacing(15) # 两个显示区域之间的间距
right_panel.addWidget(self.label2)
# 设置主布局比例
main_layout.addLayout(left_panel, 1)
main_layout.addLayout(right_panel, 3)
# 设置中心部件
central_widget = QWidget()
central_widget.setLayout(main_layout)
self.setCentralWidget(central_widget)
# 连接信号槽
self.load_model_button.clicked.connect(self.load_model)
self.image_detect_button.clicked.connect(self.select_image)
self.folder_detect_button.clicked.connect(self.detect_folder)
self.save_button.clicked.connect(self.save_detection)
self.exit_button.clicked.connect(self.exit_application)
# 增强按钮样式
button_style = """
QPushButton {
background-color: #f0f0f0;
border: 2px solid #d0d0d0;
border-radius: 6px;
color: #404040;
font: bold 13px "Microsoft YaHei";
padding: 8px 12px;
text-align: center;
transition: all 0.3s;
}
QPushButton:hover {
background-color: #e0e0e0;
border-color: #c0c0c0;
color: #202020;
}
QPushButton:pressed {
background-color: #d0d0d0;
border-color: #b0b0b0;
}
"""
for btn in [self.load_model_button, self.image_detect_button,
self.folder_detect_button, self.save_button, self.exit_button]:
btn.setStyleSheet(button_style)
def save_detection(self):
detection_type = self.detection_type
if detection_type == "image":
self.save_detection_results()
def save_detection_results(self):
if self.current_annotated_image is not None:
self.save_image(self.current_annotated_image)
def detect_folder(self):
self.label1.clear()
self.label2.clear()
folder_path = QFileDialog.getExistingDirectory(self, "选择图片文件夹")
self.flag = 1
if folder_path:
image_paths = []
for filename in os.listdir(folder_path):
if filename.lower().endswith((".jpg", ".jpeg", ".png")):
image_path = os.path.join(folder_path, filename)
image_paths.append(image_path)
for image_path in image_paths:
self.detect_image(image_path)
def save_image(self, image):
if image is not None:
file_name, _ = QFileDialog.getSaveFileName(None, "保存图片", "", "JPEG (*.jpg);;PNG (*.png);;All Files (*)")
if file_name:
cv2.imwrite(file_name, image)
def select_image(self):
image_path, _ = QFileDialog.getOpenFileName(None, "选择图片文件", "", "图片文件 (*.jpg *.jpeg *.png)")
self.flag = 0
self.detect_image(image_path)
def detect_image(self,image_path):
if image_path:
image = cv2.imread(image_path)
if image is not None:
if self.flag == 0:
results = self.model.predict(image)
elif self.flag == 1:
results = self.model.predict(image_path, save=True)
self.detection_type = "image"
if results:
annotated_image = results[0].plot()
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) #转换为 RGB
height1, width1, channel1 = image_rgb.shape
bytesPerLine1 = 3 * width1
qimage1 = QImage(image_rgb.data, width1, height1, bytesPerLine1, QImage.Format_RGB888)
pixmap1 = QPixmap.fromImage(qimage1)
self.label1.setPixmap(pixmap1.scaled(self.label1.size(), Qt.KeepAspectRatio))
annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB) #转换为 RGB
height2, width2, channel2 = annotated_image.shape
bytesPerLine2 = 3 * width2
qimage2 = QImage(annotated_image.data, width2, height2, bytesPerLine2, QImage.Format_RGB888)
pixmap2 = QPixmap.fromImage(qimage2)
self.label2.setPixmap(pixmap2.scaled(self.label2.size(), Qt.KeepAspectRatio))
self.result = results
self.current_annotated_image = results[0].plot()
self.display_statistics()
cv2.waitKey(300) # 修改图片切换时间
def stat(self):
detected_classes = []
target_range = {0, 1} # 修改为自己的类别
if self.result == None:
return None
for r in self.result:
classes = r.boxes.cls.cpu().numpy().astype(int).tolist()
# 筛选出在目标范围内的类别,并去重
detected_classes.extend([cls for cls in classes if cls in target_range])
class_counts = Counter(detected_classes)
class_counts_dic = dict(class_counts)
return class_counts_dic
def display_statistics(self):
class_counts = self.stat()
if class_counts == None:
self.output_text.setText('')
return
# 修改class_labels为自己的类别对应关系,可中文
class_labels = {
0: "helmet", 1: "vest"
}
# 构建输出字符串
output_string = ""
for class_id, count in class_counts.items():
label = class_labels.get(class_id, f"类别{class_id}") # 如果没有找到标签,则使用默认标签
output_string += f"{label}: {count} 个\n"
self.output_text.setText(output_string)
def load_model(self):
model_path, _ = QFileDialog.getOpenFileName(None, "选择模型文件", "", "模型文件 (*.pt)")
if model_path:
self.model = YOLO(model_path)
def exit_application(self):
sys.exit()
class LoginDialog(QDialog):
def __init__(self):
super().__init__()
self.setWindowTitle("@author:笑脸惹桃花")
self.setWindowIcon(QIcon("icon.png"))
self.setWindowFlags(Qt.WindowCloseButtonHint)
self.setFixedSize(500, 800)
self.setup_ui()
def setup_ui(self):
main_layout = QVBoxLayout()
main_layout.setContentsMargins(40, 30, 40, 40) # 增加边距
main_layout.setSpacing(30)
logo_frame = QFrame()
logo_layout = QHBoxLayout(logo_frame)
logo_layout.setContentsMargins(0, 0, 0, 0)
logo_label = QLabel()
logo_pixmap = QPixmap("icon.png").scaled(150, 150, Qt.KeepAspectRatio, Qt.SmoothTransformation)
logo_label.setPixmap(logo_pixmap)
logo_layout.addWidget(logo_label, 0, Qt.AlignCenter)
main_layout.addWidget(logo_frame)
title_label = QLabel("@author:笑脸惹桃花")
title_label.setFont(QFont("微软雅黑", 24, QFont.Bold))
title_label.setStyleSheet("color: #2c3e50;")
title_label.setAlignment(Qt.AlignCenter)
main_layout.addWidget(title_label)
form_frame = QFrame()
form_layout = QVBoxLayout(form_frame)
form_layout.setContentsMargins(20, 20, 20, 20)
form_layout.setSpacing(25)
username_layout = QHBoxLayout()
username_label = QLabel("账号:")
username_label.setFont(QFont("微软雅黑", 14))
username_label.setStyleSheet("color: #2c3e50;")
username_layout.addWidget(username_label)
self.username_input = QLineEdit()
self.username_input.setPlaceholderText("请输入用户名")
self.username_input.setFont(QFont("微软雅黑", 14))
self.username_input.setStyleSheet(
"QLineEdit { padding: 12px; border-radius: 6px; border: 2px solid #bdc3c7; }"
"QLineEdit:focus { border-color: #3498db; }"
)
username_layout.addWidget(self.username_input)
form_layout.addLayout(username_layout)
password_layout = QHBoxLayout()
password_label = QLabel("密码:")
password_label.setFont(QFont("微软雅黑", 14))
password_label.setStyleSheet("color: #2c3e50;")
password_layout.addWidget(password_label)
self.password_input = QLineEdit()
self.password_input.setPlaceholderText("请输入密码")
self.password_input.setEchoMode(QLineEdit.Password)
self.password_input.setFont(QFont("微软雅黑", 14))
self.password_input.setStyleSheet(
"QLineEdit { padding: 12px; border-radius: 6px; border: 2px solid #bdc3c7; }"
"QLineEdit:focus { border-color: #3498db; }"
)
password_layout.addWidget(self.password_input)
form_layout.addLayout(password_layout)
main_layout.addWidget(form_frame)
button_frame = QFrame()
button_layout = QHBoxLayout(button_frame)
button_layout.setContentsMargins(0, 0, 0, 0)
button_layout.setSpacing(15)
login_button = QPushButton("登 录")
login_button.setFont(QFont("微软雅黑", 14, QFont.Bold))
login_button.clicked.connect(self.handle_login)
login_button.setFixedSize(120, 45)
login_button.setStyleSheet("""
QPushButton {
background-color: #3498db;
color: white;
border-radius: 8px;
border: none;
}
QPushButton:hover { background-color: #2980b9; }
QPushButton:pressed { background-color: #1d6fa5; }
""")
register_button = QPushButton("注 册")
register_button.setFont(QFont("微软雅黑", 14, QFont.Bold))
register_button.clicked.connect(self.handle_register)
register_button.setFixedSize(120, 45)
register_button.setStyleSheet("""
QPushButton {
background-color: #2ecc71;
color: white;
border-radius: 8px;
border: none;
}
QPushButton:hover { background-color: #27ae60; }
QPushButton:pressed { background-color: #219a52; }
""")
exit_button = QPushButton("退 出")
exit_button.setFont(QFont("微软雅黑", 14, QFont.Bold))
exit_button.clicked.connect(self.close)
exit_button.setFixedSize(120, 45)
exit_button.setStyleSheet("""
QPushButton {
background-color: #e74c3c;
color: white;
border-radius: 8px;
border: none;
}
QPushButton:hover { background-color: #c0392b; }
QPushButton:pressed { background-color: #a93226; }
""")
button_layout.addWidget(login_button)
button_layout.addWidget(register_button)
button_layout.addWidget(exit_button)
main_layout.addWidget(button_frame, 0, Qt.AlignCenter)
hint_label = QLabel("提示:请使用您的用户名和密码登录")
hint_label.setFont(QFont("微软雅黑", 10))
hint_label.setStyleSheet("color: #95a5a6;")
hint_label.setAlignment(Qt.AlignCenter)
main_layout.addWidget(hint_label)
self.setLayout(main_layout)
def load_credentials(self):
credentials = {}
try:
with open("user.txt", "r", encoding="utf-8") as file:
for line in file:
# 分割每一行来获取用户名和密码
username, password = line.strip().split("=")
credentials[username] = password
return credentials
except FileNotFoundError:
logging.error("未找到凭证文件")
return {}
except Exception as e:
logging.error(f"加载凭证时出错: {str(e)}")
return {}
def handle_login(self):
credentials = self.load_credentials()
username = self.username_input.text().strip()
password = self.password_input.text().strip()
if username in credentials and credentials[username] == password:
logging.info(f"用户 {username} 登录成功")
QMessageBox.information(self, "成功", "登录成功!")
self.accept()
else:
logging.warning(f"用户 {username} 登录失败")
QMessageBox.warning(self, "错误", "用户名或密码错误!")
self.password_input.clear()
def handle_register(self):
credentials = self.load_credentials()
username = self.username_input.text().strip()
password = self.password_input.text().strip()
if not username or not password:
logging.warning("注册失败:用户名或密码为空")
QMessageBox.warning(self, "错误", "用户名和密码不能为空!")
return
if username in credentials:
logging.warning(f"注册失败:用户名 {username} 已存在")
QMessageBox.warning(self, "错误", "用户名已存在!")
return
credentials[username] = password
try:
with open("user.txt", "w", encoding="utf-8") as file:
for uname, pwd in credentials.items():
file.write(f"{uname}={pwd}\n")
logging.info(f"用户 {username} 注册成功")
QMessageBox.information(self, "成功", "注册成功!请登录。")
self.username_input.clear()
self.password_input.clear()
except Exception as e:
logging.error(f"保存凭证时出错: {str(e)}")
QMessageBox.critical(self, "错误", f"保存凭证时出错: {str(e)}")
if __name__ == '__main__':
app = QApplication(sys.argv)
login_dialog = LoginDialog()
if login_dialog.exec_() == QDialog.Accepted:
window = MainWindow()
window.show()
sys.exit(app.exec_());这个代码的功能是什么