一、YOLO V8
YOLO V8
是由 2023
年 ultralytics
公司开源的发布,是结合了前几代 YOLO
的融合改进版。YOLO V8
支持全方位的视觉 AI
任务,包括检测、分割、姿态估计、跟踪和分类。并且在速度和准确性方面具有无与伦比的性能。能够应用在各种对速度和精度要求较高的应用和领域。
网络结构如下图所示:
YOLO V8
相对于 YOLO V5
还是有很大的不同,例如:YOLO V8
相对于 YOLO V5
,依然使用的是CSP
的思想,不过将 V5
中的C3
模块换成了C2F
模块,以减轻模型的大小,也依旧使用 V5
架构中的SPPF
模块。但是在 PAN-FPN
层面,V8
将 V5
中的上采样阶段中的卷积结构去除了。同时借鉴了 YOLOX
的 Decoupled-Head
结构,分类和回归两个任务的 HEAD
不再共享参数等。
在模型上 V8 和 V5 类似,包括不同大小的模型,从小到大包括:yolov8n、yolov8s、yolov8m、yolov8l、yolov8x
等:
模型的比较如下:
更多的介绍可以参考官方的文档:
https://docs.ultralytics.com/de/models/yolov8/
本文借助ultralytics
中 YOLO V8
迁移训练自定义的目标检测模型,在本次的实验中,主要训练一个人脸检测模型,包括数据标注、数据拆分、训练、测试等过程。
本次采用ultralytics
公司发布的 ultralytics
框架,可以帮助开发人员高效完成数据训练和验证任务,由于 ultralytics
默认采用的为 PyTorch
框架,因此实验前请安装好 cuda
和 torch
环境,如果没有 GPU
环境,由于YOLO V8
已经足够轻量级,使用CPU
也是可以训练。
安装 ultralytics
库:
pip install ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple
ultralytics
使用文档:
https://docs.ultralytics.com/zh/quickstart/#use-ultralytics-with-python
测试 YOLO V8
的效果:
测试图片:
这里使用 yolov8n
模型,如果模型不存在会自动下载
from ultralytics import YOLO
# Load a model
model = YOLO('yolov8n.pt') # pretrained YOLOv8n model
results = model.predict('./img/1.png')
# Show results
results[0].show()
二、数据收集及标注
图像数据可以从网上找一些或者自己拍摄,我这里准备了一些 人 的图片:
这里可以准备两个目录,data/images
和 data/labels
,其中 labels
存放标注后的文件,将收集到的图像放在 images
目录下:
下面使用 labelimg
工具进行标注,如果没有安装,使用下面命令安装:
pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simple
然后在控制台输入:labelimg
打开可视化工具:
注意:数据集格式默认是 VOC
格式的,要选择为 YOLO
,我这里的人脸标签为 face
,这个后面需要使用到。
标注完成后,可以在 /data/labels
下看到标注后的文件:
三、数据拆分
这里拆分为 90%
的训练集,10%
的验证集,拆分脚本如下,
import os
import shutil
from tqdm import tqdm
# 图片地址
image_dir = "data/images/"
# 标准文件地址
label_dir = "data/labels/"
# 训练集的比例
training_ratio = 0.9
# 拆分后数据的位置
train_dir = "train_data"
def split_data():
list = os.listdir(image_dir)
all = len(list)
train_count = int(all * training_ratio)
train_images = list[0:train_count]
val_images = list[train_count:]
# 训练集目录
os.makedirs(os.path.join(train_dir, "images/train"), exist_ok=True)
os.makedirs(os.path.join(train_dir, "labels/train"), exist_ok=True)
# 验证集目录
os.makedirs(os.path.join(train_dir, "images/val"), exist_ok=True)
os.makedirs(os.path.join(train_dir, "labels/val"), exist_ok=True)
# 训练集
with open(os.path.join(train_dir, "train.txt"), "w") as file:
file.write("\n".join([train_dir + "images/train/" + image_file for image_file in train_images]))
print("save train.txt success!")
# 拷贝数据
for item in tqdm(train_images):
label_file = item.replace(".jpg", ".txt")
shutil.copy(os.path.join(image_dir, item), os.path.join(train_dir, "images/train/"))
shutil.copy(os.path.join(label_dir, label_file), os.path.join(train_dir, "labels/train/"))
# 验证集
with open(os.path.join(train_dir, "val.txt"), "w") as file:
file.write("\n".join([train_dir + "images/val/" + image_file for image_file in val_images]))
print("save val.txt success!")
# 拷贝数据
for item in tqdm(val_images):
label_file = item.replace(".jpg", ".txt")
shutil.copy(os.path.join(image_dir, item), os.path.join(train_dir, "images/val/"))
shutil.copy(os.path.join(label_dir, label_file), os.path.join(train_dir, "labels/val/"))
if __name__ == '__main__':
split_data()
可以在 train_data
中看到拆分后的数据集格式:
四、训练
使用 ultralytics 框架训练非常简单,仅需三行代码即可完成训练,不过在训练前需要编写 YAML
配置信息,主要标记数据集的位置。
创建 face.yaml
文件,写入下面内容:
path: D:/pyProject/yolov8/train_data # 数据集的根目录, 建议使用绝对路径
train: images/train # 训练集图像目录
val: images/val # 验证集图像目录
test: # test images (optional)
# 分类
names:
0: face
注意分类中的 face
就是上面标注时的标签名。
开始训练:
from ultralytics import YOLO
# 加载模型
model = YOLO('yolov8n.pt')
# 训练
model.train(
data='face.yaml', # 训练配置文件
epochs=50, # 训练的周期
imgsz=640, # 图像的大小
device=[0], # 设备,如果是 cpu 则是 device='cpu'
workers=0,
lr0=0.001, # 学习率
batch=8, # 批次大小
amp=False # 是否启用混合精度训练
)
运行后可以看到打印的网络结构:
训练中:
训练结束后可以在 runs
目录下面看到训练的结果:
其中 weights
下面的就是训练后保存的模型,这里可以先看下训练时 loss
的变化图:
五、模型测试
使用 best.pt
模型
from ultralytics import YOLO
# 加载模型
model = YOLO('runs/detect/train/weights/best.pt')
# 预测
results = model.predict('data/images/8.jpg')
# Show results
results[0].show()