使用yolov8训练焊缝缺陷检测数据集利用训练出的权重识别正常焊缝和错位、烧穿、污染等五类缺陷焊缝并做成深度学习焊缝缺陷识别系统

计算机C9硕士_算法工程师

于 2025-02-09 06:50:59 发布

阅读量1.1k

点赞数 14

分类专栏：焊接焊缝文章标签： YOLO 深度学习人工智能

本文链接：https://blog.youkuaiyun.com/2401_88440984/article/details/145527045

版权

焊接焊缝专栏收录该内容

4 篇文章

订阅专栏

使用yolov8训练焊缝缺陷检测数据集利用训练出的权重识别正常焊缝和错位、烧穿、污染等五类缺陷焊缝并做成深度学习焊缝缺陷识别系统

焊缝缺陷检测数据集在这里插入图片描述

：30000+
类型：包括良好焊缝和错位、烧穿、污染等五类缺陷焊缝，共6种焊缝类型
材质：铝合金在这里插入图片描述
1
使用YOLOv8训练一个焊缝缺陷检测系统，并将训练好的模型部署为一个实际可用的系统，你需要遵循一系列步骤。以下是一个详细的指南，涵盖了从数据集准备到模型训练再到最终系统的实现。

文字及代码仅供参考。

准备数据集、配置训练参数并执行模型训练。以下是详细的代码示例和解释，帮助你完成这个过程

为了训练一个基于YOLOv8的焊缝缺陷检测系统，你需要遵循一系列步骤来准备数据集、配置训练参数并执行模型训练。以下是详细的代码示例和解释，帮助你完成这个过程。

1. 数据集准备

假设你的数据集已经按照以下结构组织：

dataset/train/images/: 训练图像
dataset/train/labels/: 对应的标签文件（YOLO格式）
dataset/val/images/: 验证图像
dataset/val/labels/: 对应的标签文件

每个标签文件中的每一行代表一个对象，格式如下：<class_index> <x_center> <y_center> <width> <height>，其中坐标和尺寸已经归一化。

创建一个data.yaml文件来描述数据集路径和类别信息：

train: ./dataset/train/images/
val: ./dataset/val/images/

nc: 6  # 类别数量
names: ['good_weld', 'misalignment', 'burn_through', 'contamination', 'other_defect_1', 'other_defect_2']  # 类别名称

2. 安装必要的依赖

确保安装了必要的库：

pip install ultralytics opencv-python

3. 模型训练

使用YOLOv8进行模型训练。下面是一个完整的Python脚本，用于加载预训练的YOLOv8模型并开始训练过程。

训练脚本

from ultralytics import YOLO
import os

def train_yolov8_model():
    # 加载预训练的YOLOv8模型
    model = YOLO('yolov8n.pt')  # 或者选择其他变体如'yolov8s.pt', 'yolov8m.pt'等
    
    # 开始训练
    results = model.train(
        data='path/to/data.yaml',  # 数据集配置文件路径
        epochs=100,  # 根据需要调整epoch数
        imgsz=640,  # 输入图像尺寸
        batch=16,  # 批次大小
        name='weld_defect_detection',  # 实验名称
        save=True,  # 自动保存最佳模型
        exist_ok=True,  # 如果目录存在则不报错
        patience=50,  # 当验证损失不再改善时提前停止训练的轮数
        lr0=0.01,  # 初始学习率
        lrf=0.1,  # 最终学习率（lr0 * lrf）
        optimizer='SGD',  # 优化器类型
        augment=True,  # 启用数据增强
        flipud=0.0,  # 上下翻转概率
        fliplr=0.5,  # 左右翻转概率
        mosaic=1.0,  # Mosaic数据增强的概率
        mixup=0.0,  # MixUp数据增强的概率
        degrees=0.0,  # 旋转角度范围
        translate=0.1,  # 平移比例范围
        scale=0.5,  # 缩放比例范围
        shear=0.0,  # 剪切变换角度范围
        perspective=0.0,  # 透视变换比例
        dropout=0.0,  # Dropout比例
        hsv_h=0.015,  # HSV色调变化比例
        hsv_s=0.7,  # HSV饱和度变化比例
        hsv_v=0.4,  # HSV亮度变化比例
        copy_paste=0.0,  # Copy-Paste数据增强的概率
        cache=True,  # 使用缓存加速训练
        device='',  # 使用GPU或CPU ('cuda' or 'cpu')
        workers=8,  # 数据加载的工作线程数
        project='runs/detect',  # 保存结果的项目目录
        entity=None,  # WandB实体名
        upload_dataset=False,  # 是否上传数据集到WandB
        bbox_interval=-1,  # Bbox日志间隔
        artifact_alias="latest",  # 版本控制别名
        multi_scale=False,  # 多尺度训练
        sync_bn=False,  # 同步BN
        cos_lr=False,  # 使用余弦退火学习率调度
        single_cls=False,  # 单类模式
        rect=False,  # 矩形训练
        resume=False,  # 恢复训练
        freeze=[0],  # 冻结层的数量
        verbose=True,  # 显示详细输出
        seed=0,  # 随机种子
        linear_lr=False,  # 线性学习率
        overlap_mask=True,  # 掩码重叠
        mask_ratio=4,  # 掩码比例
        v5_metric=False  # 使用YOLOv5的评估指标
    )

if __name__ == "__main__":
    train_yolov8_model()

参数说明

epochs: 总共训练多少轮，默认100轮。
imgsz: 输入图像的尺寸，默认640。
batch: 每批次处理的图像数量，默认16。
name: 实验名称，用于保存结果。
save: 是否自动保存最佳模型。
exist_ok: 如果实验目录已存在是否覆盖。
patience: 当验证损失不再改善时提前停止训练的轮数。
lr0: 初始学习率。
lrf: 最终学习率（lr0 * lrf）。
optimizer: 优化器类型，支持’SGD’, ‘Adam’, 'AdamW’等。
augment: 是否启用数据增强。
flipud, fliplr: 上下翻转和左右翻转的概率。
mosaic, mixup, degrees, translate, scale, shear, perspective, dropout, hsv_h, hsv_s, hsv_v, copy_paste: 不同类型的数据增强参数。
cache: 是否使用缓存加速训练。
device: 使用GPU或CPU (‘cuda’ or ‘cpu’)。
workers: 数据加载的工作线程数。

4. 验证模型性能

训练完成后，可以使用验证集评估模型性能：

from ultralytics import YOLO

def validate_model():
    model = YOLO('runs/detect/weld_defect_detection/weights/best.pt')  # 加载最佳模型
    metrics = model.val()  # 进行验证
    print(metrics.box.map)  # 输出mAP值作为性能指标之一

if __name__ == "__main__":
    validate_model()

5. 构建焊缝缺陷检测系统

使用PyQt5构建用户界面，支持视频、图片文件选择及摄像头实时检测：

import sys
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QPushButton, QLabel, QFileDialog
from PyQt5.QtGui import QPixmap, QImage
import cv2
from ultralytics import YOLO

class WeldDefectDetectionApp(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()
        self.model = YOLO('runs/detect/weld_defect_detection/weights/best.pt')  # 加载训练好的模型

    def initUI(self):
        self.setWindowTitle('Weld Defect Detection System')
        layout = QVBoxLayout()

        self.imageLabel = QLabel(self)
        layout.addWidget(self.imageLabel)

        btn_load_image = QPushButton('Load Image', self)
        btn_load_image.clicked.connect(self.loadImage)
        layout.addWidget(btn_load_image)

        btn_start_camera = QPushButton('Start Camera', self)
        btn_start_camera.clicked.connect(self.startCamera)
        layout.addWidget(btn_start_camera)

        self.setLayout(layout)

    def loadImage(self):
        fname, _ = QFileDialog.getOpenFileName(self, 'Open file', '', "Image files (*.jpg *.png)")
        if fname:
            pixmap = QPixmap(fname)
            self.imageLabel.setPixmap(pixmap.scaled(self.imageLabel.size(), aspectRatioMode=1))
            self.detectDefects(fname)  # 调用检测函数

    def detectDefects(self, image_path):
        results = self.model.predict(image_path, conf=0.5)
        for result in results:
            for box in result.boxes:
                cls, conf = box.cls.item(), box.conf.item()
                label = self.model.model.names[int(cls)]
                print(f"Detected {label} with confidence {conf:.2f}")

    def startCamera(self):
        cap = cv2.VideoCapture(0)
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            results = self.model.predict(frame, conf=0.5)
            for result in results:
                for box in result.boxes:
                    cls, conf = box.cls.item(), box.conf.item()
                    label = self.model.model.names[int(cls)]
                    x1, y1, x2, y2 = map(int, box.xyxy)
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                    cv2.putText(frame, f'{label} {conf:.2f}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            h, w, ch = rgb_frame.shape
            bytes_per_line = ch * w
            convert_to_Qt_format = QImage(rgb_frame.data, w, h, bytes_per_line, QImage.Format_RGB888)
            p = convert_to_Qt_format.scaled(640, 480, aspectRatioMode=1)
            self.imageLabel.setPixmap(QPixmap.fromImage(p))

            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        cap.release()
        cv2.destroyAllWindows()

app = QApplication(sys.argv)
ex = WeldDefectDetectionApp()
ex.show()
sys.exit(app.exec_())

数据集准备

假设你的数据集结构如下：

dataset/train/images/: 训练图像
dataset/train/labels/: 对应的标签文件（YOLO格式）
dataset/val/images/: 验证图像
dataset/val/labels/: 对应的标签文件

每张图像都应该有对应的标签文件，标签文件中的每一行代表一个对象，格式如下：<class_index> <x_center> <y_center> <width> <height>，其中坐标和尺寸已经归一化。

创建一个data.yaml文件来描述数据集路径和类别信息：

train: ./dataset/train/images/
val: ./dataset/val/images/

nc: 6  # 类别数量
names: ['good_weld', 'misalignment', 'burn_through', 'contamination', 'other_defect_1', 'other_defect_2']  # 类别名称

模型训练

安装依赖

确保安装了必要的库：

pip install ultralytics opencv-python

训练脚本

使用YOLOv8进行训练：

from ultralytics import YOLO

def train_yolov8_model():
    # 加载预训练的YOLOv8模型
    model = YOLO('yolov8n.pt')  # 或者选择其他变体如'yolov8s.pt', 'yolov8m.pt'等
    
    # 开始训练
    results = model.train(
        data='path/to/data.yaml',  # 数据集配置文件路径
        epochs=100,  # 根据需要调整epoch数
        imgsz=640,  # 输入图像尺寸
        batch=16,  # 批次大小
        name='weld_defect_detection',  # 实验名称
        save=True,  # 自动保存最佳模型
        exist_ok=True  # 如果目录存在则不报错
    )

if __name__ == "__main__":
    train_yolov8_model()

构建焊缝缺陷检测系统

创建用户界面

使用PyQt5构建一个简易的UI，支持视频、图片文件选择及摄像头实时检测：

import sys
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QPushButton, QLabel, QFileDialog
from PyQt5.QtGui import QPixmap, QImage
import cv2
from ultralytics import YOLO

class WeldDefectDetectionApp(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()
        self.model = YOLO('runs/detect/weld_defect_detection/weights/best.pt')  # 加载训练好的模型

    def initUI(self):
        self.setWindowTitle('Weld Defect Detection System')
        layout = QVBoxLayout()

        self.imageLabel = QLabel(self)
        layout.addWidget(self.imageLabel)

        btn_load_image = QPushButton('Load Image', self)
        btn_load_image.clicked.connect(self.loadImage)
        layout.addWidget(btn_load_image)

        btn_start_camera = QPushButton('Start Camera', self)
        btn_start_camera.clicked.connect(self.startCamera)
        layout.addWidget(btn_start_camera)

        self.setLayout(layout)

    def loadImage(self):
        fname, _ = QFileDialog.getOpenFileName(self, 'Open file', '', "Image files (*.jpg *.png)")
        if fname:
            pixmap = QPixmap(fname)
            self.imageLabel.setPixmap(pixmap.scaled(self.imageLabel.size(), aspectRatioMode=1))
            self.detectDefects(fname)  # 调用检测函数

    def detectDefects(self, image_path):
        results = self.model.predict(image_path, conf=0.5)
        for result in results:
            for box in result.boxes:
                cls, conf = box.cls.item(), box.conf.item()
                label = self.model.model.names[int(cls)]
                print(f"Detected {label} with confidence {conf:.2f}")

    def startCamera(self):
        cap = cv2.VideoCapture(0)
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            results = self.model.predict(frame, conf=0.5)
            for result in results:
                for box in result.boxes:
                    cls, conf = box.cls.item(), box.conf.item()
                    label = self.model.model.names[int(cls)]
                    x1, y1, x2, y2 = map(int, box.xyxy)
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                    cv2.putText(frame, f'{label} {conf:.2f}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            h, w, ch = rgb_frame.shape
            bytes_per_line = ch * w
            convert_to_Qt_format = QImage(rgb_frame.data, w, h, bytes_per_line, QImage.Format_RGB888)
            p = convert_to_Qt_format.scaled(640, 480, aspectRatioMode=1)
            self.imageLabel.setPixmap(QPixmap.fromImage(p))

            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        cap.release()
        cv2.destroyAllWindows()

app = QApplication(sys.argv)
ex = WeldDefectDetectionApp()
ex.show()
sys.exit(app.exec_())