yolov8 训练模型

一、准备数据

1.1 收集数据

以拳皇为例，将录制的游戏视频进行抽帧。

import cv2
import os

# 视频文件路径
video_path = '1.mp4'

# 输出帧的保存目录
output_dir = './output_frames'
os.makedirs(output_dir, exist_ok=True)

# 读取视频
cap = cv2.VideoCapture(video_path)

# 获取视频帧率
fps = cap.get(cv2.CAP_PROP_FPS)

# 设置提取帧的间隔（例如每秒提取2帧）
frame_interval = int(fps / 2)

frame_count = 0
saved_frame_count = 0

while True:
    ret, frame = cap.read()
    if not ret:
        break

    if frame_count % frame_interval == 0:
        frame_filename = os.path.join(output_dir, f'frame_{saved_frame_count:05d}.jpg')
        cv2.imwrite(frame_filename, frame)
        saved_frame_count += 1

    frame_count += 1

cap.release()
print(f'Total frames saved: {saved_frame_count}')

1.2 标注数据

https://www.makesense.ai/

1.3 组织数据

将数据组织成以下目录结构：

dataset/
├── images/
│   ├── train/
│   ├── val/
│   └── test/
└── labels/
    ├── train/
    ├── val/
    └── test/

将数据集分为训练集（train）、验证集（val）和测试集（test）是训练机器学习模型的常规步骤。确保你的图片和标注文件命名一致。通常，数据集的划分比例如下：

• 训练集：70-80%

• 验证集：10-20%

• 测试集：10-20%

二、配置YOLOv8

配置文件

path: ./
train: images/train
val: images/val

names:
  0: long
  1: bashen

三、训练模型

from ultralytics import YOLO

def train_model():
    # 加载预训练的 YOLOv8 模型
    model = YOLO('yolov8n.pt')

    # 开始训练
    model.train(data='data.yaml', epochs=50, workers=1, batch=16, amp=False)

if __name__ == '__main__':
    train_model()

四、使用模型

import cv2
from ultralytics import YOLO

# 加载训练好的YOLOv8模型
model = YOLO('best.pt')

# 打开视频文件或摄像头
video_path = '1.mp4'  # 如果要使用摄像头，设置为 0
cap = cv2.VideoCapture(video_path)

if not cap.isOpened():
    print("Error: Could not open video.")
    exit()

# 获取视频帧的宽度、高度和FPS
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(cap.get(cv2.CAP_PROP_FPS))

# 定义输出视频文件（如果需要保存结果）
output_path = 'output.mp4'
out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))

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

    # 对每帧进行预测
    results = model(frame)

    # 提取预测结果并绘制边界框和标签
    annotated_frame = results[0].plot()  # 使用plot()方法绘制边界框

    # 显示结果
    cv2.imshow('YOLOv8 Real-time Detection', annotated_frame)

    # 保存结果到输出视频文件
    out.write(annotated_frame)

    # 按 'q' 键退出循环
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

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

最后：