代码收藏家技术教程 2025-01-11

Python中实现YOLO目标检测

文章目录

Python中实现YOLO目标检测

一、引言

二、环境准备

1、安装依赖

2、下载预训练模型

三、目标检测

1、图像检测

2、视频检测

四、使用示例

1、轨迹追踪

五、总结

Python中实现YOLO目标检测

一、引言

YOLO（You Only Look Once）是一种流行的实时目标检测算法，以其速度快和准确性高而闻名。在Python中，我们可以使用Ultralytics提供的YOLO模型来实现目标检测。本文将介绍如何在Python中使用YOLO进行目标检测，并提供代码示例。

二、环境准备

1、安装依赖

在开始之前，确保你的环境中安装了必要的库。主要依赖是torch和opencv-python，可以通过以下命令安装：

pip install torch opencv-python

2、下载预训练模型

在Python中使用YOLOv8进行目标检测，首先需要下载预训练的模型。Ultralytics提供了多种预训练模型，这些模型在不同的任务上进行了优化，例如目标检测、实例分割、姿态估计等。以下是如何在Python中加载预训练的YOLOv8模型的示例代码：

import torch

# 加载预训练的YOLOv8模型
# 'yolov8n' 是模型的名称，'n' 表示轻量级版本
model = torch.hub.load('ultralytics/yolov8', 'yolov8n')

这段代码通过torch.hub.load方法直接从Ultralytics的GitHub仓库下载并加载预训练的YOLOv8模型。你可以根据需要选择不同的模型版本，例如yolov8s、yolov8m、yolov8l和yolov8x，这些模型在大小和性能上有所不同，以适应不同的应用场景。

Ultralytics的官方文档提供了详细的模型列表和特点，包括每种模型支持的任务和操作模式。例如，如果你需要进行实例分割，可以选择yolov8n-seg模型；如果需要进行姿态估计，则可以选择yolov8n-pose模型。这些模型在首次使用时会自动从Ultralytics的发布页面下载。

通过这种方式，你可以轻松地在Python项目中集成YOLOv8的强大功能，无需从头开始训练模型，大大节省了时间和计算资源。

三、目标检测

1、图像检测

使用YOLO模型对单张图像进行目标检测的步骤如下：

import cv2

# 读取图像
img_path = 'path_to_your_image.jpg'
img = cv2.imread(img_path)

# 使用YOLO模型进行检测
results = model(img)

# 显示检测结果
results.show()

2、视频检测

对于视频流的目标检测，你可以使用以下代码：

video_path = 'path_to_your_video.mp4'
cap = cv2.VideoCapture(video_path)

while cap.isOpened():
    success, frame = cap.read()
    if success:
        # 在帧上运行YOLOv8追踪
        results = model(frame, persist=True)
        # 展示结果
        annotated_frame = results[0].plot()
        cv2.imshow("YOLOv8 Tracking", annotated_frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    else:
        break

cap.release()
cv2.destroyAllWindows()

四、使用示例

1、轨迹追踪

在视频中识别并追踪物体的运行轨迹，可以使用以下代码：

from collections import defaultdict
import cv2
import numpy as np

# 加载模型
model = torch.hub.load('ultralytics/yolov8', 'yolov8n')

# 打开视频
video_path = "0.mp4"
cap = cv2.VideoCapture(video_path)

# 保存历史记录
track_history = defaultdict(lambda: [])

while cap.isOpened():
    success, frame = cap.read()
    if success:
        # 追踪
        results = model(frame, persist=True)
        # 读取当前帧中物体的框
        boxes = results[0].boxes.xywh.cpu()
        # 获取每个被追踪的物体 id
        track_ids = results[0].boxes.id.int().cpu().tolist()
        annotated_frame = results[0].plot()
        # 绘制
        for box, track_id in zip(boxes, track_ids):
            x, y, w, h = box
            track = track_history[track_id]
            track.append((float(x), float(y)))
            if len(track) > 30:
                track.pop(0)
            points = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
            cv2.polylines(annotated_frame, [points], isClosed=False, color=(230, 230, 230), thickness=10)
        cv2.imshow("YOLOv8 Tracking", annotated_frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    else:
        break

cap.release()
cv2.destroyAllWindows()