使用 YOLO 和计算机视觉进行目标跟踪 | 附完整代码

点击下方卡片，关注“小白玩转Python”公众号

YOLO 是一种能够实时进行目标检测的深度学习算法。您可以使用预训练的 YOLO 模型，如 YOLOv8 或 YOLOv9，或者在需要时在自定义数据集上训练自己的模型。在本文中，我将带您了解如何使用预训练的 YOLO 模型进行目标跟踪。这是最简单的教程，我们只处理简单的目标检测。

让我们看看目标检测所涉及的步骤。

1. 我们将加载模型

2. 我们将加载视频

3. 我们将逐帧提取视频

4. 对每一帧进行预测并使用 OpenCV 显示结果

下面的图像总结了这些步骤：

目标检测步骤 让我们开始编码。

第 0 步：环境设置

我们需要两个库 ultralytics 和 opencv，您可以使用以下代码安装这些库。

pip install ultralytics opencv-python

导入所有必要的库。

import cv2 as cv
from ultralytics import YOLO

第 1 步：加载模型 

#load the model
model = YOLO('yolov8n.pt')

在 YOLO 中加载模型非常容易。我们可以直接使用 ultralytics 提供的 YOLO() 函数，并指定模型名称。一旦运行，它会自动在代码所在的文件夹中下载指定的 YOLO 模型。

第 2 步：加载视频 

我们使用 OpenCV 的 VideoCapture() 函数来加载视频。这个函数的输入是视频的路径。

path = "../videos/football_test.mp4"
#read video
vs = cv.VideoCapture(path)

第 3 步：提取帧 

一旦我们加载了视频，在下一步我们将逐个提取视频中的帧。视频是由连续拼接在一起的帧组成的。

• vs = cv.VideoCapture(path)：包含指向视频的指针。

• vs.read()：允许我们逐帧读取视频。一旦我们调用vs上的read函数，它将以迭代的方式返回当前的视频帧。

• cv.imshow('image', frame)：显示每个提取出的帧。

path = "../videos/football_test.mp4"
#read video
vs = cv.VideoCapture(path)


#load the model
model = YOLO('yolov8n.pt')
while True:
    (grabbed,frame) = vs.read()
    if not grabbed:
        break
    cv.imshow('image', frame)
    cv.waitKey(500)

第 4 步：应用预测 

模型的predict()方法是用于对每一帧进行预测的方法。预测结果存储在results变量中。在while循环中，我们从视频中提取每一帧，并在每一帧上进行预测。简单来说，我正在尝试在视频中的每一帧中检测对象。

results = model.predict(frame,stream=False)


path = "../videos/football_test.mp4"
    #read video
    vs = cv.VideoCapture(path)
    #load the model
    model = YOLO('yolov8n.pt')
    while True:
        (grabbed,frame) = vs.read()
        if not grabbed:
            break
        results = model.predict(frame,stream=False)
        cv.imshow('image', frame)
        cv.waitKey(500)

第 5 步：绘制边界框

为了绘制边界框，我创建了一个函数，该函数将输入图像、数据（包含预测结果）、检测到的类别名称，并返回带有检测到的对象边界框的图像。drawBox(data, image, name)：data是预测详情，image是帧，name是检测到的类别名称，例如：人、汽车等。

def drawBox(data,image,name):
    x1, y1, x2, y2, conf, id = data
    p1 = (int(x1), int(y1))
    p2 = (int(x2), int(y2))
    cv.rectangle(image, p1, p2, (0, 0, 255), 3)
    cv.putText(image, name, p1, cv.FONT_HERSHEY_SIMPLEX, 3, (0, 0, 255), 3)
    return image
    
results = model.predict(frame,stream=False)
for result in results:
     for data in result.boxes.data.tolist():
           #print(data)
           id = data[5]
           drawBox(data, frame,detection_classes[id])

由于结果包含了帧中所有检测到的对象的详细信息，我们在上述代码中对每个检测到的对象进行解码。对于每个检测，我们调用`drawBox(data, frame, detection_classes[id])`来在检测周围绘制边界框。

第 6 步：完整代码

您也可以在此处访问完整代码：https://github.com/sunnykumar1516/YOLO/blob/main/ReadVideoApplyYoloV8.py。

import cv2 as cv
from ultralytics import YOLO


def drawBox(data,image,name):
    x1, y1, x2, y2, conf, id = data
    p1 = (int(x1), int(y1))
    p2 = (int(x2), int(y2))
    cv.rectangle(image, p1, p2, (0, 0, 255), 3)
    cv.putText(image, name, p1, cv.FONT_HERSHEY_SIMPLEX, 3, (0, 0, 255), 3)
    return image


detection_classes= []
 path = "../videos/football_test.mp4"
    #read video
 vs = cv.VideoCapture(path)
    #load the model
 model = YOLO('yolov8n.pt')
 while True:
    (grabbed,frame) = vs.read()
     if not grabbed:
            break
     results = model.predict(frame,stream=False)
     detection_classes = results[0].names
     for result in results:
         for data in result.boxes.data.tolist():
                #print(data)
                id = data[5]
                drawBox(data, frame,detection_classes[id])

完整代码参考：https://github.com/sunnykumar1516/YOLO/blob/main

·  END  ·

HAPPY LIFE

本文仅供学习交流使用，如有侵权请联系作者删除