基于YOLOv8的缺陷检测任务模型训练

文章目录

  • 一、引言
  • 二、环境说明
  • 三、缺陷检测任务模型训练详解
    • [3.1 PCB数据集](#3.1 PCB数据集)
      • [3.1.1 数据集简介](#3.1.1 数据集简介)
      • [3.1.2 数据集下载](#3.1.2 数据集下载)
      • [3.1.3 构建yolo格式的数据集](#3.1.3 构建yolo格式的数据集)
    • [3.2 基于ultralytics训练YOLOv8](#3.2 基于ultralytics训练YOLOv8)
      • [3.2.1 安装依赖包](#3.2.1 安装依赖包)
      • [3.2.2 ultralytics的训练规范说明](#3.2.2 ultralytics的训练规范说明)
      • [3.2.3 创建训练配置文件](#3.2.3 创建训练配置文件)
      • [3.2.4 下载预训练模型](#3.2.4 下载预训练模型)
      • [3.2.5 训练模型](#3.2.5 训练模型)
      • [3.2.6 训练结果](#3.2.6 训练结果)
  • 四、基于YOLO8的PCB缺陷检测系统

一、引言

  • 在当今快速发展的工业生产领域,自动化和智能化已成为提高生产效率和产品质量的关键因素。缺陷检测作为智能制造的重要组成部分,对于确保产品质量和降低生产成本具有重要意义。
  • 随着计算机视觉技术的不断进步,基于深度学习的对象检测算法,尤其是YOLO(You Only Look Once)系列算法,因其速度快、精度高而在工业缺陷检测中得到了广泛应用。YOLOv8作为该系列的最新版本之一,以其更高效的网络结构和更准确的检测性能,为缺陷检测任务带来了新的突破。
  • 本文将深入探讨基于YOLOv8的缺陷检测任务训练过程,从数据准备、模型训练到性能评估,全面解析如何利用这一先进技术提升工业缺陷检测的自动化水平。通过实际案例分析,本文旨在为相关领域的研究者和工程师提供实用的技术参考和指导。

注意:学习本文的训练过程需要相关的前置知识,如有需要可参考文章:Python基础(千锋篇)专栏介绍Jupyter NotebookPyTorch

二、环境说明

  1. 本文介绍的训练过程,使用ModelScope 提供的Notebook的线上环境。分为CPU环境/GPU环境,均可使用。

ModelScope--我的Notebook

  1. 本文的YOLOv8的训练过程,使用了ultralytics 的训练技术方案。相关文档可参考:https://docs.ultralytics.com/modes/train/

三、缺陷检测任务模型训练详解

完整代码资源:https://download.csdn.net/download/weixin_44063529/89645813

3.1 PCB数据集

3.1.1 数据集简介

PCB缺陷检测数据集:印刷电路板(PCB)瑕疵数据集由北京大学发布,其中包含1386张图像以及6种缺陷(缺失孔,鼠标咬伤,开路,短路,杂散,伪铜),用于检测,分类和配准任务。我们选取了其中适用与检测任务的693张图像,下载内容包括JSON和VOC格式版。

3.1.2 数据集下载

shell 复制代码
!git lfs install
!git clone https://www.modelscope.cn/datasets/ModelBulider/PCB_DATASET_JSON.git

!tar -xvf ./PCB_DATASET_JSON/PCB_DATASET_JSON.tar

3.1.3 构建yolo格式的数据集

  1. 创建训练集与测试集的标签文件
python 复制代码
import json
import os

def refactor_data_format(json_path, yolo_paths):
    with open(json_path) as f:
        data = json.load(f)

    imgs = {}
    for img in data['images']:
        imgs[img['id']] = {
            'h': img['height'],
            'w': img['width'],
            'file_name': img['file_name'],
        }


    tmp = ''
    for anno in data['annotations']:
        print(imgs[anno['image_id']]['file_name'])
        txt_path = os.path.join(yolo_paths, imgs[anno['image_id']]['file_name'].split('.')[0] + '.txt')
        with open(txt_path, 'w') as txt_file:
            if imgs[anno['image_id']] != tmp:
                # xywh --> xywh(归一化)
                bbox = [anno['bbox'][0] / imgs[anno['image_id']]['w'],
                        anno['bbox'][1] / imgs[anno['image_id']]['h'],
                        anno['bbox'][2] / imgs[anno['image_id']]['w'],
                            anno['bbox'][3] / imgs[anno['image_id']]['h']]

                cls_id = anno['category_id'] - 1
                # 保存
                txt_file.write(str(cls_id) + ' ' +" ".join([str(a) for a in bbox])+"\n")  # 生成格式0 cx,cy,w,h
                tmp = imgs[anno['image_id']]

            else:
                # xywh --> xywh(归一化)
                bbox = [anno['bbox'][0] / imgs[anno['image_id']]['w'],
                        anno['bbox'][1] / imgs[anno['image_id']]['h'],
                        anno['bbox'][2] / imgs[anno['image_id']]['w'],
                            anno['bbox'][3] / imgs[anno['image_id']]['h']]

                cls_id = anno['category_id'] - 1
                # 保存
                txt_file.write(str(cls_id) + ' ' +" ".join([str(a) for a in bbox])+"\n")  # 生成格式0 cx,cy,w,h    

# 生成训练与测试的label文件(.txt)
train_json_path = "./PCB_DATASET/Annotations/train.json"
train_yolo_paths = "./datasets/PCB_DATASET/labels/train2024"

val_json_path = "./PCB_DATASET/Annotations/val.json"
val_yolo_paths = "./datasets/PCB_DATASET/labels/val2024"

if not os.path.exists(train_yolo_paths):
    os.makedirs(train_yolo_paths)

if not os.path.exists(val_yolo_paths):
    os.makedirs(val_yolo_paths)

refactor_data_format(train_json_path, train_yolo_paths)
refactor_data_format(val_json_path, val_yolo_paths)
  1. 创建训练集与测试集的数据集
python 复制代码
import os
import shutil

Images_path = r'./PCB_DATASET/images'  # 源图路径

train_labels = r'./datasets/PCB_DATASET/labels/train2024'  # train标签路径
val_labels = r'./datasets/PCB_DATASET/labels/val2024'  # val标签路径

train_images = r'./datasets/PCB_DATASET/images/train2024'  # 保存train图像路径
val_images = r'./datasets/PCB_DATASET/images/val2024'  # 保存val图像路径

if not os.path.exists(train_images):
    os.makedirs(train_images)

if not os.path.exists(val_images):
    os.makedirs(val_images)

# 判断文件夹是否存在,不存在即创建
if not os.path.exists(train_images):
    os.mkdir(train_images)
if not os.path.exists(val_images):
    os.mkdir(val_images)


# 按照标签名移动对应图像
for label_name in os.listdir(train_labels):
    img_name = label_name[:-3] + 'jpg'
    shutil.move(os.path.join(Images_path, img_name), os.path.join(train_images, img_name))


for label_name in os.listdir(val_labels):
    img_name = label_name[:-3] + 'jpg'
    shutil.move(os.path.join(Images_path, img_name), os.path.join(val_images, img_name))

3.2 基于ultralytics训练YOLOv8

3.2.1 安装依赖包

shell 复制代码
!pip install ultralytics

3.2.2 ultralytics的训练规范说明

  1. YOLO训练标签示例如下(转换目标):
text 复制代码
类别id 归一化后框中心点x坐标 归一化后框中心点y坐标 归一化后框的宽度 归一化后框的高度
  1. 训练YOLOv8时,需要存放数据集的目录结构如下(图像与标签路径必须对应):
text 复制代码
├─datasets
│  ├─images
│  │  ├─train2024   # jpg/png 训练集图片
│  │  └─val2024     # jpg/png 验证集图片
│  ├─labels
│  │  ├─train2024   # txt 训练集标签
│  │  └─val2024     # txt 验证集标签

3.2.3 创建训练配置文件

python 复制代码
content = """path: ./PCB_DATASET # 数据集的根目录
train: images/train2024 # train images
val: images/val2024 # val images

# 分类类别
names:
  0: missing_hole
  1: mouse_bite
  2: open_circuit
  3: short
  4: spur
  5: spurious_copper
"""

with open(r"dataset.yaml", mode='w', encoding='UTF-8') as f:
    f.write(content)

3.2.4 下载预训练模型

python 复制代码
# 下载预训练的yolo8模型
import os
from modelscope.hub.api import HubApi
api = HubApi()
api.login('Your Modelscope SDK 令牌') # 通过访问:https://www.modelscope.cn/my/myaccesstoken 进行获取,需要提交下载申请

# 不同大小YOLO8模型下载
from modelscope import snapshot_download
model_dir = snapshot_download('ModelBulider/yolo8', local_dir="./modelscope_downloads")


yolo_version = 'yolov8l' # 选择训练使用yolo8的版本,可以在modelscope_downloads中查看所有的YOLO8版本
yolov8_model_path = os.path.abspath(os.path.join(model_dir, f'{yolo_version}.pt'))
print(yolov8_model_path)
  • YOLO8包含了各种大小版本的模型,各模型大小、性能等参数比较如下:
  • 如需要下载,可CSDN私信联系博主,佛系回复敬请谅解~
  • 可在yolo_version 处选择特定版本的YOLOv8预训练模型进行加载

3.2.5 训练模型

python 复制代码
import os
from ultralytics import YOLO

def download_train_file():
    """
    下载训练所需的文件
    """
    os.system("cp ./modelscope_downloads/Arial.ttf /root/.config/Ultralytics/")

if __name__ == '__main__':
    download_train_file()
    # Load a model
    # model = YOLO("yolov8n.yaml")  # build a new model from scratch
    model = YOLO(f'{yolo_version}.yaml').load(yolov8_model_path)  # load a pretrained model (recommended for training)

    # Use the model
    model.train(data="dataset.yaml", imgsz=640, batch=16, workers=8, cache=True, epochs=100)  # train the model
    metrics = model.val()  # evaluate model performance on the validation set
    # results = model("ultralytics\\assets\\bus.jpg")  # predict on an image
    path = model.export(format="onnx", opset=13)  # export the model to ONNX format
  • 训练日志:

3.2.6 训练结果

  • 训练产生的数据会生成在runs/detect/train路径下。其中,weight种存储了,训练过程中表现最好的模型参数(best.pt)以及最后一次训练得到的模型参数(last.pt)
  • 训练结果验证
python 复制代码
from ultralytics import YOLO

# Load a model
model = YOLO(r'runs/detect/train/weights/best.pt')  # pretrained YOLOv8n model

# Run batched inference on a list of images
results = model([r"datasets/PCB_DATASET/images/train2024/01_missing_hole_01.jpg", 
                 r"datasets/PCB_DATASET/images/train2024/01_missing_hole_17.jpg"])  # return a list of Results objects

# Process results list
for idex, result in enumerate(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
    obb = result.obb  # Oriented boxes object for OBB outputs
    result.show()  # display to screen
    result.save(filename=f"result_{idex}.jpg")  # save to disk


四、基于YOLO8的PCB缺陷检测系统

  • 使用训练得到模型,开发缺陷检测系统

在线体验:https://www.modelscope.cn/studios/ModelBulider/yolo_pcb_dec

相关推荐
人机与认知实验室14 分钟前
宽度学习与深度学习
人工智能·深度学习·学习
新智元16 分钟前
AI 永生时代来临!DeepMind「生成幽灵」让逝者赛博重生
人工智能·openai
HyperAI超神经19 分钟前
【vLLM 学习】Aqlm 示例
java·开发语言·数据库·人工智能·学习·教程·vllm
cnbestec23 分钟前
欣佰特携数十款机器人相关前沿产品,亮相第二届人形机器人和具身智能行业盛会
人工智能·机器人
爱的叹息23 分钟前
关于 梯度下降算法、线性回归模型、梯度下降训练线性回归、线性回归的其他训练算法 以及 回归模型分类 的详细说明
人工智能·算法·回归·线性回归
EasyGBS24 分钟前
室外摄像头异常自检指南+视频监控系统EasyCVR视频质量诊断黑科技
大数据·人工智能·音视频
Conan х31 分钟前
第1 篇:你好,时间序列!—— 开启时间数据探索之旅
人工智能·python·神经网络·机器学习·信息可视化
悟能不能悟32 分钟前
Coze平台 创建AI智能体的详细步骤指南
人工智能
CoderOnly44 分钟前
【图问答】DeepSeek-VL 论文阅读笔记
计算机视觉·大模型
阿星AI工作室1 小时前
飞书文档秒变高颜值网站!扣子空间MCP杀疯了,小白3步生成商务风主页!
人工智能