概述
上节使用Labelimg标注斯诺克图片,分成以下几部分:
-
选择一段斯诺克比赛视频
-
将视频按照帧间隔截取图片,保存本地
-
用Labelimg工具将图片中所有的球标注出来。
1、将视频内容截取使用python cv2库,其代码如下:
import cv2
v=cv2.VideoCapture("./data/1.mp4")
n=0
step=500
while True:
ret ,frame=v.read()
if not ret:
break
n+=1
if n%step==0:
cv2.imwrite("./data/img/"+str(n)+".jpg",frame)
每间隔500帧,将视频中的图片报错下来。 2、Labelimg工具标注
- 建立标签,斯诺克八种颜色球代表八个标签
- 按照标签标注球位置
标注完成,生成的数据如下图:
classes.txt代表标签即标签排列 其余文件是对于图片内标注标签即球在的图片中位置,其内容如下
red
bule
yellow
green
white
pink
black
brown
如:1500.txt
7 0.667188 0.170833 0.020313 0.033333
5 0.500000 0.593750 0.017188 0.031944
1 0.501563 0.402778 0.017188 0.027778
0 0.418750 0.582639 0.021875 0.031944
0 0.462891 0.627083 0.022656 0.043056
6 0.500391 0.777083 0.025781 0.034722
4 0.345313 0.679861 0.018750 0.040278
第一行7 0.667188 0.170833 0.020313 0.033333,7是序号,表示classes中标签排序7的位置,即black,剩下四位是标签在图片中的位置大小。
训练前数据准备
数据准备,是要按照yolov8要求,将数据放置到对应位置。准备两个文件夹
- images,用于存放图片,新建train、val文件夹
- train训练集图片
- val验证集图片
- labels存放标签,新建train、val文件夹
- train训练集标签文件
- val验证集标签文件
整理后的数据,如下图:
数据集描述文件
数据集描述文件,是yolov8指定训练数据集位置以及标签内容。在yolov8根目录新建yolo-snooker.yaml,其内容如下
# 数据集描述文件
path: snooker
train: images/train
val: images/val
# Classes
names:
0: red
1: bule
2: yellow
3: green
4: white
5: pink
6: black
7: brown
其中path指定路径。train训练集路径,val验证集路径。
开始训练
新建snooker.py
- 导入yolov8模型
- 创建训练模块
- 指定数据集表述文件 代码如下
from ultralytics import YOLO
model=YOLO("yolov8n.pt","detect")
result=model.train(data="./yolo-snooker.yaml",workers=1,epochs=50,batch=16)
print(result)
其中workers现场数,epochs训练50轮,batch一次加载训练16张图片
点击运行,就可以开始训练了,如下图
训练完成后,会在runs目录下,生成train训练结果,其中weights就是训练出来的模型,best.pt是所要结果。
验证best.pt模型
vscode 打开cmd,进入conda虚拟环境,输入yolo预测命令
- conda activate yolov8
- yolo detect predict model=runs/detect/train/weights/best.pt source=./datasets/1.mp4 show=True