1. 数据集

使用NEU-DET（钢材表面缺陷数据集）作为后续模型初步验证的数据集，但该数据下载后的数据格式为.XML格式，不能直接在YOLO系列模型中进行应用，所以需要对该数据集进行进一步的处理。

NET-DET下载链接：

https://www.kaggle.com/datasets/sovitrath/neu-steel-surface-defect-detect-trainvalid-split

2. YOLO数据格式

由于本人代码能力有限，所能够想到的方法过于繁琐，仅供参考！

由于NEU-DET数据集下载如下图所示。

其中 train_annotations和valid_annotations为标注好的文件，而train_images和valid_images为原始图像文件。

下载完数据集后，根据上述数据集的格式，需要将train_annotations和valid_annotations文件中的标记好的文件放在一个文件夹中，然后将所有的原始图像问价也放在一个文件夹中。

准备工作完成后，开始进行数据格式的转化，本次转化共需要两个python代码，一个代码是将所有的.xml文件转化为.txt文件，第二个代码则是根据原始图像。

2.1 .xml转.txt 

具体代码如下所示：

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

# 目标类别字典（可以根据需要修改）
classes = ['crazing', 'inclusion', 'patches','pitted_surface','rolled-in_scale','scratchess']  # 根据实际情况修改类别名称

def convert_annotation(xml_file, txt_file, image_width, image_height):
    # 解析xml文件
    tree = ET.parse(xml_file)
    root = tree.getroot()
    
    with open(txt_file, '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)
            
            # 获取物体的边界框信息
            bndbox = obj.find('bndbox')
            xmin = int(bndbox.find('xmin').text)
            ymin = int(bndbox.find('ymin').text)
            xmax = int(bndbox.find('xmax').text)
            ymax = int(bndbox.find('ymax').text)
            
            # 转换为YOLO格式
            x_center = (xmin + xmax) / 2.0
            y_center = (ymin + ymax) / 2.0
            width = xmax - xmin
            height = ymax - ymin
            
            # 归一化到[0, 1]范围
            x_center /= image_width
            y_center /= image_height
            width /= image_width
            height /= image_height
            
            # 写入txt文件
            f.write(f"{cls_id} {x_center} {y_center} {width} {height}\n")

def convert_xml_to_yolo_format(xml_folder, image_folder, output_folder):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    
    # 遍历xml文件夹
    for xml_file in os.listdir(xml_folder):
        if not xml_file.endswith('.xml'):
            continue
        
        # 获取对应的图像文件
        image_file = xml_file.replace('.xml', '.jpg')  # 假设图像为.jpg格式
        image_path = os.path.join(image_folder, image_file)
        
        if not os.path.exists(image_path):
            continue
        
        # 获取图像的宽高
        from PIL import Image
        img = Image.open(image_path)
        image_width, image_height = img.size
        
        # 转换XML到YOLO格式
        xml_path = os.path.join(xml_folder, xml_file)
        txt_file = os.path.join(output_folder, xml_file.replace('.xml', '.txt'))
        convert_annotation(xml_path, txt_file, image_width, image_height)

# 示例使用
xml_folder = '路径1'  #此处需要进行修改，修改为刚刚所有的.xml文件存在的文件夹的位置 XML文件夹路径
image_folder = '路径2'  # 图像文件夹路径
output_folder = '路径3'  # 输出的TXT文件夹路径，此处输出的就是我们所要的.txt格式的标注文件

convert_xml_to_yolo_format(xml_folder, image_folder, output_folder)

通过上述代码我们会得到一个文件夹，里面是根据.xml文件转化为的.txt标注文件。至此，第一步格式转化完成。

2.2 得到YOLO系列所需的数据结构格式 

YOLO系列的数据结构格式如下所示： 

但此时只得到了一个.txt的标注文件，并没有形成这样的数据结构，接下来需要得到该格式。

首先，需要再建立一个文件夹，其内容格式如下图所示：

其中，images里面是所有的原始图像文件，labels里面是2.1中所得到的.txt文件，classes.txt是该数据集中的类型标注，如下图所示：

注意：

1.此处有一个错误，最后一个类别scratchess类没有标注类型，这导致后面的模型训练中，这一类的识别出现了错误。

2.要将目标种类的名称和标注文件以及图像中的名称对应起来，否则会在后期的训练过程中产生意想不到的错误！！！

准备好上述文件后，使用下列代码进行最后的格式转化：

import os
import random
import shutil
import time
import yaml
 
 
 
class YOLOTrainDataSetGenerator:
    def __init__(self, origin_dataset_dir, train_dataset_dir, train_ratio=0.7, val_ratio=0.15, test_ratio=0.15,
                 clear_train_dir=False):
        # 设置随机数种子
        random.seed(1233)
 
        self.origin_dataset_dir = origin_dataset_dir
        self.train_dataset_dir = train_dataset_dir
        self.train_ratio = train_ratio
        self.val_ratio = val_ratio
        self.test_ratio = test_ratio
        self.clear_train_dir = clear_train_dir
 
        assert self.train_ratio > 0.5, 'train_ratio must larger than 0.5'
        assert self.val_ratio > 0.01, 'train_ratio must larger than 0.01'
        assert self.test_ratio > 0.01, 'test_ratio must larger than 0.01'
        total_ratio = round(self.train_ratio + self.val_ratio + self.test_ratio)
        assert total_ratio == 1.0, 'train_ratio + val_ratio + test_ratio must equal 1.0'
 
    def generate(self):
        time_start = time.time()
        # 原始数据集的图像目录，标签目录，和类别文件路径
        origin_image_dir = os.path.join(self.origin_dataset_dir, 'images')
        origin_label_dir = os.path.join(self.origin_dataset_dir, 'labels')
        origin_classes_file = os.path.join(self.origin_dataset_dir, 'classes.txt')
        if not os.path.exists(origin_classes_file):
            return
        else:
            origin_classes = {}
            with open(origin_classes_file, mode='r') as f:
                for cls_id, cls_name in enumerate(f.readlines()):
                    cls_name = cls_name.strip()
                    if cls_name != '':
                        origin_classes[cls_id] = cls_name
 
        # 获取所有原始图像文件名（包括后缀名）
        origin_image_filenames = os.listdir(origin_image_dir)
 
        # 随机打乱文件名列表
        random.shuffle(origin_image_filenames)
 
        # 计算训练集、验证集和测试集的数量
        total_count = len(origin_image_filenames)
        train_count = int(total_count * self.train_ratio)
        val_count = int(total_count * self.val_ratio)
        test_count = total_count - train_count - val_count
 
        # 定义训练集文件夹路径
        if self.clear_train_dir and os.path.exists(self.train_dataset_dir):
            shutil.rmtree(self.train_dataset_dir, ignore_errors=True)
        train_dir = os.path.join(self.train_dataset_dir, 'train')
        val_dir = os.path.join(self.train_dataset_dir, 'val')
        test_dir = os.path.join(self.train_dataset_dir, 'test')
        train_image_dir = os.path.join(train_dir, 'images')
        train_label_dir = os.path.join(train_dir, 'labels')
        val_image_dir = os.path.join(val_dir, 'images')
        val_label_dir = os.path.join(val_dir, 'labels')
        test_image_dir = os.path.join(test_dir, 'images')
        test_label_dir = os.path.join(test_dir, 'labels')
 
        # 创建训练集输出文件夹
        os.makedirs(train_image_dir, exist_ok=True)
        os.makedirs(train_label_dir, exist_ok=True)
        os.makedirs(val_image_dir, exist_ok=True)
        os.makedirs(val_label_dir, exist_ok=True)
        os.makedirs(test_image_dir, exist_ok=True)
        os.makedirs(test_label_dir, exist_ok=True)
 
        # 将图像和标签文件按设定的ratio划分到训练集，验证集，测试集中
        for i, filename in enumerate(origin_image_filenames):
            if i < train_count:
                output_image_dir = train_image_dir
                output_label_dir = train_label_dir
            elif i < train_count + val_count:
                output_image_dir = val_image_dir
                output_label_dir = val_label_dir
            else:
                output_image_dir = test_image_dir
                output_label_dir = test_label_dir
            src_img_name_no_ext = os.path.splitext(filename)[0]
            src_image_path = os.path.join(origin_image_dir, filename)
            src_label_path = os.path.join(origin_label_dir, src_img_name_no_ext + '.txt')
            if os.path.exists(src_label_path):
                # 复制图像文件
                dst_image_path = os.path.join(output_image_dir, filename)
                shutil.copy(src_image_path, dst_image_path)
                # 复制标签文件
                src_label_path = os.path.join(origin_label_dir, src_img_name_no_ext + '.txt')
                dst_label_path = os.path.join(output_label_dir, src_img_name_no_ext + '.txt')
                shutil.copy(src_label_path, dst_label_path)
            else:
                pass
        train_dir = os.path.normpath(train_dir)
        val_dir = os.path.normpath(val_dir)
        test_dir = os.path.normpath(test_dir)
        data_dict = {
            'train': train_dir,
            'val': val_dir,
            'test': test_dir,
            'nc': len(origin_classes),
            'names': origin_classes
        }
 
        yaml_file_path = os.path.normpath(os.path.join(self.train_dataset_dir, 'data.yaml'))
        with open(yaml_file_path, mode='w') as f:
            yaml.safe_dump(data_dict, f, default_flow_style=False, allow_unicode=True)
    
 
if __name__ == '__main__':
    g_origin_dataset_dir = 'F:/computer algorithm/YOLOV8/data/test/archive/data_set_partitioning'    #此处，修改为自己的xml文件所在位置（修改） 且文件内的布局需要一定的格式
    g_train_dataset_dir = 'F:/computer algorithm/YOLOV8/data/test/archive/success_set_partitioning'    #此处，修改为自己的转化完成后txt文件所要存放的位置（需修改）
    g_train_ratio = 0.7
    g_val_ratio = 0.15
    g_test_ratio = 0.15
    yolo_generator = YOLOTrainDataSetGenerator(g_origin_dataset_dir, g_train_dataset_dir, g_train_ratio, g_val_ratio,
                                               g_test_ratio, True)
    yolo_generator.generate()

 至此，完成VOC数据格式转化为YOLO系列数据格式。

3. VOC数据格式

3.1 VOC数据格式简介

Pascal VOC数据集由五个部分构成：JPEGImages，Annotations，ImageSets，SegmentationClass和SegmentationObject。其格式结构如下所示  

 由于本人主要侧重的是目标检测方向，故而在ImageSets中只关注Main内的文件，Main文件内存放四个文本文件，分别是test.txt、train.txt、trainval.txt和val.txt，Main的文件结构及每个文件存放的内容如下所示。

VOC数据集的标签主要是以XML文件形式进行存放。XML文件的格式标注格式如下所示：

<annotation>
  <folder>1</folder> # 图片所处文件夹
  <filename>1.bmp</filename> # 图片名
  <path>~/frcnn-image/61/ADAS/image/frcnn-image/17/77258.bmp</path>
  <source>  #图片来源相关信息
    <database>Unknown</database>  
  </source>
  <size> #图片尺寸
    <width>640</width>
    <height>480</height>
    <depth>3</depth>
  </size>
  <segmented>0</segmented>  #是否有分割label
  <object> 包含的物体
    <name>car</name>  #物体类别
    <pose>Unspecified</pose>  #物体的姿态
    <truncated>0</truncated>  #物体是否被部分遮挡（>15%）
    <difficult>0</difficult>  #是否为难以辨识的物体， 主要指要结体背景才能判断出类别的物体。虽有标注， 但一般忽略这类物体
    <bndbox>  #物体的bound box
      <xmin>2</xmin>     #左
      <ymin>156</ymin>   #上
      <xmax>111</xmax>   #右
      <ymax>259</ymax>   #下
    </bndbox>
  </object>
</annotation>

3.2 生成VOC数据集

将所有的XML文件与源图像文件方便放置在两个并列的文件夹中，如下所示：

图中，images中存放源图像文件，而labes中存放XML文件，然后执行如下代码：

import os
import shutil
import random
from glob import glob

def create_voc_structure(root_dir):
    voc_dirs = [
        "Annotations",
        "JPEGImages",
        "ImageSets/Main"
    ]
    
    for d in voc_dirs:
        os.makedirs(os.path.join(root_dir, "VOC2012", d), exist_ok=True)

def copy_files(xml_dir, img_dir, voc_root):
    annotation_dst = os.path.join(voc_root, "VOC2012", "Annotations")
    image_dst = os.path.join(voc_root, "VOC2012", "JPEGImages")
    
    xml_files = glob(os.path.join(xml_dir, "*.xml"))
    img_files = glob(os.path.join(img_dir, "*.jpg"))
    
    img_names = set(os.path.splitext(os.path.basename(img))[0] for img in img_files)
    xml_names = set(os.path.splitext(os.path.basename(xml))[0] for xml in xml_files)
    
    common_names = img_names.intersection(xml_names)
    
    for name in common_names:
        shutil.copy(os.path.join(xml_dir, name + ".xml"), annotation_dst)
        shutil.copy(os.path.join(img_dir, name + ".jpg"), image_dst)
    
    return list(common_names)

def split_dataset(file_names, voc_root, train_ratio=0.7, val_ratio=0.2, test_ratio=0.1):
    random.shuffle(file_names)
    total = len(file_names)
    
    train_end = int(total * train_ratio)
    val_end = train_end + int(total * val_ratio)
    
    train_files = file_names[:train_end]
    val_files = file_names[train_end:val_end]
    test_files = file_names[val_end:]
    trainval_files = train_files + val_files
    
    main_dir = os.path.join(voc_root, "VOC2012", "ImageSets", "Main")
    
    def write_txt(file_path, names):
        with open(file_path, "w") as f:
            for name in names:
                f.write(name + "\n")
    
    write_txt(os.path.join(main_dir, "train.txt"), train_files)
    write_txt(os.path.join(main_dir, "val.txt"), val_files)
    write_txt(os.path.join(main_dir, "test.txt"), test_files)
    write_txt(os.path.join(main_dir, "trainval.txt"), trainval_files)

def main(xml_dir, img_dir, voc_root):
    create_voc_structure(voc_root)
    file_names = copy_files(xml_dir, img_dir, voc_root)
    split_dataset(file_names, voc_root)
    print("VOC 数据集整理完成！")

# 使用示例
xml_dir = "F:\data\VOC\data\labes"  # 替换为XML文件所在路径
img_dir = "F:\data\VOC\data\images"  # 替换为图片文件所在路径
voc_root = "F:\data\VOC\\test"  # 替换为VOCdevkit根目录

main(xml_dir, img_dir, voc_root)

 即可得到VOC格式的数据集文件。

4. COCO数据格式

4.1 COCO数据集简介

COCO是常用的数据集格式，主要用于目标检测、语意分割等视觉任务等，COCO采用JSON格式来存储标注信息，相较于VOC格式，其支持更加复杂的标注。

4.2 生成COCO数据集

将xml文件和源图像文件分别放置在两个单独的文件夹中，然后执行下列代码。

import os
import json
import shutil
import xml.etree.ElementTree as ET
from sklearn.model_selection import train_test_split

# 配置数据集路径
xml_dir = r"F:\data\\NEU-DET\DATA\\labes"  # XML 文件存放路径
image_dir = r"F:\data\\NEU-DET\DATA\\images"  # 图像文件存放路径
output_dir = r"F:/data/NEU-DET/COCO/test"  # 目标数据集目录

# 目标目录结构
train_image_dir = os.path.join(output_dir, "train2017")
val_image_dir = os.path.join(output_dir, "val2017")
annotations_dir = os.path.join(output_dir, "annotations")

# 创建必要的目录
os.makedirs(train_image_dir, exist_ok=True)
os.makedirs(val_image_dir, exist_ok=True)
os.makedirs(annotations_dir, exist_ok=True)

# 获取所有 XML 文件
xml_files = [f for f in os.listdir(xml_dir) if f.endswith(".xml")]

# 划分训练集和验证集
train_files, val_files = train_test_split(xml_files, test_size=0.2, random_state=42)

# 解析 XML
def parse_voc_xml(xml_file):
    tree = ET.parse(os.path.join(xml_dir, xml_file))
    root = tree.getroot()

    filename = root.find("filename").text
    width = int(root.find("size/width").text)
    height = int(root.find("size/height").text)

    image_info = {
        "id": int(filename.split(".")[0].split("_")[-1]),  # 这里可以改成一个自增 ID
        "file_name": filename,
        "width": width,
        "height": height,
        "license": 1,  # 默认 license ID
    }

    annotations = []
    for obj in root.findall("object"):
        category_name = obj.find("name").text
        bndbox = obj.find("bndbox")
        xmin = int(bndbox.find("xmin").text)
        ymin = int(bndbox.find("ymin").text)
        xmax = int(bndbox.find("xmax").text)
        ymax = int(bndbox.find("ymax").text)

        annotations.append({
            "image_id": image_info["id"],
            "category_name": category_name,  # 先存类别名称，稍后映射为数字
            "bbox": [xmin, ymin, xmax - xmin, ymax - ymin],
            "area": (xmax - xmin) * (ymax - ymin),
            "iscrowd": 0,
            "segmentation": [],  # COCO 需要 segmentation 字段，暂时为空
        })

    return image_info, annotations


# 生成 COCO JSON
def create_coco_json(xml_file_list, json_filename, image_output_dir):
    coco_json = {
        "info": {
            "description": "Steel Defect Dataset",
            "url": "http://cocodataset.org",
            "version": "1.0",
            "year": 2017,
            "contributor": "COCO Consortium",
            "date_created": "2017/09/01"
        },
        "licenses": [
            {
                "url": "http://creativecommons.org/licenses/by-nc-sa/2.0/",
                "id": 1,
                "name": "Attribution-NonCommercial-ShareAlike License"
            }
        ],
        "images": [],
        "categories": [],
        "annotations": []
    }

    category_mapping = {}
    category_id = 1
    annotation_id = 1
    image_id_counter = 1  # 统一编号

    for xml_file in xml_file_list:
        image_info, annotations = parse_voc_xml(xml_file)
        image_info["id"] = image_id_counter  # 统一 ID
        image_id_counter += 1

        # 确保文件名有扩展名
        if "." not in image_info["file_name"]:
            image_info["file_name"] += ".jpg"  # 根据实际情况调整

        src_path = os.path.join(image_dir, image_info["file_name"])
        dst_path = os.path.join(image_output_dir, image_info["file_name"])

        print(f"Trying to copy from: {src_path} to {dst_path}")

        if os.path.exists(src_path):
            shutil.copy(src_path, dst_path)
        else:
            print(f"Warning: File not found - {src_path}")

        coco_json["images"].append(image_info)

        for ann in annotations:
            category_name = ann.pop("category_name")  # 获取类别名称
            if category_name not in category_mapping:
                category_mapping[category_name] = category_id
                coco_json["categories"].append({
                    "id": category_id,
                    "name": category_name,
                    "supercategory": "defect"  # 可以改为实际类别
                })
                category_id += 1

            ann["category_id"] = category_mapping[category_name]
            ann["id"] = annotation_id
            annotation_id += 1
            coco_json["annotations"].append(ann)

    # 保存 JSON
    with open(os.path.join(annotations_dir, json_filename), "w") as f:
        json.dump(coco_json, f, indent=4)

# 生成 COCO JSON 并整理数据
create_coco_json(train_files, "instances_train2017.json", train_image_dir)
create_coco_json(val_files, "instances_val2017.json", val_image_dir)

print("数据转换完成！")