实现训练、测试、调用摄像头实时检测和追踪的接口调用脚本

目录

介绍环境准备训练模型图片测试视频测试目标跟踪指标评估结论

介绍

YOLOv8是单阶段目标检测的集大成之作,它可以帮助我们在图像或视频中自动识别和定位感兴趣的目标物体。本文将介绍如何调用YOLOv8的API实现训练、测试和调用摄像头进行实时检测和追踪。

环境准备

import netron

import torch

from PIL import Image

import onnx

import sys

import os

import numpy as np

from pathlib import Path

from typing import Union

import cv2

from ultralytics import YOLO

训练模型

在YOLOv8中使用下面的代码即可实现公共数据集的自动下载和模型训练。如果你需要是使用自制数据集进行训练,首先需要增加新数据集的配置文件xx.yaml,需要准备训练数据集和标注文件。将图像文件和相应的标注文件存储在一个文件夹中,并按照一种标准格式进行标注,YOLOv8默认使用的是YOLO格式标注。然后,使用目标检测框架提供的训练代码,指定训练参数和路径,开始训练模型。具体的训练命令取决于您选用的目标检测框架,以下是一个示例:

def train():

# 加载模型配置文件,这里使用v8的m模型结构

model = YOLO('yolov8m.yaml')

# 做预训练

# model = YOLO('yolov8x.pt')

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

# 训练模型

model.train(data="coco.yaml", epochs=100, imgsz=640)

上述代码中加载模型配置文件时,如果加载的是 .yaml 文件格式,采用的是从头训练的方案。如果加载的是 .pt 格式则采用的是对应的预训练方案,这种方式在训练时会额外下载对应模型的预训练权重(如果下载失败可考虑科学上网),在训练模型代码中配置数据集文件data、训练迭代次数epochs和输入图片尺寸imgsz。通过执行上述命令,训练代码将读取图像和标注文件,然后使用指定的配置文件和预训练权重进行模型训练。训练过程可能需要一些时间,具体取决于数据集的大小和计算资源的性能。

图片测试

在训练完成后,我们可以使用训练的模型权重进行测试。以下是测试图片代码:

def test_img():

# 训练好的模型权重路径

model = YOLO("YOLOv8/runs/detect/train1/weights/best.pt")

# 测试图片的路径

img = cv2.imread("YOLOv8/7.jpg")

res = model(img)

ann = res[0].plot()

while True:

cv2.imshow("yolo", ann)

if cv2.waitKey(1) & 0xFF == ord('q'):

break

# 设置保存图片的路径

cur_path = sys.path[0]

print(cur_path, sys.path)

if os.path.exists(cur_path):

cv2.imwrite(cur_path + "out.jpg", ann)

else:

os.mkdir(cur_path)

cv2.imwrite(cur_path + "out.jpg", ann)

以上代码主要实现了一个测试图片的功能 。首先,通过导入训练好的模型权重来初始化YOLO模型。然后,读取指定路径下的测试图片,并使用YOLO模型对图片进行推理,得到目标检测结果。接下来,通过将检测结果中的目标物体绘制在图片上形成标注框的方式,将标注框的图片显示出来。用户可以通过按下键盘上的"q"键来退出显示。最后,将带有标注框的图片保存到当前路径下作为"out.jpg"的图片文件。 ️

视频测试

使用训练的模型权重进行测试。以下是测试视频代码

def test_video():

model = YOLO("YOLOv8/runs/detect/train1/weights/best.pt")

# 测试视频存放目录

pa = "/home/you/Downloads/l.mp4"

cap = cv2.VideoCapture(pa)

# 调用设备自身摄像头

# cap = cv2.VideoCapture(0) # -1

# 设置视频尺寸

size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),)

# 第一个参数是将检测视频存储的路径

out = cv2.VideoWriter('save.mp4', cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 40, size)

while cap.isOpened():

ret, frame = cap.read()

if ret:

res = model(frame)

ann = res[0].plot()

cv2.imshow("yolo", ann)

out.write(ann)

if cv2.waitKey(1) & 0xFF == ord('q'):

break

cv2.destroyAllWindows()

cap.release()

以上代码首先加载了一个模型权重,然后通过cv2.VideoCapture打开指定路径下的视频文件,获取视频的尺寸。接下来,创建了一个cv2.VideoWriter对象,用于将检测结果保存成新的视频文件。然后进入一个循环中,不断读取视频的帧,并利用加载的模型对每一帧进行检测。检测结果通过plot方法绘制在帧上,并通过cv2.imshow显示在窗口中。同时,将处理后的帧写入到保存视频的对象中,最后通过按下"q"键退出循环,关闭窗口并释放资源。此代码实现了对视频进行目标检测,并在检测的帧上绘制了检测结果,并保存了检测后的视频。

目标追踪

这段代码实现了使用YOLOv8模型实现botsort目标跟踪,并在每一帧上绘制了检测框和目标ID,并保存了跟踪后的视频。

def tracker():

pa = "/home/you/Downloads/l.mp4"

cap = cv2.VideoCapture(pa)

size = (int(cap .get(cv2.CAP_PROP_FRAME_WIDTH)),int(cap .get(cv2.CAP_PROP_FRAME_HEIGHT)),)

model = YOLO("YOLOv8/runs/detect/train1/weights/best.pt")

flag = 0

out = cv2.VideoWriter('save.mp4', cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 40, size)

while True:

if flag<1:

flag += 1

continue

else:

flag +=1

ret, frame = cap.read()

if not ret:

break

results = model.track(frame, persist=True)

boxes = results[0].boxes.xyxy.cpu().numpy().astype(int)

# result.boxes.id.cpu().numpy().astype(int)

try:

ids = results[0].boxes.id.cpu().numpy().astype(int)

for box, id in zip(boxes, ids):

cv2.rectangle(frame, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)

cv2.putText(

frame,

f"Id {id}",

(box[0], box[1]),

cv2.FONT_HERSHEY_SIMPLEX,

1,

(0, 0, 255),

2,

)

except Exception as e:

print(e)

cv2.imshow("frame", frame)

out.write(frame)

if cv2.waitKey(1) & 0xFF == ord("q"):

break

以上代码功能是利用训练好的YOLOv8模型对指定路径的视频进行目标跟踪,并在每一帧上绘制检测框和目标ID,并将处理后的帧保存成新的视频。函数内部首先通过cv2.VideoCapture打开指定路径下的视频文件,获取视频的尺寸。然后加载YOLOv8模型的权重。接下来进入一个循环,循环条件为True。在循环中,首先判断flag的值,当flag小于1时,继续下一次循环,直到flag变为1。这样做是为了跳过视频的第一帧,因为第一帧通常会导致较大的计算误差。然后通过cap.read()读取视频的下一帧,并判断ret的值。如果ret为False,表示视频已经读取完毕,跳出循环。否则,对该帧进行目标跟踪,将跟踪结果的边界框坐标和目标ID提取出来。接着,在原始帧上绘制检测框和目标ID,并通过cv2.imshow展示出来。同时,将处理后的帧写入到保存视频的对象中。最后,通过按下"q"键退出循环,关闭窗口并释放资源。

指标评估

YOLOv8对科研工作者的友好体现在他有完整的指标评估代码。以下的函数能调用内置的指标评估接口:

def predict():

# Load a model

# model = YOLO('yolov8n.pt') # 加载官方的模型权重作评估

model = YOLO('YOLOv8/runs/detect/your/weights/best.pt') # 加载自定义的模型权重作评估

# 评估

metrics = model.val() # 不需要传参,这里定义的模型会自动在训练的数据集上作评估

print(metrics.box.map) # map50-95

print(metrics.box.map50) # map50

print(metrics.box.map75) # map75

print(metrics.box.maps) # 包含每个类别的map50-95列表

☑️代码内部首先导入了"YOLO"类,并通过实例化该类来加载指定的模型权重。可以通过两种方式加载模型权重,一种是通过加载官方的模型权重进行评估,另一种是通过加载自定义的模型权重进行评估。加载完成后,通过调用模型的"val"方法对模型进行评估。评估结果包括了不同阈值下的mAP(mean average precision)指标,分别是map50-95、map50、map75,以及包含每个类别map50-95的列表。最后,通过打印输出这些评估指标,以便用户查看模型评估结果。可以根据需要,选择不同的模型权重进行评估,并获取相应的评估指标。 以下是用于计算 YOLOv8 模型评估指标的类API,我整理成了表格,可以参考调用✨:

AttributesDescriptionp每个类别的准确率。shapes:(nc,)r每个类别的召回率。shapes:(nc,)f1每个类别的F1分数。shapes:(nc,)all_ap所有类别的 AP 分数和所有 IoU 阈值。shapes:(nc,10)ap_class_index每个 AP 分数的类别索引。shapes:(nc,)nc类别数量。

MethodsDescriptionap50()所有类别的 AP 的 IoU 阈值为 0.5。returns: AP 分数列表。shape:(nc,) 或 []ap()所有类别的 AP 的 IoU 阈值从 0.5 到 0.95。returns: AP 分数列表。shape:(nc,) 或 []mp()所有类别的平均精度。returns:浮点型mr()所有类别的平均回忆率。returns:浮点型map50()所有类别的 IoU 阈值为 0.5 时的平均 AP。returns:浮点型map75()所有类别的 IoU 阈值为 0.75 时的平均 AP。returns:浮点型map()所有类别的 IoU 阈值的平均 AP 范围为 0.5 到 0.95。returns:浮点型mean_results()结果平均值。returns: mp、mr、map50、mapclass_result(i)类感知结果。returns: p[i]、r[i]、ap50[i]、ap[i]maps()每个类别的mAP. returns: mAP分数列表, shape: (nc,).fitness()将训练模型建模为指标的加权组合。returns:浮点型update(results)使用新的评估结果更新指标属性。

结论

本文介绍了如何使用YOLOv8目标检测代码实现训练、测试和调用摄像头进行实时检测和追踪的接口调用脚本。通过按照给定的步骤进行操作,您可以轻松地使用目标检测代码实现这些功能,并将其集成到自己的项目中。一切以YOLOv8的官方参考文件为准。

感谢您阅读本篇博客文章! 如果您觉得这篇文章对您有帮助,不妨给我一个关注点赞以示支持!您的支持是我创作的最大动力。= 如果您对这个领域还有更多的疑问或者想要探讨更深入的话题,欢迎在评论区留下您的想法和问题,我将竭诚为您解答。欄 另外,如果您有任何问题或者希望单独交流的内容,也欢迎通过私信与我联系,我将尽快回复您。 再次感谢您的阅读和支持!期待与您在下一篇博客中再次相会!✨

以下是完整代码

import netron

import torch

from PIL import Image

import onnx

import sys

import os

import numpy as np

from pathlib import Path

from typing import Union

import cv2

from ultralytics import YOLO

def train():

# 加载模型配置文件,这里使用v8的m模型结构

model = YOLO('yolov8m.yaml')

# 做预训练

# model = YOLO('yolov8x.pt')

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

# 训练模型

model.train(data="coco.yaml", epochs=100, imgsz=640)

def onnx():

# 使用onnx导出文件

# model = YOLO('yolov8n.pt') # load an official model

model = YOLO('YOLOv8/runs/detect/train1/weights/best.pt') # load a custom trained

# Export the model

model.export(format='onnx')

def test_img():

# 训练好的模型权重路径

model = YOLO("YOLOv8/runs/detect/train1/weights/best.pt")

# 测试图片的路径

img = cv2.imread("YOLOv8/7.jpg")

res = model(img)

ann = res[0].plot()

while True:

cv2.imshow("yolo", ann)

if cv2.waitKey(1) & 0xFF == ord('q'):

break

# 设置保存图片的路径

cur_path = sys.path[0]

print(cur_path, sys.path)

if os.path.exists(cur_path):

cv2.imwrite(cur_path + "out.jpg", ann)

else:

os.mkdir(cur_path)

cv2.imwrite(cur_path + "out.jpg", ann)

def predict():

from ultralytics import YOLO

# Load a model

# model = YOLO('yolov8n.pt') # 加载官方的模型权重作评估

model = YOLO('YOLOv8/runs/detect/your/weights/best.pt') # 加载自定义的模型权重作评估

# 评估

metrics = model.val() # 不需要传参,这里定义的模型会自动在训练的数据集上作评估

print(metrics.box.map) # map50-95

print(metrics.box.map50) # map50

print(metrics.box.map75) # map75

print(metrics.box.maps) # 包含每个类别的map50-95列表

def test_video():

model = YOLO("YOLOv8/runs/detect/train1/weights/best.pt")

# 测试视频存放目录

pa = "/home/you/Downloads/l.mp4"

cap = cv2.VideoCapture(pa)

# 调用设备自身摄像头

# cap = cv2.VideoCapture(0) # -1

# 设置视频尺寸

size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),)

# 第一个参数是将检测视频存储的路径

out = cv2.VideoWriter('save.mp4', cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 40, size)

while cap.isOpened():

ret, frame = cap.read()

if ret:

res = model(frame)

ann = res[0].plot()

cv2.imshow("yolo", ann)

out.write(ann)

if cv2.waitKey(1) & 0xFF == ord('q'):

break

cv2.destroyAllWindows()

cap.release()

def tracker():

pa = "/home/you/Downloads/l.mp4"

cap = cv2.VideoCapture(pa)

size = (int(cap .get(cv2.CAP_PROP_FRAME_WIDTH)),int(cap .get(cv2.CAP_PROP_FRAME_HEIGHT)),)

model = YOLO("YOLOv8/runs/detect/train1/weights/best.pt")

flag = 0

out = cv2.VideoWriter('save.mp4', cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 40, size)

while True:

if flag<1:

flag += 1

continue

else:

flag +=1

ret, frame = cap.read()

if not ret:

break

results = model.track(frame, persist=True)

boxes = results[0].boxes.xyxy.cpu().numpy().astype(int)

# result.boxes.id.cpu().numpy().astype(int)

try:

ids = results[0].boxes.id.cpu().numpy().astype(int)

for box, id in zip(boxes, ids):

cv2.rectangle(frame, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)

cv2.putText(

frame,

f"Id {id}",

(box[0], box[1]),

cv2.FONT_HERSHEY_SIMPLEX,

1,

(0, 0, 255),

2,

)

except Exception as e:

print(e)

cv2.imshow("frame", frame)

out.write(frame)

if cv2.waitKey(1) & 0xFF == ord("q"):

break

# train()

# test_video()

test_img()

# predict()

# tracker()

# onnx()

# 下面是使用netron导出模型结构

# netron.start("YOLOv8/runs/detect/train1/weights/best.onnx")

查看原文