实现训练、测试、调用摄像头实时检测和追踪的接口调用脚本
目录
- 介绍
- 训练模型
- 图片测试
- 视频流测试
- 目标跟踪
- 结论
介绍
YOLOv8是单阶段目标检测的集大成之作，它可以帮助我们在图像或视频中自动识别和定位感兴趣的目标物体。本文将介绍如何调用YOLOv8的API实现训练、测试和调用摄像头进行实时检测和追踪。
训练模型
首先，我们需要准备训练数据集和标注文件。将图像文件和相应的标注文件存储在一个文件夹中，并按照一种标准格式进行标注，如YOLO格式或COCO格式。然后，使用目标检测框架提供的训练代码，指定训练参数和路径，开始训练模型。在YOLOv8中使用下面的代码即可实现数据集的自动下载和模型训练。
具体的训练命令取决于您选用的目标检测框架，以下是一个示例：
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)
视频流测试
使用训练的模型权重进行测试。以下是测试视频代码
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'):breakcv2.destroyAllWindows()cap.release()设置检测的视频路径，运行程序，实时保存视频内容，按 q 退出。
目标追踪
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 = 0out = cv2.VideoWriter('save.mp4', cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 40, size)while True:if flag<1:flag += 1continueelse:flag +=1ret, frame = cap.read()if not ret:breakresults = 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以上的代码使用默认的botsort进行视频跟踪
结论
本文介绍了如何使用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 YOLOdef 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 modelmodel = YOLO('YOLOv8/runs/detect/train1/weights/best.pt') # load a custom trained# Export the modelmodel.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 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'):breakcv2.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 = 0out = cv2.VideoWriter('save.mp4', cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 40, size)while True:if flag<1:flag += 1continueelse:flag +=1ret, frame = cap.read()if not ret:breakresults = 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()
# tracker()
# onnx()# 下面是使用netron导出模型结构
# netron.start("YOLOv8/runs/detect/train1/weights/best.onnx")
