使用YOLOv8进行实时推理 分类

564 阅读2分钟

简介

YOLOv8是ultralytics公司在YOLOv5上的一个更新版本,目前支持图像分类、物体检测和实例分割任务。本篇主要内容是如何使用YOLOv8进行实时物体检测。ultralytics GitHub地址

  1. 首先需要在项目目录下下载ultralytics(这里因为一些原因选择使用gitcode地址下载)

git clone https://gitcode.com/ultralytics/ultralytics.git

  1. 在ultralytics gihub地址下可以下载预测权重,这里选择yolov8n.pt(yolov8是个模型簇,从小到大包括:yolov8n、yolov8s、yolov8m、yolov8l、yolov8x)

image.png

3.终端进入ultralytics目录下,pip install -r requirements.txt安装所需依赖

image.png

4.使用PyQt5进行屏幕截图,将图片格式转变为opencv可以读取的格式方便回显。最后使用model进行推理,最后将推理结果使用opencv进行回显这时候简单的推理就实现了。(没有使用result的show方法来回显是因为比较慢)

import sys
import numpy as np
from PyQt5.QtWidgets import QApplication
from PyQt5.QtGui import QImage
from ultralytics import YOLO
import torch
import cv2

# 创建设备对象,指定使用 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = YOLO("yolov8n.pt").to(device)

app = QApplication(sys.argv)

while True:
    # 获取整个屏幕的截图
    screen = app.primaryScreen()
    screenshot = screen.grabWindow(app.desktop().winId())

    # 将截图转换为OpenCV格式
    screenshot_image = screenshot.toImage()
    screenshot_image = screenshot_image.convertToFormat(QImage.Format_RGB888)
    width = screenshot.width()
    height = screenshot.height()
    ptr = screenshot_image.constBits()
    ptr.setsize(screenshot_image.byteCount())
    screenshot_cv = np.array(ptr).reshape(height, width, 3).copy()
    image = cv2.cvtColor(screenshot_cv, cv2.COLOR_BGR2RGB)
    results = model(image)
    annotated_frame = results[0].plot()
    # 显示截图
    cv2.imshow('Screenshot', annotated_frame)
    # 检测按键,按下 'q' 键退出循环
    if cv2.waitKey(1) == ord('q'):
        break
# 释放资源
cv2.destroyAllWindows()

5.如果想对推理结果进行进一步操作可以参考官方给出的参数对result进一步操作

pred: 推理结果的列表。每个元素都是一个字典,包含了检测到的物体的相关信息,如类别、置信度、边界框坐标等。
xyxy: 获取边界框的坐标信息。这可能是一个 `Numpy` 数组,其中每一行包含一个边界框的坐标信息。
pandas: 将推理结果转换为 Pandas DataFrame 格式,方便进行进一步的数据处理和分析。
numpy(): 将推理结果转换为 Numpy 数组。
labels: 获取检测到的物体的类别标签。
pred: 获取预测的边界框信息。
plot(): 绘制推理结果,将边界框和类别标签叠加在图像上。
show(): 显示图像和推理结果。
save(): 将带有推理结果的图像保存到文件中。