赞
踩
近年来,计算机视觉技术,特别是基于深度学习的目标检测方法,在许多应用领域得到了广泛的关注和应用。YOLO(You Only Look Once)作为一种高效的实时目标检测算法,凭借其速度快、精度高的优势,在学术界和工业界都取得了显著的成果。YOLOv8是YOLO系列最新版本,融合了许多先进的技术和优化策略,进一步提升了目标检测的性能。
本项目旨在利用YOLOv8模型,开发一个集成图像和视频目标检测的桌面应用程序,帮助用户方便地进行目标检测任务。该应用程序主要针对研究人员、工程师以及其他需要进行目标检测任务的用户,提供了一种简单易用的工具,能够快速检测图像或视频中的目标,并展示检测结果。
1. **开发一个用户友好的桌面应用程序**:
- 利用PyQt5开发图形用户界面(GUI),使用户可以方便地导入图像和视频,进行目标检测,并查看检测结果。
- 提供基本的图像和视频导航功能,如加载图像/视频、查看上一张/下一张图像或视频帧等。
2. **集成YOLOv8模型进行目标检测**:
- 使用YOLOv8预训练模型进行目标检测,确保高效、准确地检测图像和视频中的目标。
- 实现目标检测结果的可视化,将检测框和标签叠加在图像或视频帧上,帮助用户直观地了解检测结果。
3. **支持模型训练与优化**:
- 提供训练YOLOv8模型的功能,允许用户使用自定义数据集训练模型。
- 在训练过程中保存最佳模型,确保用户能够获取最优的目标检测模型。
4. **提供结果保存功能**:
- 允许用户将检测结果保存为图像或视频文件,方便后续分析和处理。
项目由两个主要部分组成:前端和后端。
前端部分主要是图形用户界面(GUI)的实现,使用PyQt5开发。主要功能包括:
1. **图像和视频的加载与显示**:
- 导入图像:允许用户选择文件夹并加载其中的所有图像。
- 导入视频:允许用户选择视频文件并逐帧加载。
- 显示图像和视频帧:在GUI中显示原始图像或视频帧,以及检测后的结果。
2. **检测结果的导航和查看**:
- 上一个/下一个按钮:用于导航图像列表或视频帧。
- 开始检测按钮:触发目标检测过程,显示检测结果。
3. **结果的保存与移除**:
- 保存按钮:允许用户将检测结果保存为图像或视频文件。
- 移除按钮:从当前加载的列表中移除图像或视频,清除显示内容。
4. **应用退出功能**:
- 退出按钮:关闭应用程序。
后端部分主要是YOLOv8模型的集成与目标检测的实现。主要功能包括:
1. **模型加载与预测**:
- 加载YOLOv8预训练模型。
- 对导入的图像或视频帧进行目标检测,生成检测结果。
2. **模型训练与优化**:
- 使用指定的数据集和超参数进行模型训练。
- 在训练过程中保存最佳模型。
3. **检测结果的处理与保存**:
- 将检测结果(包括检测框和标签)叠加在图像或视频帧上。
- 保存处理后的图像或视频文件。
通过本项目,我们希望为需要进行目标检测任务的用户提供一个高效、易用的工具。该工具不仅能够快速、准确地检测图像和视频中的目标,还支持用户根据自定义数据集训练模型,满足不同应用场景的需求。未来,我们可以进一步扩展该项目的功能,如支持更多的模型类型、优化检测速度和精度、增加更多的图像和视频处理功能等。
数据预处理是为了提高模型训练效果和加速训练过程,对数据进行的一系列处理操作。常用的数据预处理方法包括:
图像缩放:
归一化:
YOLO模型的标签格式非常简单,每个标注文件对应一个图像,包含多个标注信息。每行代表一个目标,格式如下:
- class_id x_center y_center width height
- 49 0.642859 0.0792187 0.148063 0.148062
数据来源:自行标注200条+https://github.com/ultralytics/ultralytics在gihub上面公开数据集128条
loss变化:
模型训练:
yolo detect train data=datasets/mubiao/my_data.yaml model=yolov8n.yaml pretrained=ultralytics/yolov8n.pt epochs=50 batch=4 lr0=0.01 resume=True
模型评估效果:
运行后参数的变化:精度,召回率......可视化
在终端中运行了训练代码,转换成一般python代码后:
- from ultralytics import YOLO
-
- def train_yolov8():
- # 加载预训练模型
- model = YOLO('ultralytics/yolov8n.pt') #这些代码是根据终端训练模型写出详细的训练代码
-
- # 开始训练
- model.train(
- data='datasets/mubiao/my_data.yaml', # 数据集配置文件路径
- model='yolov8n.yaml', # 模型配置文件路径
- epochs=50, # 训练的轮数
- batch=4, # 批处理大小
- lr0=0.01, # 初始学习率
- resume=True, # 是否从上次中断处恢复训练
- save_period=1 # 每个epoch保存一次模型
- )
-
- # 另存为最优模型
- best_model_path = model.ckpt_path.replace('last', 'best')
- model.save(best_model_path)
-
- if __name__ == "__main__":
- train_yolov8
以上代码就是训练完之后选择最好的模型(best)作为我们的一个项目的使用
- import cv2 # OpenCV库,用于图像处理
- import numpy as np # NumPy库,用于数组操作
- from ultralytics import YOLO # YOLO库,用于加载YOLO模型
-
- class YOLOApp(QWidget): # 定义YOLOApp类,继承自QWidget
- def __init__(self): # 初始化方法
- super().__init__() # 调用父类的初始化方法
-
- self.model = YOLO('best.pt') # 加载YOLO模型
- self.classnameList = self.model.names # 获取类别名列表
- self.imagePaths = [] # 图像路径列表
- self.currentImageIndex = -1 # 当前图像索引
- self.videoPath = None # 视频路径
- self.cap = None # 视频捕获对象
- self.timer = None # 定时器
- self.frameIndex = 0 # 当前帧索引
- self.frames = [] # 视频帧列表
- self.detectedFrames = [] # 检测后的帧列表
- self.initUI() # 初始化UI
-
- def detectObjects(self): # 检测对象的方法
- if self.videoPath: # 如果有视频路径
- self.detectVideo() # 检测视频
- elif hasattr(self, 'imagePath'): # 如果有图像路径
- img = cv2.imread(self.imagePath) # 读取图像
- self.detectFrame(img) # 检测图像中的对象
-
- def detectFrame(self, img): # 检测图像中的对象
- results = self.model.predict(img, stream=True) # 使用模型进行预测
- used_labels = [] # 已使用标签的位置
- for result in results:
- boxes = result.boxes.cpu().numpy() # 获取检测框
- for box in boxes:
- r = box.xyxy[0].astype(int) # 获取框的坐标
- cv2.rectangle(img, (r[0], r[1]), (r[2], r[3]), (0, 0, 255), 2) # 绘制红色矩形框
- classID = int(box.cls[0]) # 获取类别ID
- label = self.classnameList[classID] # 获取标签名称
-
- label_pos = (r[0], r[1] - 10) # 标签位置
- while any(self.overlap(label_pos, used_label_pos) for used_label_pos in used_labels): # 检查是否重叠
- label_pos = (label_pos[0], label_pos[1] - 15) # 调整标签位置
- used_labels.append(label_pos) # 保存已使用标签位置
-
- cv2.putText(img, label, label_pos, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2) # 绘制标签
-
- height, width, channel = img.shape # 获取图像形状
- bytesPerLine = 3 * width # 每行字节数
- qImg = QImage(img.data, width, height, bytesPerLine, QImage.Format_RGB888).rgbSwapped() # 转换图像格式
- self.labelDetected.setPixmap(QPixmap.fromImage(qImg)) # 显示检测后的图像
- self.detectedImage = img # 保存检测后的图像
-
- def detectVideo(self): # 检测视频中的对象
- self.cap = cv2.VideoCapture(self.videoPath) # 打开视频文件
- self.detectedFrames = [] # 初始化检测后的帧列表
- while True:
- ret, frame = self.cap.read() # 读取视频帧
- if not ret:
- break # 如果读取失败,则退出循环
- detected_frame = self.detectFrameForVideo(frame) # 检测帧中的对象
- self.detectedFrames.append(detected_frame) # 保存检测后的帧
- self.cap.release() # 释放视频捕获对象
- self.playDetectedVideo() # 播放检测后的视频
-
- def detectFrameForVideo(self, frame): # 对视频帧进行对象检测
- results = self.model.predict(frame, stream=True) # 使用模型进行预测
- used_labels = [] # 已使用标签的位置
- for result in results:
- boxes = result.boxes.cpu().numpy() # 获取检测框
- for box in boxes:
- r = box.xyxy[0].astype(int) # 获取框的坐标
- cv2.rectangle(frame, (r[0], r[1]), (r[2], r[3]), (0, 0, 255), 2) # 绘制红色矩形框
- classID = int(box.cls[0]) # 获取类别ID
- label = self.classnameList[classID] # 获取标签名称
-
- label_pos = (r[0], r[1] - 10) # 标签位置
- while any(self.overlap(label_pos, used_label_pos) for used_label_pos in used_labels): # 检查是否重叠
- label_pos = (label_pos[0], label_pos[1] - 15) # 调整标签位置
- used_labels.append(label_pos) # 保存已使用标签位置
-
- cv2.putText(frame, label, label_pos, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2) # 绘制标签
- return frame # 返回检测后的帧
学习率 (lr0=0.01
)
批次大小 (batch=4
)
训练轮数 (epochs=50
)
预训练权重 (pretrained=ultralytics/yolov8n.pt
)
- import sys
- from PyQt5.QtWidgets import QApplication, QLabel, QWidget, QVBoxLayout, QPushButton, QFileDialog, QHBoxLayout, QDesktopWidget
- from PyQt5.QtGui import QPixmap, QImage
- import requests
- import cv2
- import os
- import numpy as np
-
- class YOLOApp(QWidget): # 定义主窗口类YOLOApp,继承自QWidget
- def __init__(self):
- super().__init__()
-
- self.imagePaths = [] # 存储图像路径的列表
- self.currentImageIndex = -1 # 当前图像索引,初始值为-1表示没有图像
- self.videoPath = None # 视频路径,初始为None
- self.cap = None # 视频捕获对象,初始为None
- self.timer = None # 定时器对象,初始为None
- self.frameIndex = 0 # 当前视频帧索引,初始值为0
- self.frames = [] # 存储视频帧的列表
- self.detectedFrames = [] # 存储检测后的视频帧的列表
- self.initUI() # 调用initUI方法初始化UI
-
- def initUI(self): # 初始化UI的方法
- self.setWindowTitle('YOLOv8 目标检测') # 设置窗口标题
-
- # 获取屏幕的大小
- screen = QDesktopWidget().screenGeometry()
- self.setGeometry(0, 0, screen.width(), screen.height()) # 设置窗口为全屏大小
-
- self.layout = QVBoxLayout() # 创建一个垂直布局
-
- self.btnLayout = QHBoxLayout() # 创建一个水平布局
-
- self.btnLoadImage = QPushButton('导入图片', self) # 创建“导入图片”按钮
- self.btnLoadImage.clicked.connect(self.loadImages) # 连接按钮点击信号到loadImages方法
- self.btnLayout.addWidget(self.btnLoadImage) # 将按钮添加到水平布局中
-
- self.btnLoadVideo = QPushButton('导入视频', self) # 创建“导入视频”按钮
- self.btnLoadVideo.clicked.connect(self.loadVideo) # 连接按钮点击信号到loadVideo方法
- self.btnLayout.addWidget(self.btnLoadVideo) # 将按钮添加到水平布局中
-
- self.btnPrev = QPushButton('上一个', self) # 创建“上一个”按钮
- self.btnPrev.clicked.connect(self.showPrev) # 连接按钮点击信号到showPrev方法
- self.btnLayout.addWidget(self.btnPrev) # 将按钮添加到水平布局中
-
- self.btnNext = QPushButton('下一个', self) # 创建“下一个”按钮
- self.btnNext.clicked.connect(self.showNext) # 连接按钮点击信号到showNext方法
- self.btnLayout.addWidget(self.btnNext) # 将按钮添加到水平布局中
-
- # 将按钮布局添加到主布局中
- self.layout.addLayout(self.btnLayout)
-
- # 创建一个水平布局用于放置图像
- self.imageLayout = QHBoxLayout()
-
- # 创建一个标签用于显示原始图像
- self.labelOriginal = QLabel(self)
- self.imageLayout.addWidget(self.labelOriginal) # 将标签添加到图像布局中
-
- # 创建一个标签用于显示检测后的图像
- self.labelDetected = QLabel(self)
- self.imageLayout.addWidget(self.labelDetected) # 将标签添加到图像布局中
-
- # 将图像布局添加到主布局中
- self.layout.addLayout(self.imageLayout)
-
- # 设置窗口的主布局
- self.setLayout(self.layout)
-
- def loadImages(self): # 加载图像的方法
- options = QFileDialog.Options() # 创建文件对话框选项
- folder = QFileDialog.getExistingDirectory(self, "选择文件夹", "", options=options) # 打开文件夹选择对话框
- if folder: # 如果选择了文件夹
- # 获取文件夹中所有图像文件的路径并排序
- self.imagePaths = [os.path.join(folder, file) for file in os.listdir(folder) if file.endswith(('.png', '.jpg', '.jpeg'))]
- self.imagePaths.sort()
- self.currentImageIndex = 0 # 将当前图像索引设置为0
- self.videoPath = None # 清空视频路径
- self.frames = [] # 清空视频帧列表
- self.showImage() # 显示当前图像
-
- def loadVideo(self): # 加载视频的方法
- options = QFileDialog.Options() # 创建文件对话框选项
- fileName, _ = QFileDialog.getOpenFileName(self, "选择视频文件", "", "视频文件 (*.mp4 *.avi *.mkv)", options=options) # 打开视频文件选择对话框
- if fileName: # 如果选择了视频文件
- self.videoPath = fileName # 设置视频路径
- self.cap = cv2.VideoCapture(self.videoPath) # 创建视频捕获对象
- self.frames = [] # 清空视频帧列表
- while True: # 循环读取视频帧
- ret, frame = self.cap.read() # 读取一帧
- if not ret: # 如果读取失败则退出循环
- break
- self.frames.append(frame) # 将帧添加到视频帧列表
- self.frameIndex = 0 # 将当前帧索引设置为0
- self.cap.release() # 释放视频捕获对象
- self.imagePaths = [] # 清空图像路径列表
- self.labelOriginal.clear() # 清空原始图像标签
- self.showFrame() # 显示当前视频帧
-
- def showImage(self): # 显示图像的方法
- if 0 <= self.currentImageIndex < len(self.imagePaths): # 如果当前图像索引有效
- self.imagePath = self.imagePaths[self.currentImageIndex] # 获取当前图像路径
- pixmap = QPixmap(self.imagePath) # 加载图像为QPixmap对象
- self.labelOriginal.setPixmap(pixmap) # 在标签上显示图像
- self.labelDetected.clear() # 清除检测后的图像标签内容
-
- def showFrame(self): # 显示视频帧的方法
- if 0 <= self.frameIndex < len(self.frames): # 如果当前帧索引有效
- frame = self.frames[self.frameIndex] # 获取当前帧
- height, width, channel = frame.shape # 获取帧的高宽和通道数
- bytesPerLine = 3 * width # 计算每行的字节数
- qImg = QImage(frame.data, width, height, bytesPerLine, QImage.Format_RGB888).rgbSwapped() # 将帧转换为QImage对象
- self.labelOriginal.setPixmap(QPixmap.fromImage(qImg)) # 在标签上显示原始视频帧
- self.labelDetected.clear() # 清除检测后的图像标签内容
-
- def showPrev(self): # 显示前一个图像或视频帧的方法
- if self.imagePaths: # 如果有图像路径
- if self.currentImageIndex > 0: # 如果当前图像索引大于0
- self.currentImageIndex -= 1 # 减少图像索引
- self.showImage() # 显示前一个图像
- elif self.frames: # 如果有视频帧
- if self.frameIndex > 0: # 如果当前帧索引大于0
- self.frameIndex -= 1 # 减少帧索引
- self.showFrame() # 显示前一个视频帧
-
- def showNext(self): # 显示下一个图像或视频帧的方法
- if self.imagePaths: # 如果有图像路径
- if self.currentImageIndex < len(self.imagePaths) - 1: # 如果当前图像索引小于图像数量-1
- self.currentImageIndex += 1 # 增加图像索引
- self.showImage() # 显示下一个图像
- elif self.frames: # 如果有视频帧
- if self.frameIndex < len(self.frames) - 1: # 如果当前帧索引小于视频帧数量-1
- self.frameIndex += 1 # 增加帧索引
- self.showFrame() # 显示下一个视频帧
前端实现过程
1. 导入必要的库
代码首先导入了实现图像和视频处理、界面开发所需的库,包括`PyQt5`、`cv2`、`numpy`等。
2. 创建主窗口类
定义了一个名为`YOLOApp`的类,继承自`QWidget`,表示主窗口。
3. 初始化类
在`__init__`方法中,初始化一些变量,如图像路径列表、当前图像索引、视频路径、视频捕获对象、帧列表等。然后调用`initUI`方法来设置用户界面。
4. 初始化用户界面
`initUI`方法中设置窗口标题,获取屏幕大小并将窗口设置为全屏。创建一个垂直布局`QVBoxLayout`来容纳其他控件,并创建一个水平布局`QHBoxLayout`来容纳按钮。
创建了“导入图片”、“导入视频”、“上一个”、“下一个”按钮,并将它们添加到按钮布局中。为每个按钮设置点击事件处理函数。
然后创建两个标签`QLabel`,一个用于显示原始图像,一个用于显示检测后的图像,并将它们添加到一个水平布局中。将该布局添加到主布局中,并设置主窗口的布局。
5. 加载图像功能
`loadImages`方法中,通过文件对话框选择图像文件夹,获取所有图像路径并排序。设置当前图像索引为0,清空视频路径和帧列表,然后调用`showImage`方法显示当前图像。
6. 加载视频功能
`loadVideo`方法中,通过文件对话框选择视频文件,创建视频捕获对象,循环读取视频帧并存储在帧列表中。设置当前帧索引为0,清空图像路径列表并清除原始图像标签,然后调用`showFrame`方法显示当前视频帧。
7. 显示图像和视频帧
`showImage`方法中,如果当前图像索引有效,获取当前图像路径并加载图像,在标签上显示原始图像,并清除检测后的图像标签内容。
`showFrame`方法中,如果当前帧索引有效,获取当前帧并将其转换为`QImage`对象,在标签上显示原始视频帧,并清除检测后的图像标签内容。
8. 导航功能
`showPrev`和`showNext`方法分别实现了显示前一个和下一个图像或视频帧的功能,通过修改当前图像或帧索引,调用`showImage`或`showFrame`方法来显示相应的图像或帧。
总结
本代码实现了一个图像和视频目标检测的前端应用程序,用户可以通过简洁的界面加载图像和视频,浏览和查看检测结果。PyQt5提供了良好的用户体验,OpenCV用于图像和视频处理,结合YOLOv8模型实现目标检测功能,为目标检测应用提供了一个完整的解决方案。
- import sys
- from PyQt5.QtWidgets import QApplication, QLabel, QWidget, QVBoxLayout, QPushButton, QFileDialog, QHBoxLayout, QDesktopWidget
- from PyQt5.QtGui import QPixmap, QImage
- import cv2
- from ultralytics import YOLO
- import os
- import numpy as np
-
- class YOLOApp(QWidget):
- def __init__(self):
- super().__init__()
-
- self.model = YOLO('best.pt')
- self.classnameList = self.model.names
- self.imagePaths = []
- self.currentImageIndex = -1
- self.videoPath = None
- self.cap = None
- self.timer = None
- self.frameIndex = 0
- self.frames = []
- self.detectedFrames = []
- self.initUI()
-
- def initUI(self):
- self.setWindowTitle('YOLOv8 目标检测')
-
- # 获取屏幕的大小
- screen = QDesktopWidget().screenGeometry()
- self.setGeometry(0, 0, screen.width(), screen.height()) # 设置窗口为全屏大小
-
- self.layout = QVBoxLayout()
-
- self.btnLayout = QHBoxLayout()
-
- self.btnLoadImage = QPushButton('导入图片', self)
- self.btnLoadImage.clicked.connect(self.loadImages)
- self.btnLayout.addWidget(self.btnLoadImage)
-
- self.btnLoadVideo = QPushButton('导入视频', self)
- self.btnLoadVideo.clicked.connect(self.loadVideo)
- self.btnLayout.addWidget(self.btnLoadVideo)
-
- self.btnPrev = QPushButton('上一个', self)
- self.btnPrev.clicked.connect(self.showPrev)
- self.btnLayout.addWidget(self.btnPrev)
-
- self.btnNext = QPushButton('下一个', self)
- self.btnNext.clicked.connect(self.showNext)
- self.btnLayout.addWidget(self.btnNext)
-
- self.btnDetect = QPushButton('开始检测', self)
- self.btnDetect.clicked.connect(self.detectObjects)
- self.btnLayout.addWidget(self.btnDetect)
-
- self.btnSave = QPushButton('保存', self)
- self.btnSave.clicked.connect(self.save)
- self.btnLayout.addWidget(self.btnSave)
-
- self.btnRemove = QPushButton('移除', self)
- self.btnRemove.clicked.connect(self.remove)
- self.btnLayout.addWidget(self.btnRemove)
-
- self.btnExit = QPushButton('退出', self)
- self.btnExit.clicked.connect(self.closeApp)
- self.btnLayout.addWidget(self.btnExit)
-
- self.layout.addLayout(self.btnLayout)
-
- self.imageLayout = QHBoxLayout()
-
- self.labelOriginal = QLabel(self)
- self.imageLayout.addWidget(self.labelOriginal)
-
- self.labelDetected = QLabel(self)
- self.imageLayout.addWidget(self.labelDetected)
-
- self.layout.addLayout(self.imageLayout)
-
- self.setLayout(self.layout)
-
- def loadImages(self):
- options = QFileDialog.Options()
- folder = QFileDialog.getExistingDirectory(self, "选择文件夹", "", options=options)
- if folder:
- self.imagePaths = [os.path.join(folder, file) for file in os.listdir(folder) if file.endswith(('.png', '.jpg', '.jpeg'))]
- self.imagePaths.sort()
- self.currentImageIndex = 0
- self.videoPath = None
- self.frames = []
- self.showImage()
-
- def loadVideo(self):
- options = QFileDialog.Options()
- fileName, _ = QFileDialog.getOpenFileName(self, "选择视频文件", "", "视频文件 (*.mp4 *.avi *.mkv)", options=options)
- if fileName:
- self.videoPath = fileName
- self.cap = cv2.VideoCapture(self.videoPath)
- self.frames = []
- while True:
- ret, frame = self.cap.read()
- if not ret:
- break
- self.frames.append(frame)
- self.frameIndex = 0
- self.cap.release()
- self.imagePaths = []
- self.labelOriginal.clear()
- self.showFrame()
-
- def showImage(self):
- if 0 <= self.currentImageIndex < len(self.imagePaths):
- self.imagePath = self.imagePaths[self.currentImageIndex]
- pixmap = QPixmap(self.imagePath)
- self.labelOriginal.setPixmap(pixmap) # 在左边显示原图
- self.labelDetected.clear() # 清除右边检测图的内容
-
- def showFrame(self):
- if 0 <= self.frameIndex < len(self.frames):
- frame = self.frames[self.frameIndex]
- height, width, channel = frame.shape
- bytesPerLine = 3 * width
- qImg = QImage(frame.data, width, height, bytesPerLine, QImage.Format_RGB888).rgbSwapped()
- self.labelOriginal.setPixmap(QPixmap.fromImage(qImg)) # 在左边显示原图
- self.labelDetected.clear() # 清除右边检测图的内容
-
- def showPrev(self):
- if self.imagePaths:
- if self.currentImageIndex > 0:
- self.currentImageIndex -= 1
- self.showImage()
- elif self.frames:
- if self.frameIndex > 0:
- self.frameIndex -= 1
- self.showFrame()
-
- def showNext(self):
- if self.imagePaths:
- if self.currentImageIndex < len(self.imagePaths) - 1:
- self.currentImageIndex += 1
- self.showImage()
- elif self.frames:
- if self.frameIndex < len(self.frames) - 1:
- self.frameIndex += 1
- self.showFrame()
-
- def detectObjects(self):
- if self.videoPath:
- self.detectVideo()
- elif hasattr(self, 'imagePath'):
- img = cv2.imread(self.imagePath)
- self.detectFrame(img)
-
- def detectFrame(self, img):
- results = self.model.predict(img, stream=True)
- used_labels = []
- for result in results:
- boxes = result.boxes.cpu().numpy()
- for box in boxes:
- r = box.xyxy[0].astype(int)
- cv2.rectangle(img, (r[0], r[1]), (r[2], r[3]), (0, 0, 255), 2) # 改成红色
- classID = int(box.cls[0])
- label = self.classnameList[classID]
-
- # 检查是否重叠并调整标签位置
- label_pos = (r[0], r[1] - 10)
- while any(self.overlap(label_pos, used_label_pos) for used_label_pos in used_labels):
- label_pos = (label_pos[0], label_pos[1] - 15)
- used_labels.append(label_pos)
-
- cv2.putText(img, label, label_pos, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2) # 改成红色
-
- # 转换图像格式以便在PyQt中显示
- height, width, channel = img.shape
- bytesPerLine = 3 * width
- qImg = QImage(img.data, width, height, bytesPerLine, QImage.Format_RGB888).rgbSwapped()
- self.labelDetected.setPixmap(QPixmap.fromImage(qImg)) # 在右边显示检测后的图片
- self.detectedImage = img # 保存检测后的图像
-
- def detectVideo(self):
- self.cap = cv2.VideoCapture(self.videoPath)
- self.detectedFrames = []
- while True:
- ret, frame = self.cap.read()
- if not ret:
- break
- detected_frame = self.detectFrameForVideo(frame)
- self.detectedFrames.append(detected_frame)
- self.cap.release()
- self.playDetectedVideo()
-
- def detectFrameForVideo(self, frame):
- results = self.model.predict(frame, stream=True)
- used_labels = []
- for result in results:
- boxes = result.boxes.cpu().numpy()
- for box in boxes:
- r = box.xyxy[0].astype(int)
- cv2.rectangle(frame, (r[0], r[1]), (r[2], r[3]), (0, 0, 255), 2)
- classID = int(box.cls[0])
- label = self.classnameList[classID]
-
- label_pos = (r[0], r[1] - 10)
- while any(self.overlap(label_pos, used_label_pos) for used_label_pos in used_labels):
- label_pos = (label_pos[0], label_pos[1] - 15)
- used_labels.append(label_pos)
-
- cv2.putText(frame, label, label_pos, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
- return frame
-
- def playDetectedVideo(self):
- self.frameIndex = 0
- self.labelOriginal.clear() # 清除左边的原视频
- self.timer = self.startTimer(30)
-
- def timerEvent(self, event):
- if self.frameIndex < len(self.detectedFrames):
- frame = self.detectedFrames[self.frameIndex]
- height, width, channel = frame.shape
- bytesPerLine = 3 * width
- qImg = QImage(frame.data, width, height, bytesPerLine, QImage.Format_RGB888).rgbSwapped()
- self.labelDetected.setPixmap(QPixmap.fromImage(qImg))
- self.frameIndex += 1
- else:
- self.killTimer(self.timer)
-
- def overlap(self, pos1, pos2, threshold=10):
- return abs(pos1[0] - pos2[0]) < threshold and abs(pos1[1] - pos2[1]) < threshold
-
- def save(self):
- if hasattr(self, 'detectedImage'):
- options = QFileDialog.Options()
- filePath, _ = QFileDialog.getSaveFileName(self, "保存图片", "", "JPEG (*.jpg;*.jpeg);;PNG (*.png)", options=options)
- if filePath:
- cv2.imwrite(filePath, self.detectedImage)
- elif self.detectedFrames:
- options = QFileDialog.Options()
- filePath, _ = QFileDialog.getSaveFileName(self, "保存视频", "", "MP4 (*.mp4);;AVI (*.avi)", options=options)
- if filePath:
- height, width, layers = self.detectedFrames[0].shape
- fourcc = cv2.VideoWriter_fourcc(*'mp4v') if filePath.endswith('.mp4') else cv2.VideoWriter_fourcc(*'XVID')
- out = cv2.VideoWriter(filePath, fourcc, 20.0, (width, height))
- for frame in self.detectedFrames:
- out.write(frame)
- out.release()
-
- def remove(self):
- if self.imagePaths:
- del self.imagePaths[self.currentImageIndex]
- if self.currentImageIndex >= len(self.imagePaths):
- self.currentImageIndex = len(self.imagePaths) - 1
- self.showImage() if self.imagePaths else self.labelOriginal.clear()
- self.labelDetected.clear()
- elif self.frames:
- self.frames = []
- self.detectedFrames = []
- self.frameIndex = 0
- self.labelOriginal.clear()
- self.labelDetected.clear()
-
- def closeApp(self):
- self.close()
-
- def closeEvent(self, event):
- event.accept()
-
- if __name__ == '__main__':
- if not QApplication.instance():
- app = QApplication(sys.argv)
- else:
- app = QApplication.instance()
- ex = YOLOApp()
- ex.show()
- sys.exit(app.exec_())
运行前端界面:
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。