【YOLO11实战】NEU-DET，钢材缺陷识别实战

1.下载数据

http://faculty.neu.edu.cn/songkechen/zh_CN/zdylm/263270/list/

配置环境
使用的python版本是3.8.10

certifi==2025.1.31
charset-normalizer==3.4.1
colorama==0.4.6
coloredlogs==15.0.1
contourpy==1.1.1
cycler==0.12.1
filelock==3.16.1
flatbuffers==25.2.10
fonttools==4.57.0
fsspec==2025.3.0
humanfriendly==10.0
idna==3.10
importlib_resources==6.4.5
Jinja2==3.1.6
kiwisolver==1.4.7
MarkupSafe==2.1.5
matplotlib==3.7.5
mpmath==1.3.0
networkx==3.1
numpy==1.24.4
onnx==1.17.0
onnxruntime-gpu==1.19.2
onnxslim==0.1.50
opencv-python==4.11.0.86
packaging==24.2
pandas==2.0.3
pillow==10.4.0
protobuf==5.29.4
psutil==7.0.0
py-cpuinfo==9.0.0
pyparsing==3.1.4
pyreadline3==3.5.4
python-dateutil==2.9.0.post0
pytz==2025.2
PyYAML==6.0.2
requests==2.32.3
scipy==1.10.1
seaborn==0.13.2
six==1.17.0
sympy==1.13.3
torch==2.4.1+cu118
torchaudio==2.4.1+cu118
torchvision==0.19.1+cu118
tqdm==4.67.1
typing_extensions==4.13.2
tzdata==2025.2
ultralytics==8.3.107
ultralytics-thop==2.0.14
urllib3==2.2.3
zipp==3.20.2

2. 转化数据为yolo格式

手动把Annotation和images划分成train数据集和valid数据集，

valid数据集可以每类放一点，一共6种缺陷类型：

• crazing
• inclusion
• patches
• pitted_surface
• rolled-in_scale
• scratches

我的valid数据集是 随机选的一些数据

data.yaml:

train: ./NEU/NEU-DET/train/images
train_xml: ./NEU/NEU-DET/train/labels
val: ./NEU/NEU-DET/valid/images
val_xml: ./NEU/NEU-DET/valid/labels

nc: 6
classes:
  - crazing
  - inclusion
  - patches
  - pitted_surface
  - rolled-in_scale
  - scratches

xml转yolo

import os
import xml.etree.ElementTree as ET
import os
import yaml

def load_config(config_path):
    """从YAML文件加载配置"""
    try:
        with open(config_path, 'r') as f:
            config = yaml.safe_load(f)
            return config
    except Exception as e:
        print(f"加载配置文件失败: {str(e)}")
        exit(1)

def xml_to_yolo(xml_path, output_dir, classes):
    tree = ET.parse(xml_path)
    root = tree.getroot()

    # 获取图片尺寸
    size = root.find('size')
    img_width = int(size.find('width').text)
    img_height = int(size.find('height').text)

    # 创建输出文件
    txt_filename = os.path.splitext(os.path.basename(xml_path))[0] + '.txt'
    txt_path = os.path.join(output_dir, txt_filename)

    with open(txt_path, 'w') as f:
        for obj in root.iter('object'):
            # 获取类别
            cls = obj.find('name').text
            if cls not in classes:
                continue
            cls_id = classes.index(cls)

            # 获取边界框坐标
            bbox = obj.find('bndbox')
            xmin = float(bbox.find('xmin').text)
            ymin = float(bbox.find('ymin').text)
            xmax = float(bbox.find('xmax').text)
            ymax = float(bbox.find('ymax').text)

            # 转换为YOLO格式
            x_center = (xmin + xmax) / 2 / img_width
            y_center = (ymin + ymax) / 2 / img_height
            width = (xmax - xmin) / img_width
            height = (ymax - ymin) / img_height

            f.write(f"{cls_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")


# 使用示例


config = load_config("data.yaml")
print(config)
classes = config['classes'] # 你的类别列表
print(classes)
xml_dir = config['val_xml']  # ！！！！！需要改成train_xml再运行一次！！！！！
output_dir = config['val_xml']  # ！！！！！需要改成train_xml再运行一次！！！！！
print(xml_dir)
print(output_dir)

for xml_file in os.listdir(xml_dir):
    if xml_file.endswith('.xml'):
        xml_path = os.path.join(xml_dir, xml_file)
        xml_to_yolo(xml_path, output_dir, classes)

运行后，labels里面得到

这样的数据

然后新建一个train.py

from ultralytics import YOLO
if __name__ == '__main__':
    model = YOLO("yolo11n.pt")

    model.train(data="data.yaml", epochs=100, imgsz=640)

点击运行：

train_batch0:

train_batch_9812

验证集标签：

验证集通过模型识别到的