在视频中使用 YOLOv5 跟踪对象并打印边界框坐标

2024-03-09 06:27:56

在计算机视觉领域，目标检测是一项关键技术，广泛应用于监控系统、自动驾驶、机器人导航等领域。YOLOv5 作为目前最先进的目标检测算法之一，以其高效和准确的特性受到了广泛关注。本文将详细介绍如何使用 YOLOv5 检测视频中的动态对象，并打印出它们的边界框坐标。

问题描述

在实际应用中，我们经常需要从视频中提取特定对象的边界框坐标。例如，在安防监控中，我们需要实时跟踪并记录可疑人员的移动轨迹。传统的目标检测方法往往无法满足实时性和准确性的要求，而 YOLOv5 则能够很好地解决这一问题。

解决方案

1. 导入必要的库

首先，我们需要导入一些必要的库，包括 OpenCV 和 YOLOv5。

import cv2
from ultralytics import YOLO

2. 加载 YOLOv5 模型

接下来，我们需要加载预训练的 YOLOv5 模型。YOLOv5 提供了多种不同大小的模型，可以根据实际需求选择合适的模型。

model = YOLO('yolov5s.pt')  # 选择合适的预训练模型

3. 打开视频文件

使用 OpenCV 的 VideoCapture 函数打开要分析的视频文件。

cap = cv2.VideoCapture("path_to_video.mp4")

4. 遍历视频帧

使用 while 循环遍历视频中的每一帧。

while cap.isOpened():
    success, frame = cap.read()
    if not success:
        break

5. 运行 YOLOv5 跟踪

对每一帧，使用 YOLOv5 的 track 方法进行目标检测和跟踪。

    results = model.track(frame, persist=True, classes=[0, 1])  # 指定要检测的类别

6. 可视化结果

使用 OpenCV 的绘图函数可视化检测到的对象和它们的边界框。

    for result in results:
        boxes = result.boxes.cpu().numpy()
        for box in boxes:
            x1, y1, x2, y2 = box.xyxy[0].astype(int)
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)

7. 打印边界框坐标

对于每个检测到的对象，打印其边界框坐标和帧号。

    for result in results:
        boxes = result.boxes.cpu().numpy()
        for box in boxes:
            x1, y1, x2, y2 = box.xyxy[0].astype(int)
            frame_number = cap.get(cv2.CAP_PROP_POS_FRAMES)
            print(f"Frame: {frame_number} | Bounding Box Coordinates: ({x1}, {y1}, {x2}, {y2})")

示例代码

以下是完整的示例代码：

import cv2
from ultralytics import YOLO

# 加载 YOLOv5 模型
model = YOLO('yolov5s.pt')

# 打开视频文件
cap = cv2.VideoCapture("path_to_video.mp4")

while cap.isOpened():
    success, frame = cap.read()
    if not success:
        break

    # 运行 YOLOv5 跟踪
    results = model.track(frame, persist=True, classes=[0, 1])

    # 可视化结果
    for result in results:
        boxes = result.boxes.cpu().numpy()
        for box in boxes:
            x1, y1, x2, y2 = box.xyxy[0].astype(int)
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)

    # 打印边界框坐标
    for result in results:
        boxes = result.boxes.cpu().numpy()
        for box in boxes:
            x1, y1, x2, y2 = box.xyxy[0].astype(int)
            frame_number = cap.get(cv2.CAP_PROP_POS_FRAMES)
            print(f"Frame: {frame_number} | Bounding Box Coordinates: ({x1}, {y1}, {x2}, {y2})")

    # 显示帧
    cv2.imshow('Frame', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放资源
cap.release()
cv2.destroyAllWindows()