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)