基于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

相关推荐
澜舟孟子开源社区12 分钟前
“AI玩手机”原理揭秘:大模型驱动的移动端GUI智能体
人工智能·科技·agi
Mr.鱼19 分钟前
opencv undefined reference to `cv::noarray()‘ 。window系统配置opencv,找到opencv库,但连接不了
人工智能·opencv·计算机视觉
ATpiu23 分钟前
免费微调自己的大模型(llama-factory微调llama3.1-8b)
人工智能·机器学习·llama
凌虚(失业了求个工作)26 分钟前
RAG 示例:使用 langchain、Redis、llama.cpp 构建一个 kubernetes 知识库问答
人工智能·redis·python·langchain·llama
逝去的紫枫1 小时前
Python PIL:探索图像处理的无限可能
图像处理·人工智能·python
sp_fyf_20242 小时前
计算机前沿技术-人工智能算法-大语言模型-最新研究进展-2024-11-05
人工智能·深度学习·神经网络·算法·机器学习·语言模型·数据挖掘
小火炉Q2 小时前
02 python基础 python解释器安装
人工智能·python·神经网络·机器学习·网络安全·自然语言处理
钰见梵星2 小时前
深度学习优化算法
人工智能·深度学习·算法
难念的码2 小时前
Skill 语言语法基础
人工智能·后端
是Winky啊2 小时前
【论文阅读】WGSR
论文阅读·深度学习·计算机视觉·超分辨率重建