Yolov8训练识别模型

本文手把手教你用YoloV8训练自己的数据集并实现物体识别

操作环境：

系统：Windows10

Python：3.9

Pytorch：2.2.2+cu121

环境安装

• 安装CUDA以及cudnn

 参考博客《Windows安装CUDA 12.1及cudnn》（https://www.cnblogs.com/RiverRiver/p/18103991）

• 安装torch, torchvision对应版本，这里先下载好，直接安装
  

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

•  pip安装yolo包

pip3 install ultralytics

• pip安装数据标注工具

pip install labelimg

数据准备

提前准备好需要训练的图片数据（尽量多一点），我这里以验证码的形状为例，如下图：

• 命令行输入 labelimg 打开数据标注工具，数据集类型切换成YOLO，然后依次完成标注即可

 点击Create RectBox开始标注，将需要识别的图形框起来，框起来后需要输入标签（注意：同一类型物体要用一个标签）。如下图：

•  标注划分

标注好之后，使用下面的脚本划分训练集、验证集，注意设置正确的图片和txt路径：

import os
import random
import shutil

# 设置文件路径和划分比例
root_path = "D:\\dataset"
# 标注过的图片存放目录
image_dir = "D:\\dataset\\images"
# 标注过生成的txt存放目录
label_dir = "D:\\dataset\\labels"
train_ratio = 0.7
val_ratio = 0.2
test_ratio = 0.1

# 创建训练集、验证集和测试集目录
os.makedirs("images/train", exist_ok=True)
os.makedirs("images/val", exist_ok=True)
os.makedirs("images/test", exist_ok=True)
os.makedirs("labels/train", exist_ok=True)
os.makedirs("labels/val", exist_ok=True)
os.makedirs("labels/test", exist_ok=True)

# 获取所有图像文件名
image_files = os.listdir(image_dir)
total_images = len(image_files)
random.shuffle(image_files)

# 计算划分数量
train_count = int(total_images * train_ratio)
val_count = int(total_images * val_ratio)
test_count = total_images - train_count - val_count

# 划分训练集
train_images = image_files[:train_count]
for image_file in train_images:
    label_file = image_file[:image_file.rfind(".")] + ".txt"
    shutil.copy(os.path.join(image_dir, image_file), "images/train/")
    shutil.copy(os.path.join(label_dir, label_file), "labels/train/")

# 划分验证集
val_images = image_files[train_count:train_count+val_count]
for image_file in val_images:
    label_file = image_file[:image_file.rfind(".")] + ".txt"
    shutil.copy(os.path.join(image_dir, image_file), "images/val/")
    shutil.copy(os.path.join(label_dir, label_file), "labels/val/")

# 划分测试集
test_images = image_files[train_count+val_count:]
for image_file in test_images:
    label_file = image_file[:image_file.rfind(".")] + ".txt"
    shutil.copy(os.path.join(image_dir, image_file), "images/test/")
    shutil.copy(os.path.join(label_dir, label_file), "labels/test/")

# 生成训练集图片路径txt文件
with open("train.txt", "w") as file:
    file.write("\n".join([root_path + "images/train/" + image_file for image_file in train_images]))

# 生成验证集图片路径txt文件
with open("val.txt", "w") as file:
    file.write("\n".join([root_path + "images/val/" + image_file for image_file in val_images]))

# 生成测试集图片路径txt文件
with open("test.txt", "w") as file:
    file.write("\n".join([root_path + "images/test/" + image_file for image_file in test_images]))

print("数据划分完成！")

运行后会生成划分好的数据集如下：

 训练与预测

• 开始训练

训练脚本如下：

from ultralytics import YOLO

# Load a model
model = YOLO('yolov8n.yaml')

results = model.train(data='shield.yml', epochs=1300, imgsz=640, device=[0],
                      workers=0, lr0=0.001, batch=128, amp=False)

shield.yml内容如下，注意修改自己的数据集路径即可：

# Ultralytics YOLO 

标签：val,模型,Yolov8,train,file,images,path,识别,image	

From： https://www.cnblogs.com/RiverRiver/p/18103892