OpenCV与AI深度学习 | 实战 | YOLO11自定义数据集训练实现缺陷检测 (标注+训练+预测 保姆级教程)

本文来源公众号“OpenCV与AI深度学习”，仅用于学术分享，侵权删，干货满满。

原文链接：实战 | YOLO11自定义数据集训练实现缺陷检测 (标注+训练+预测 保姆级教程)

导  读

    本文将手把手教你用YOLO11训练自己的数据集并实现缺陷检测。

安装环境

YOLO11的介绍和使用这里不再赘述，请参考下面两篇文章即可：

OpenCV与AI深度学习 | YOLOv11来了：将重新定义AI的可能性_opencv ai-CSDN博客

OpenCV与AI深度学习 | YOLO11介绍及五大任务推理演示(目标检测，图像分割，图像分类，姿态检测，带方向目标检测)_yolo11n-pose.pt-CSDN博客

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

pip install torch-1.13.1+cu116-cp38-cp38-win_amd64.whlpip install torchvision-0.14.1+cu116-cp38-cp38-win_amd64.whl

安装好后可以查看是否安装成功，上面安装的gpu版本，查看指令与结果：

import torchprint(torch.__version__)print(torch.cuda

【2】安装ultralytics

pip install ultralytics

【3】下载YOLO11预训练模型：

https://github.com/ultralytics/ultralytics

    本文使用YOLO11s，大家可以自行选择不同模型测试。

【4】运行demo测试安装是否成功：

from ultralytics import YOLO

# Load a model
model = YOLO("yolo11s.pt")

results = model('c1.jpg',save=True)

results[0].show()

标注/制作数据集

【1】下载缺陷检测数据集

    本文使用DAGM 2007数据集其中的Dataset4类别的缺陷图：

    标注文件可以自己根据Label文件夹的标签图写脚本来生成，也可以使用labelImg自己标注。

【2】使用labelImg标注样本

    标注工具使用labelimg即可，直接pip安装：

pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simple

    安装完成后，命令行直接输入labelimg，回车即可打开labelimg，数据集类型切换成YOLO，然后依次完成标注即可。

【3】标注划分

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

# -*- coding: utf-8 -*-
import os
import random
import shutil

# 设置文件路径和划分比例
root_path = "./dataset/"
image_dir = "./temp/images/"
label_dir = "./temp/labels/"
train_ratio = 0.7
val_ratio = 0.2
test_ratio = 0.1

# 创建训练集、验证集和测试集目录
os.makedirs(root_path+"images/train", exist_ok=True)
os.makedirs(root_path+"images/val", exist_ok=True)
os.makedirs(root_path+"images/test", exist_ok=True)
os.makedirs(root_path+"labels/train", exist_ok=True)
os.makedirs(root_path+"labels/val", exist_ok=True)
os.makedirs(root_path+"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), root_path+"images/train/")
    shutil.copy(os.path.join(label_dir, label_file), root_path+"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), root_path+"images/val/")
    shutil.copy(os.path.join(label_dir, label_file), root_path+"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), root_path+"images/test/")
    shutil.copy(os.path.join(label_dir, label_file), root_path+"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("数据划分完成！")

    接着会生成划分好的数据集如下：

训练与预测

【1】开始训练

    训练脚本如下：

from ultralytics import YOLO

# Load a model
model = YOLO('yolo11s.pt')  

results = model.train(data='defects.yaml', epochs=100, imgsz=640, device=[0],
                      workers=0,lr0=0.001,batch=8,amp=False)

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

path: E:/Practice/Python/PyTorch/dataset/ # dataset root dir
train: E:/Practice/Python/PyTorch/dataset/images/train
val: E:/Practice/Python/PyTorch/dataset/images/val
test: # test images (optional)

# Classes
names:
  0: defects

    CPU训练将device=[0]改为device='cpu'即可

    训练完成后再runs/detect/train文件夹下生成如下内容：

    在weights文件夹下生成两个模型文件，直接使用best.pt即可。

【2】预测推理

    预测脚本如下：

from ultralytics import YOLO

# Load a model
#model = YOLO("yolo11s.pt")
model = YOLO("best.pt")

results = model('./test/0050.PNG',conf=0.5,save=True)

results[0].show()

    使用数据集Test文件夹中的异常和正常图片测试：

THE END !

文章结束，感谢阅读。您的点赞，收藏，评论是我继续更新的动力。大家有推荐的公众号可以评论区留言，共同学习，一起进步。