一、项目背景
本项目旨在创建一个综合性的视频处理工具,融合了目标检测、姿态估计和视频分割等多种计算机视觉技术。通过提供直观的用户界面,使用户能够轻松地对视频进行处理和分析,满足在监控、视频内容创作、科研等领域对视频深入理解和处理的需求。
二、项目目的
目的在于创建一个基于 PyQt5 的图形用户界面(GUI)应用程序,能够方便地打开视频文件,并提供了三种主要的处理功能:追踪、姿态估计和分割。用户可以通过界面上的按钮选择所需的操作,并从预定义的模型选项中选择合适的模型进行处理。整个应用旨在为用户提供一个直观、易用的视频处理工具,使其能够在一个界面中完成多种与计算机视觉相关的任务,而无需复杂的命令行操作或专业知识。同时,通过实时显示处理后的视频帧,让用户能够直观地看到处理效果。
三、代码概述
这段代码主要使用了 Python 的多个库和模块,包括 os
、sys
、cv2
用于处理文件和图像操作,PyQt5
用于构建图形用户界面(GUI),以及 ultralytics
中的 YOLO
模型进行目标检测等任务。
四、关键代码解析
1.模型选项的定义:
这部分定义了不同任务所对应的模型路径,为后续根据用户选择加载模型提供了基础。
# 定义模型选择的选项
model_options = {
"检测模型": "yolov8n.pt",
"姿态估计模型": "yolov8n-pose.pt",
"分割模型": "yolov8n-seg.pt"
2.MainWindow
类的构造函数:
这里完成了主窗口的初始化工作,包括窗口属性设置、布局安排、按钮创建以及与各种操作对应的事件关联
class MainWindow(QMainWindow):
def __init__(self):
3.打开视频的相关办法:
def open_video(self):
"""
打开视频文件的方法
- 弹出文件选择对话框,获取用户选择的视频文件路径
- 如果未选择文件,弹出错误提示框并返回
- 尝试使用 OpenCV 打开选择的视频文件
- 如果打开成功,启用相关操作按钮,并调用显示视频的方法
- 如果打开失败,弹出错误提示框
"""
video_path, _ = QFileDialog.getOpenFileName(self, "选择视频文件", "", "Video Files (*.mp4 *.avi)")
if not video_path:
QMessageBox.warning(self, "错误", "未选择视频文件")
return
try:
self.cap = cv2.VideoCapture(video_path)
self.start_tracking_button.setEnabled(True)
self.start_estimation_button.setEnabled(True)
self.start_segmentation_button.setEnabled(True)
self.show_video()
except:
QMessageBox.warning(self, "错误", "视频文件错误")
4.各种任务的开始和执行方法:
def start_tracking(self):
"""
开始追踪的方法
- 如果视频捕获对象存在
- 设置追踪标志为 `True`,禁用开始追踪按钮,启用结束追踪按钮
- 根据选择的模型选项获取模型路径并加载模型
- 调用追踪视频的方法
"""
if self.cap is not None:
self.is_tracking = True
self.start_tracking_button.setEnabled(False)
self.stop_tracking_button.setEnabled(True)
selected_option = self.model_combobox.currentText()
model_path = model_options[selected_option]
# 加载模型
self.model = YOLO(model_path, task='track')
self.track_video(self.cap)
开始姿态估计的方法:
def start_estimation(self):
"""
开始姿态估计的方法
- 如果视频捕获对象存在
- 设置姿态估计标志为 `True`,禁用开始姿态估计按钮
- 根据选择的模型选项获取模型路径并加载模型
- 调用姿态估计处理的方法
"""
if self.cap is not None:
self.is_estimating = True
self.start_estimation_button.setEnabled(False)
selected_option = self.model_combobox.currentText()
model_path = model_options[selected_option]
# 加载模型
self.model = YOLO(model_path, task='pose')
self.estimate_pose()
开始分割的方法:
def start_segmentation(self):
"""
开始分割的方法
1. 如果视频捕获对象存在:
- 设置分割标志为 True
- 禁用开始分割按钮
- 根据当前选择的模型选项获取模型路径
- 加载分割模型
- 调用执行分割的方法
"""
if self.cap is not None:
self.is_segmenting = True
self.start_segmentation_button.setEnabled(False)
selected_option = self.model_combobox.currentText()
model_path = model_options[selected_option]
# 加载模型
self.model = YOLO(model_path, task='segment')
self.perform_segmentation()
主程序入口:
if __name__ == "__main__":
"""
主程序入口
1. 创建 PyQt5 应用程序对象
2. 创建主窗口对象
3. 显示主窗口
4. 启动应用程序的主事件循环,并在退出时返回退出状态
"""
app = QApplication(sys.argv)
window = MainWindow()
window.show()
sys.exit(app.exec_())
前段界面窗口 :
前端的追踪模型运行窗口:
超详细项目代码分析:
import os
import sys
import cv2
from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from ultralytics import YOLO
"""
上述代码导入了所需的模块和库,包括操作系统相关模块、系统模块、OpenCV 库、PyQt5 的相关模块以及 YOLO 模型库
"""
# 定义模型选择的选项
model_options = {
"检测模型": "yolov8n.pt",
"姿态估计模型": "yolov8n-pose.pt",
"分割模型": "yolov8n-seg.pt"
}
"""
创建一个字典 `model_options`,其中包含了不同任务的模型路径
"""
class MainWindow(QMainWindow):
def __init__(self):
"""
主窗口的构造函数
- 调用父类 `QMainWindow` 的构造函数
- 设置窗口标题和大小
- 创建中心部件,并将其设置为窗口的中心部件
- 创建垂直布局,并将其设置为中心部件的布局
- 创建水平布局,用于放置各种按钮和组合框
- 创建并设置打开视频按钮,关联点击事件 `open_video`
- 创建并设置开始追踪、结束追踪、开始姿态估计、开始分割和退出按钮,分别关联相应的点击事件
- 创建模型选择组合框,添加选项,并将其添加到水平布局
- 将水平布局添加到垂直布局
- 创建用于显示视频的标签,并添加到垂直布局
- 初始化视频捕获对象、各种操作的标志以及模型对象
"""
super(MainWindow, self).__init__()
self.setWindowTitle("基于YOLOv8的目标检测")
self.setGeometry(100, 100, 800, 600)
self.central_widget = QWidget()
self.setCentralWidget(self.central_widget)
self.layout = QVBoxLayout()
self.central_widget.setLayout(self.layout)
self.horizontal_layout = QHBoxLayout()
self.open_button = QPushButton("打开视频")
self.open_button.setFixedSize(100, 30)
self.open_button.clicked.connect(self.open_video)
self.horizontal_layout.addWidget(self.open_button)
self.start_tracking_button = QPushButton("开始追踪")
self.start_tracking_button.setEnabled(False)
self.start_tracking_button.setFixedSize(100, 30)
self.start_tracking_button.clicked.connect(self.start_tracking)
self.horizontal_layout.addWidget(self.start_tracking_button)
self.stop_tracking_button = QPushButton("结束追踪")
self.stop_tracking_button.setEnabled(False)
self.stop_tracking_button.setFixedSize(100, 30)
self.stop_tracking_button.clicked.connect(self.stop_tracking)
self.horizontal_layout.addWidget(self.stop_tracking_button)
self.start_estimation_button = QPushButton("开始姿态估计")
self.start_estimation_button.setEnabled(False)
self.start_estimation_button.setFixedSize(100, 30)
self.start_estimation_button.clicked.connect(self.start_estimation)
self.horizontal_layout.addWidget(self.start_estimation_button)
self.start_segmentation_button = QPushButton("开始分割")
self.start_segmentation_button.setEnabled(False)
self.start_segmentation_button.setFixedSize(100, 30)
self.start_segmentation_button.clicked.connect(self.start_segmentation)
self.horizontal_layout.addWidget(self.start_segmentation_button)
self.exit_button = QPushButton("退出")
self.exit_button.setFixedSize(100, 30)
self.exit_button.clicked.connect(self.close)
self.horizontal_layout.addWidget(self.exit_button)
self.model_combobox = QComboBox()
self.model_combobox.addItems(model_options.keys())
self.horizontal_layout.addWidget(self.model_combobox)
self.layout.addLayout(self.horizontal_layout)
self.video_label = QLabel()
self.layout.addWidget(self.video_label)
self.cap = None
self.is_tracking = False
self.is_estimating = False
self.is_segmenting = False
self.model = None
def open_video(self):
"""
打开视频文件的方法
- 弹出文件选择对话框,获取用户选择的视频文件路径
- 如果未选择文件,弹出错误提示框并返回
- 尝试使用 OpenCV 打开选择的视频文件
- 如果打开成功,启用相关操作按钮,并调用显示视频的方法
- 如果打开失败,弹出错误提示框
"""
video_path, _ = QFileDialog.getOpenFileName(self, "选择视频文件", "", "Video Files (*.mp4 *.avi)")
if not video_path:
QMessageBox.warning(self, "错误", "未选择视频文件")
return
try:
self.cap = cv2.VideoCapture(video_path)
self.start_tracking_button.setEnabled(True)
self.start_estimation_button.setEnabled(True)
self.start_segmentation_button.setEnabled(True)
self.show_video()
except:
QMessageBox.warning(self, "错误", "视频文件错误")
def show_video(self):
"""
持续读取并显示视频帧的方法
- 只要视频捕获对象处于打开状态,循环读取帧
- 如果读取成功,将帧转换为 `QImage` 和 `QPixmap`,并设置到视频标签中显示,处理应用程序事件
- 如果读取失败,退出循环并释放视频捕获对象
"""
while self.cap.isOpened():
success, frame = self.cap.read()
if success:
qt_image = QImage(frame.data, frame.shape[1], frame.shape[0], QImage.Format_RGB888)
pixmap = QPixmap.fromImage(qt_image)
self.video_label.setPixmap(pixmap)
QApplication.processEvents()
else:
break
self.cap.release()
def start_tracking(self):
"""
开始追踪的方法
- 如果视频捕获对象存在
- 设置追踪标志为 `True`,禁用开始追踪按钮,启用结束追踪按钮
- 根据选择的模型选项获取模型路径并加载模型
- 调用追踪视频的方法
"""
if self.cap is not None:
self.is_tracking = True
self.start_tracking_button.setEnabled(False)
self.stop_tracking_button.setEnabled(True)
selected_option = self.model_combobox.currentText()
model_path = model_options[selected_option]
# 加载模型
self.model = YOLO(model_path, task='track')
self.track_video(self.cap)
def stop_tracking(self):
"""
停止追踪的方法
- 设置追踪标志为 `False`,启用开始追踪按钮,禁用结束追踪按钮
"""
self.is_tracking = False
self.start_tracking_button.setEnabled(True)
self.stop_tracking_button.setEnabled(False)
def track_video(self, cap):
"""
实际执行追踪和显示追踪结果的方法
- 只要视频捕获对象处于打开状态且追踪标志为 `True`,循环读取帧
- 如果读取成功,使用模型进行追踪处理,获取并显示标注后的帧
- 如果读取失败,退出循环,释放视频捕获对象和关闭 OpenCV 窗口
"""
while cap.isOpened() and self.is_tracking:
success, frame = cap.read()
if success:
results = self.model.track(frame, persist=True)
annotated_frame = results[0].plot()
qt_image = QImage(annotated_frame.data, annotated_frame.shape[1], annotated_frame.shape[0], QImage.Format_RGB888)
pixmap = QPixmap.fromImage(qt_image)
self.video_label.setPixmap(pixmap)
QApplication.processEvents()
else:
break
cap.release()
cv2.destroyAllWindows()
def start_estimation(self):
"""
开始姿态估计的方法
- 如果视频捕获对象存在
- 设置姿态估计标志为 `True`,禁用开始姿态估计按钮
- 根据选择的模型选项获取模型路径并加载模型
- 调用姿态估计处理的方法
"""
if self.cap is not None:
self.is_estimating = True
self.start_estimation_button.setEnabled(False)
selected_option = self.model_combobox.currentText()
model_path = model_options[selected_option]
# 加载模型
self.model = YOLO(model_path, task='pose')
self.estimate_pose()
def estimate_pose(self):
"""
执行姿态估计并显示结果的方法
循环条件:只要视频捕获对象处于打开状态并且正在进行姿态估计
1. 从视频捕获对象中读取一帧图像
2. 如果读取成功:
- 使用加载的模型对帧进行姿态估计处理
- 获取处理后的标注帧
- 将标注帧转换为 QImage 和 QPixmap 格式
- 将 QPixmap 设置到视频标签中显示
- 处理应用程序的事件以确保界面更新
3. 如果读取失败,退出循环
4. 最后释放视频捕获对象
"""
while self.cap.isOpened() and self.is_estimating:
success, frame = self.cap.read()
if success:
results = self.model(frame)
annotated_frame = results[0].plot()
qt_image = QImage(annotated_frame.data, annotated_frame.shape[1], annotated_frame.shape[0], QImage.Format_RGB888)
pixmap = QPixmap.fromImage(qt_image)
self.video_label.setPixmap(pixmap)
QApplication.processEvents()
else:
break
self.cap.release()
def start_segmentation(self):
"""
开始分割的方法
1. 如果视频捕获对象存在:
- 设置分割标志为 True
- 禁用开始分割按钮
- 根据当前选择的模型选项获取模型路径
- 加载分割模型
- 调用执行分割的方法
"""
if self.cap is not None:
self.is_segmenting = True
self.start_segmentation_button.setEnabled(False)
selected_option = self.model_combobox.currentText()
model_path = model_options[selected_option]
# 加载模型
self.model = YOLO(model_path, task='segment')
self.perform_segmentation()
def perform_segmentation(self):
"""
执行分割并显示结果的方法
循环条件:只要视频捕获对象处于打开状态并且正在进行分割
1. 从视频捕获对象中读取一帧图像
2. 如果读取成功:
- 使用加载的模型对帧进行分割处理
- 获取处理后的标注帧
- 将标注帧转换为 QImage 和 QPixmap 格式
- 将 QPixmap 设置到视频标签中显示
- 处理应用程序的事件以确保界面更新
3. 如果读取失败,退出循环
4. 最后释放视频捕获对象
"""
while self.cap.isOpened() and self.is_segmenting:
success, frame = self.cap.read()
if success:
results = self.model(frame)
annotated_frame = results[0].plot()
qt_image = QImage(annotated_frame.data, annotated_frame.shape[1], annotated_frame.shape[0], QImage.Format_RGB888)
pixmap = QPixmap.fromImage(qt_image)
self.video_label.setPixmap(pixmap)
QApplication.processEvents()
else:
break
self.cap.release()
if __name__ == "__main__":
"""
主程序入口
1. 创建 PyQt5 应用程序对象
2. 创建主窗口对象
3. 显示主窗口
4. 启动应用程序的主事件循环,并在退出时返回退出状态
"""
app = QApplication(sys.argv)
window = MainWindow()
window.show()
sys.exit(app.exec_())
总结:
通过以上代码的实现,我们成功构建了一个集成了多种计算机视觉任务的 PyQt5 应用程序。用户可以通过简单的操作选择模型和执行相应的任务,并在界面中实时查看处理结果
本文文章链接:YOLOV8的目标检测系统-CSDN博客