YOLOv8实现热轧带钢表面缺陷目标检测（附代码和数据集）

1.数据集

NEU-DET数据集收集了六种典型的热轧带钢表面缺陷，即：

• 轧制鳞片（Rolled-in_scale）
• 斑块（Patches）
• 裂纹（Crazing）
• 坑洼面（Pitted_surface）
• 夹杂（Inclusion）
• 划痕（Scratches）

数据库包括1800张灰度图像，6种不同的典型表面缺陷样本各300个。数据集中images文件存放的是jpg格式的缺陷图片，annatation中存放的是缺陷图片对应的标注为xml格式文件，包括缺陷的类型、位置和大小等信息。

使用YOLOv8进行缺陷检测，需要的数据集格式为txt。因此需要把xml格式转为txt格式，转换代码：(注意：需要提前在train文件夹和test文件夹中分别创建labels命名的文件夹)

import os
import xml.etree.ElementTree as ET


def convert_xml_to_yolo(xml_file, classes, output_folder):
    tree = ET.parse(xml_file)
    root = tree.getroot()

    image_filename = root.find('filename').text
    image_width = int(root.find('size').find('width').text)
    image_height = int(root.find('size').find('height').text)

    yolo_labels = []

    for obj in root.findall('object'):
        class_name = obj.find('name').text
        if class_name not in classes:
            continue

        class_id = classes.index(class_name)
        bbox = obj.find('bndbox')

        x_center = (int(bbox.find('xmin').text) + int(bbox.find('xmax').text)) / 2 / image_width
        y_center = (int(bbox.find('ymin').text) + int(bbox.find('ymax').text)) / 2 / image_height
        w = (int(bbox.find('xmax').text) - int(bbox.find('xmin').text)) / image_width
        h = (int(bbox.find('ymax').text) - int(bbox.find('ymin').text)) / image_height

        yolo_labels.append(f"{class_id} {x_center} {y_center} {w} {h}")

    output_path = os.path.join(output_folder, os.path.splitext(image_filename)[0] + ".txt")

    with open(output_path, 'w') as out_file:
        for line in yolo_labels:
            out_file.write(line + '\n')

# Define your classes
classes = ['crazing', 'inclusion', 'patches', 'pitted_surface', 'rolled-in_scale', 'scratches']  # Add your own class names

# Path to the directory containing XML files
train_xml_dir = 'D:/pycharm code/ultralytics-main/data/NEU-DET/train/annatation'
train_output_dir = 'D:/pycharm code/ultralytics-main/data/NEU-DET/train/labels'  # Output directory for YOLO format files

test_xml_dir = 'D:/pycharm code/ultralytics-main/data/NEU-DET/test/annatation'
test_output_dir = 'D:/pycharm code/ultralytics-main/data/NEU-DET/test/labels'  # Output directory for YOLO format files

# Create the output directory if it doesn't exist
os.makedirs(train_output_dir , exist_ok=True)
os.makedirs(test_output_dir, exist_ok=True)

# Convert all XML files to YOLO format and save to the output directory
for xml_file in os.listdir(train_xml_dir ):
    if xml_file.endswith('.xml'):
        xml_path = os.path.join(xml_dir, xml_file)
        convert_xml_to_yolo(xml_path, classes, train_output_dir )
# 测试集
for xml_file in os.listdir(test_xml_dir):
    if xml_file.endswith('.xml'):
        xml_path = os.path.join(test_xml_dir, xml_file)
        convert_xml_to_yolo(xml_path, classes, test_output_dir)

2.下载预训练模型

（1）下载yolov8n训练模型，下载地址如下：

https://github.com/ultralytics/ultralyticshttps://github.com/ultralytics/ultralytics

（n、s、m、l、x五种不同规模的预训练模型选择一个即可），放到ultralytics-main目录下即可。

（2）创建yaml模型文件

在自己创建的data文件夹中，创建data.yaml文件，存放自己的训练集和测试集地址，检测项目类别数，检测项目类别：（对于本数据集共有6中类别，因此nc：6。）（种类名称放到names：）

（3）找到YOLO v8\ultralytics-main\ultralytics/yolo/cfg文件路径下的default.yaml文件复制到自己数据集目录下并打开进行如下修改：

• task设置为detect；
• mode设置为train；
• model后设置为预训练模型yolov8l.pt的文件路径；（我下载的yolov8l.pt文件）
• data后设置为（2）中所述的data.yaml文件的绝对路径；
• epoch设置为自己所需的训练轮数；
• batch根据自己电脑配置进行设置。

（4）开始训练，在终端中输入以下命令回车，即可开始训练：

yolo cfg=dataSets/NEU-DET/default.yaml

训练过程如下：

(5) 进行predict预测

1. 将default.yaml文件中mode改为predict。
2. 将model由预训练模型更改为自己训练后的最优模型best.pt的绝对路径。
3. 将source改为待预测图片所在文件夹的路径。

4. 在终端中输入以下命令回车，即可开始预测：

   yolo cfg=dataSets/NEU-DET/default.yaml

运行预测完毕后可以看到每张图片的预测结果和最终的预测结果保存路径如下，也可以直接预测视频：

预测推理结果如下： 

有问题咨询qq：498609554 

微信：