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 🚀, AGPL-3.0 license
# COCO8 dataset (first 8 images from COCO train2017) by Ultralytics
# Documentation: https://docs.ultralytics.com/datasets/detect/coco8/
# Example usage: yolo train data=coco8.yaml
# parent
# ├── ultralytics
# └── datasets
# └── coco8 ← downloads here (1 MB)
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: E:\Code\Python\yolov8 # dataset root dir
train: E:\Code\Python\yolov8/images/train # train images (relative to 'path') 4 images
val: E:\Code\Python\yolov8/images/val # val images (relative to 'path') 4 images
test: # test images (optional)
# Classes 类别填写标记时的标签多个类型的话按照顺序0，1，2，3....添加
names:
0: shield
# Download script/URL (optional)
# download: https://ultralytics.com/assets/coco8.zip
训练完成后会再runs/detect/train文件夹下生成如下内容：
在weights文件夹下生成两个模型文件，直接使用best.pt即可。
预测推理
- 预测脚本如下
from ultralytics import YOLO # Load a model model = YOLO('E:\\Code\\Python\\yolov8\\runs\\detect\\train\\weights\\best.pt') # pretrained YOLOv8n model # Run batched inference on a list of images results = model(['.\\images\\test\\Screenshot_20230118_210923_com.tencent.mobileqq.jpg','.\\images\\test\\Screenshot_20230118_210936_com.tencent.mobileqq.jpg', './images/test/ax1.png']) # return a list of Results objects # Process results list for result in results: boxes = result.boxes # Boxes object for bounding box outputs masks = result.masks # Masks object for segmentation masks outputs keypoints = result.keypoints # Keypoints object for pose outputs probs = result.probs # Probs object for classification outputs result.show() # display to screen result.save(filename='result.jpg') # save to disk
预测结果：
