点击下方卡片，关注“机器视觉与AI深度学习”
视觉/图像重磅干货，第一时间送达!
导 读
本文将手把手教你用YOLO11训练自己的数据集并实现缺陷检测。
安装环境
YOLO11的介绍和使用这里不再赘述，请参考公众号下面两篇文章即可：
YOLO11介绍及五大任务推理演示(目标检测，图像分割，图像分类，姿态检测，带方向目标检测)
【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，大家可以自行选择不同模型测试。
from ultralytics import YOLO# Load a modelmodel = YOLO("yolo11s.pt")results = model('c1.jpg',save=True)results[0].show()
标注/制作数据集
【1】下载缺陷检测数据集
【2】使用labelImg标注样本
pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simple【3】标注划分
# -*- coding: utf-8 -*-import osimport randomimport shutil# 设置文件路径和划分比例root_path = "./dataset/"image_dir = "./temp/images/"label_dir = "./temp/labels/"train_ratio = 0.7val_ratio = 0.2test_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 modelmodel = YOLO('yolo11s.pt')results = model.train(data='defects.yaml', epochs=100, imgsz=640, device=[0],workers=0,lr0=0.001,batch=8,amp=False)
path: E:/Practice/Python/PyTorch/dataset/ # dataset root dirtrain: E:/Practice/Python/PyTorch/dataset/images/trainval: E:/Practice/Python/PyTorch/dataset/images/valtest: # test images (optional)# Classesnames:0: defects
【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()
